Swift memoization helpers

Memoize.swift

// Source: https://www.hackingwithswift.com/plus/high-performance-apps/using-memoization-to-speed-up-slow-functions
// Source: https://medium.com/@mvxlr/swift-memoize-walk-through-c5224a558194

typealias MemFn<Input, Output> = (Input) -> Output

// work with any sort of input and output as long as the input is hashable
// accept a function that takes Input and returns Output, and return a function with the same signature
func memoize<Input: Hashable, Output>(
    _ work: @escaping MemFn<Input, Output>
) -> MemFn<Input, Output> {
    // our item cache
    var cache = [Input: Output]()

    // send back a new closure that does our calculation
    return { input in
        if let cached = cache[input] {
            return cached
        }

        let result = work(input)
        cache[input] = result
        return result
    }
}

typealias MemRecFn<Input, Output> = (MemFn<Input, Output>, Input) -> Output

func memoizeRec<Input: Hashable, Output>(
    work: @escaping MemRecFn<Input, Output>
) -> MemFn<Input, Output> {
    var cache = [Input: Output]()

    func wrap(_ input: Input) -> Output {
        if let cached = cache[input] {
            return cached
        }

        let result = work(wrap, input)
        cache[input] = result
        return result
    }

    return wrap
}

// Examples
let memoizedSqrt = memoize(sqrt)

let memoizedFactorial = memoizeRec {
    factorial, x in
    x == 0 ? 1 : x * factorial(x-1)
}