Demo of a possible approach to memoization of arbitrary function return values

gistfile1.swift

//: # Memoization demo

import UIKit

/*:
  These are infrastructure classes for storing the cached function return values based on object instance, function name and input arguments.
  
  There is still work to do here. It is not thread safe.
*/
typealias MemoKey = String

class MemoValueStore<T> {
    var values = [MemoKey:T]()
}

// Map of Class instance -> Function name -> [inputs] -> Value
// var memos = [ObjectIdentifier:[String:[String:Any]?]]()
class MemoStore {
    private static var globalStores = [ObjectIdentifier:MemoStore]()

    private var storesByFunctionName = [String:Any]()
    
    class func storeForInstance(anyObject:AnyObject) -> MemoStore {
        let objectIdentifier = reflect(anyObject).objectIdentifier!
        if let store = globalStores[objectIdentifier] {
            return store
        } else {
            let newStore = MemoStore()
            globalStores[objectIdentifier] = newStore
            return newStore
        }
    }
    
    func storeForFunction<T>(functionName:String) -> MemoValueStore<T> {
        if let store = storesByFunctionName[functionName] {
            return store as! MemoValueStore<T>
        } else {
            let newStore = MemoValueStore<T>()
            storesByFunctionName[functionName] = newStore
            return newStore
        }
    }
}

/*:
  ## The protocol to adopt in classes that wish to memoize
*/
protocol Memoizing : class {
}

/*:
  ## The protocol extension that provides the implementation
*/
extension Memoizing {
    private func storageForFunction<T: Any>(functionName:String) -> MemoValueStore<T> {
        let s = self
        let memoStore = MemoStore.storeForInstance(s)
        return memoStore.storeForFunction(functionName)
    }
    
    func memoize<T : Any>(functionName: String, keys:[String], valueCreator: () -> T) -> T {
        let storage: MemoValueStore<T> = storageForFunction(functionName)
        // Don't like this, would like a more efficient way to have compound keys, e.g. equality on array 
        // of Hashable but can't work out how to wangle this yet.
        let compoundKey = "|".join(keys)
        var value: T? = storage.values[compoundKey]
        if let value = value {
            return value
        } else {
            value = valueCreator()
            storage.values[compoundKey] = value
        }
        return value!
    }
}

/*:
  ## Here we have our demo code that uses the memoization
*/

/*:
  This is the "expensive" value that we use to simulate a slow operation that we'd like to cache the result of
*/
struct ExpensiveValue {
    let value: String

    init(value:String) {
        self.value = "\(value):stored"
        sleep(5)
    }
}

/*:
  ## Using it in a class
  
  You simply add the `Memoizing` protocol to a class and in your function implementations call the `memoize` function
  provided by the Memoizing protocol extension. You pass in the function name (or some other identifying key for the function)
  and the list of keys that distinguish the current inputs from other invocations. These must be all the inputs that
  result in a change in the output of the function. Then you just pass your function body as a closure.
*/
class Whatever : Memoizing {

    func somethingExpensive(inputValue:String) -> ExpensiveValue {
        return memoize(__FUNCTION__, keys:[inputValue]) {
            return ExpensiveValue(value:inputValue)
        }
    }
    
    func somethingElseExpensive(inputValue:String) -> ExpensiveValue {
        return memoize(__FUNCTION__, keys:[inputValue]) {
            return ExpensiveValue(value:"alternative-\(inputValue)")
        }
    }
}


/*:
  Some test script to create two different instances of Whatever, ensure they don't pollute each other's caches
  and that different inputs return the correct different cached values
*/
var object = Whatever()
// First hit is slow
print("Calculating 'Testing'")
var result = object.somethingExpensive("Testing")
assert(result.value == "Testing:stored")

// First hit for this arg list is slow
print("Calculating 'Testing another'")
result = object.somethingExpensive("Testing another")
assert(result.value == "Testing another:stored")

print("Now hitting them up again, should be fast")
// These should be fast
result = object.somethingExpensive("Testing")
assert(result.value == "Testing:stored")

result = object.somethingExpensive("Testing another")
assert(result.value == "Testing another:stored")

result = object.somethingExpensive("Testing")
assert(result.value == "Testing:stored")

result = object.somethingExpensive("Testing another")
assert(result.value == "Testing another:stored")
print("Done - was it fast?")

var object2 = Whatever()
// First hit is slow
print("Calculating #2 'Testing'")
result = object2.somethingElseExpensive("Testing")
assert(result.value == "alternative-Testing:stored")

// First hit for this arg list is slow
print("Calculating #2 'Testing another'")
result = object2.somethingElseExpensive("Testing another")
assert(result.value == "alternative-Testing another:stored")

print("Now hitting them up again, should be fast")
// These should be fast
result = object2.somethingElseExpensive("Testing")
assert(result.value == "alternative-Testing:stored")

result = object2.somethingElseExpensive("Testing another")
assert(result.value == "alternative-Testing another:stored")

result = object2.somethingElseExpensive("Testing")
assert(result.value == "alternative-Testing:stored")

result = object2.somethingElseExpensive("Testing another")
assert(result.value == "alternative-Testing another:stored")
print("Done - was it fast?")