Functional-UIImage-generation.swift

## Functional-UIImage-generation.swift
// Generating UIImages functionally in Swift
// By Andrew Bulhak  http://dev.null.org/acb/  github:andrewcb
// Distributed under the Apache Licence
//
// For more information, see http://tech.null.org/item/201603210210_functionally_uiimage

/* --------------------------------------------------------------
 A UIImage extension providing the means of functionally generating images; the general form is:

 let myImage = UIImage (width: w, height: h) { (x, y) in return pixelValue }

 The trailing closure returns either a Float (to generate greyscale images), a 3-tuple of
 Floats (for RGB images) or a 4-tuple (for RGBA).
   -------------------------------------------------------------- */

import UIKit

// Helper class to assist in filling array buffers from a function
public struct BufferWriter<T> {
    var buffer: [T]
    var index: Array<T>.Index
    init(count: Int, repeatedValue: T) {
        self.buffer = [T](count: count, repeatedValue: repeatedValue)
        self.index = self.buffer.startIndex
    }
    mutating func write(value: T) {
        if self.index != self.buffer.endIndex {
            self.buffer[index] = value
            self.index = self.index.successor()
        }
    }
}

public extension UIImage {

    /* Most of the heavy lifting takes place here */
    private static func makeCGImage(width width:Int, height: Int, bytesPerPixel: Int, space:CGColorSpace?, bitmapInfo: CGBitmapInfo, @noescape writefunc: (Int, Int, inout BufferWriter<UInt8>)->())  -> CGImageRef? {
        var bufferWriter = BufferWriter<UInt8>(count: width*height*bytesPerPixel, repeatedValue:0)
        for i in 0..<height {
            for j in 0..<width {
                writefunc(j, i, &bufferWriter)
            }
        }

        let provider = CGDataProviderCreateWithCFData(NSData(bytes: &bufferWriter.buffer, length: width*height*bytesPerPixel))

        return CGImageCreate(
            width, height,
            8, 8*bytesPerPixel,
            width*bytesPerPixel,
            space,
            bitmapInfo,
            provider,
            nil,
            false,
            CGColorRenderingIntent.RenderingIntentDefault)
    }


    /// Generate a greyscale bitmap from a (x,y)->(grey level) function
    convenience init?(width: Int, height: Int, @noescape function:(Int, Int)->Float) {
        guard let cgimage = (UIImage.makeCGImage(width: width, height: height, bytesPerPixel: 1, space: CGColorSpaceCreateDeviceGray(), bitmapInfo: CGBitmapInfo(rawValue: CGImageAlphaInfo.None.rawValue)) { (x, y, inout writer: BufferWriter<UInt8>)   in writer.write(UInt8(function(x,y) * 255.0))})
            else { return nil }
        self.init(CGImage: cgimage)
    }

    /// Generate a RGB bitmap from a (x,y)->(r,g,b) function
    convenience init?(width: Int, height: Int, @noescape function:(Int, Int)->(Float, Float, Float)) {
        guard let cgimage = (UIImage.makeCGImage(width: width, height: height, bytesPerPixel: 3, space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: CGBitmapInfo(rawValue: CGImageAlphaInfo.None.rawValue)) {
            (x, y, inout writer: BufferWriter<UInt8>) in
            let (r,g,b) = function(x, y)
            writer.write(UInt8(r * 255.0))
            writer.write(UInt8(g * 255.0))
            writer.write(UInt8(b * 255.0))
            }) else { return nil }

        self.init(CGImage: cgimage)
    }

    /// Generate a RGBA bitmap from a (x,y)->(r,g,b,a) function
    convenience init?(width: Int, height: Int, premultiplied: Bool = false, @noescape function:(Int,Int)->(Float,Float,Float,Float)) {
        let bitmapInfo = CGBitmapInfo(rawValue: (premultiplied ? CGImageAlphaInfo.PremultipliedLast : CGImageAlphaInfo.Last).rawValue)
        guard let cgimage = (UIImage.makeCGImage(width: width, height: height, bytesPerPixel: 4, space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: bitmapInfo) {
            (x, y, inout writer: BufferWriter<UInt8>) in
            let (r,g,b,a) = function(x, y)
            writer.write(UInt8(r * 255.0))
            writer.write(UInt8(g * 255.0))
            writer.write(UInt8(b * 255.0))
            writer.write(UInt8(a * 255.0))
            }) else { return nil }

        self.init(CGImage: cgimage)
    }
}

/* --------------------------------------------------------------
  Some examples
  -------------------------------------------------------------- */


// generate a simple 2x2 checkerboard, with the tiles being 4 pixels square:

let checker1 = UIImage(width:8, height:8) { (x,y) in return Float(((x/4)+(y/4))%2) }

// a function which generates arbitrary checkerboards

func checkerboard(width w:Int, height h:Int, rows:Int, cols:Int) -> UIImage {
    return UIImage(width: w*cols, height: h*rows) { (x,y) in return Float(((x/w)+(y/h))%2) }!
}

let checker2 = checkerboard(width:8, height:8, rows:8, cols:8)

// the Mandelbrot set

func mandelbrot(width w:Int, height h:Int, maxIterations:Int=1000) -> UIImage {
    return UIImage(width:w, height:h) {(px, py) -> (Float,Float,Float) in
        let x0 = ((Float(px)/Float(w))*3.5) - 2.5
        let y0 = ((Float(py)/Float(h))*2.0) - 1.0
        var x: Float = 0.0, y: Float = 0.0
        var iteration = -1
        while iteration<maxIterations {
            if (x*x + y*y >= 4.0) { break }
            (y,x) = (2*x*y + y0, x*x - y*y + x0)
            iteration += 1
        }
        return ( (Float(iteration)*0.1)%1.0, (Float(iteration)*0.0333)%1.0, (Float(iteration)*0.01)%1.0)
    }!
}

let mandel = mandelbrot(width: 140, height: 80, maxIterations:100)