memoize.swift

/// memoize function from Advanced Swift.
/// https://developer.apple.com/videos/wwdc/2014/?id=404
///
/// let fib = memoize {(fib:(Int)->Int, x:Int) in (x <= 1) ? 1 : fib(x - 1) + fib(x - 2)}
///
/// let 𝜑 = Double(fib(45)) / Double(fib(44)) // ≅ 1.61803398874989...
public func memoize<A : Hashable, R>(body: ((A)->R, A)->R) -> (A)->R {
  var memo:[A: R] = [:]
  var memoized: ((A)->R)!
  memoized = {(a) in
    if let r = memo[a] { return r }
    let r = body(memoized, a)
    memo[a] = r
    return r
  }
  return memoized
}

// About `var memoized: ((A)->R)!`
// - Local functions can't be recursive, so `func memoized` will not work
//   "Local functions cannot reference themselves"
// - Definition uses itself, so `let memoized:(A->R) = {}()` will not work
//   "Variable used within its own initial value"


// MARK: Overloads
/// memoize for functions with 2 arguments
public func memoize<A : Hashable, B : Hashable, R>(body: ((A, B)->R, A, B)->R)
  -> (A, B)->R
{
  var memo:[A: [B: R]] = [:]
  var memoized: ((A, B)->R)!
  memoized = {(a, b) in
    if let r = memo[a]?[b] { return r }
    
    if memo[a] == nil { memo[a] = [:] }
    let r = body(memoized, a, b)
    memo[a]![b] = r
    return r
  }
  return memoized
}

/// memoize for functions with 3 arguments
public func memoize<A : Hashable, B : Hashable, C : Hashable, R>
  (body: ((A, B, C)->R, A, B, C)->R)
  -> (A, B, C)->R
{
  var memo:[A: [B: [C: R]]] = [:]
  var memoized: ((A, B, C)->R)!
  memoized = {(a, b, c) in
    if let r = memo[a]?[b]?[c] { return r }
    
    if memo[a] == nil { memo[a] = [:] }
    if memo[a]![b] == nil { memo[a]![b] = [:] }
    let r = body(memoized, a, b, c)
    memo[a]![b]![c] = r
    return r
  }
  return memoized
}

/// memoize for functions with 4 arguments
public func memoize<A : Hashable, B : Hashable, C : Hashable, D : Hashable, R>
  (body: ((A, B, C, D)->R, A, B, C, D)->R)
  -> (A, B, C, D)->R
{
  var memo:[A: [B: [C: [D: R]]]] = [:]
  var memoized: ((A, B, C, D)->R)!
  memoized = {(a, b, c, d) in
    if let r = memo[a]?[b]?[c]?[d] { return r }
    
    if memo[a] == nil { memo[a] = [:] }
    if memo[a]![b] == nil { memo[a]![b] = [:] }
    if memo[a]![b]![c] == nil { memo[a]![b]![c] = [:] }
    let r = body(memoized, a, b, c, d)
    memo[a]![b]![c]![d] = r
    return r
  }
  return memoized
}