Memoization in Swift 2

The WWDC 2014 talk "Advanced Swift" (https://developer.apple.com/videos/wwdc/2014/#404) included a sample 'memoization function' which could speed up performance of certain functions by storing function results in a dictionary. Here's my implementation of it:

func memoize<Input: Hashable, Output>(function: Input -> Output) -> (Input -> Output) {
    var lookupTable = [Input : Output]()
    return {input in
        if lookupTable[input] == nil {lookupTable[input] = function(input)}
        return lookupTable[input]!
    }

}

This implementation is very similar to the first implementation of memoization presented in the Advanced Swift talk. However, that first implementation had an issue with creating recursive functions:

let factorialMemoized = memoize {n in return n < 2 ? 1 : factorialMemoized(n-1)*n} //Error: Variable used within its own initial value when using 'Take One' implementation of memoization

At the time, the workaround used was to create a more complicated memoization function. However, because Swift now supports declaring a 'let' without initializing it, the first implementation from the talk can be used to create memoized functions with a little work by the caller:

let memoizedFactorial: Int->Int;
do {
    func factorial(n: Int) -> Int {return n < 2 ? 1 : factorial(n-1)*n}
    memoizedFactorial = memoize(factorial)
}

The trick is that the temporary function 'factorial' is removed upon exiting the scope of the 'do' block, but memoizedFactorial remains because it was declared in the global scope.

Update 2: The above trick will not work.

The mistake here is that while memoizedFactorial is memoized, factorial is not. So calling memoizedFactorial(3) and then memoizedFactorial(3) again would yield a memoized result. But calling memoizedFactorial(4) would result in a call to factorial(4) which does not run the memoization routine every time it has to compute, so it would do the entire set of recursive calls of factorial(3).

It would look like this: memoizedFactorial → memoization check → factorial → factorial . . .
The lack of a memoization check between factorials means that even if the value of factorial(3) is stored in the dictionary, factorial(4) will run the factorial(3) computation instead of getting the stored value. Hence it's not proper memoization and will fail to improve performance as it should.

I realized this while skimming an article from an MIT course site linked from a Stack Overflow question . . . it gave an example of this problem in Python, and I thought "well, clearly Swift's Safe Typing™ will save the day, so this memoize can't have that problem, right?" But not even Swift Typing™ can save me from stupidity. Darn.

The good news is that the function I created *below* should still work. The reason for this is that instead of calling *itself* directly, the method being memoized, like factorial(recur, n) calls one of its parameters. (recur, in this case) And the memoized function that I return passes itself as that parameter, so it is called to run the memoization routine at every recursive call. So it shouldn't have the same problem as memoizedFactorial above.

It would look like this: factorial(memoFunc, n) → memoFunc(n) : memoizationCheck → factorial(memoFunc, n) → memoFunc(n): memoizationCheck → factorial...
Notice that before factorial is recursively called the memoization check is performed. Just as it should be.

update: Here's an interesting implementation of factorial that alleviates the work the caller needs to do. Uses the same call syntax as the 'Take Two' implementation from the Advanced Swift video, but creates and returns a function instead of using an implicitly unwrapped closure.

func recursiveMemoize<Input: Hashable, Output>(function: (Input->Output, Input) -> Output) -> (Input) -> Output {
    var lookupTable = [Input : Output]()
    func memoized(input: Input) -> Output {
        if lookupTable[input] == nil {lookupTable[input] = function(memoized, input)}
        lookupTable
        return lookupTable[input]!
    }
    return memoized

}

Comments

Popular posts from this blog

The Hidden Garageband Sampler (and other features)

How to add a folder to Launchpad