Memoizing example for learners

Using Swift
1

I have found a little gem here.

@tevelee: Thank you for posting it.

// [https://forums.swift.org/t/advent-of-code-2024/76301/57]

func memoize <In: Hashable, Out> (_ f: @escaping (In) -> Out) -> (In) -> Out {
  var memo: [In: Out] = [:]
  return {
    if let result = memo[$0] {
      return result
    } else {
      let result = f($0)
      memo[$0] = result
      return result
    }
  }
}

Could an expert disect this function for the benefit of learners?

Usage Example

func f (u:Int) -> String {
    "\(u)x\(u) = \(u * u)"
}

func g (u:Int) -> String {
    "\(u)x\(u)x\(u) = \(u * u * u)"
}

@main
struct Driver {
    static func main () {
        do {
            let u = memoize (f)
            
            print (u (2))
            print (u (2))
            print (u (2))
            print (u (7))
            print (u (7))
        }
        
        do {
            let u = memoize (g)
            
            print (u (2))
            print (u (2))
            print (u (2))
            print (u (7))
            print (u (7))
        }
    }
}

Output

2x2 = 4
2x2 = 4
2x2 = 4
7x7 = 49
7x7 = 49
2x2x2 = 8
2x2x2 = 8
2x2x2 = 8
7x7x7 = 343
7x7x7 = 343
2

The function returns a closure which captures memo variable. The variable serves a cache for past results. What's the part of the code you don't understand? Note memo is allocated in heap, that's why it can be accessed after memoize() returns.

1 Like
3

this is a cache access function that leverages Dictionary’s hashing implementation to track the precomputed values.

func memoize <In: Hashable, Out> 
    (_ f: @escaping (In) -> Out) -> (In) -> Out {
//                  ~~~~~~~~~~~     ~~~~~~~~~~~
//               input function     memoized function

the signature of memoize says it takes an input function of type (In) -> Out, and returns a modified version of that function of type (In) -> Out. in other words, they have the same type, which is a prerequisite for using memoize(f) as a drop-in replacement for f.

  var memo: [In: Out] = [:]

this line declares the backing storage of the memoized function. it’s a regular Swift Dictionary. this will be captured by the returned function.

  return {
        ...
  }

this is the minimal syntax needed to declare a local function (as a closure literal). it could be written more explicitly like this

  return {
        (i:In) -> Out in
  }

but then you would lose the ability to use the $0 abbreviation inside the closure body.

    if let result = memo[$0] {
//     ~~~~~~~~~~~~~~~~~~~~~ Have we computed this value already?
      return result
    } else {
      let result = f($0)
//     ~~~~~~~~~~~~~~~~~ Compute the value now...
      memo[$0] = result
//    ~~~~~~~~~~~~~~~~~ Save the value in the Dictionary for the 
//                      next access...
      return result
    }

this is a dictionary cache access pattern. it could be written more efficiently like this

  return {
    (i: In) -> Out in {
      let value: Out = $0 ?? f(i)
      $0 = value 
      return value
    }(&memo[i])
  }

which would avoid re-hashing In twice, which may be expensive for a lot of real world types.

also, notice that f is captured in the closure, which means f will outlive the setup of memoize. thus, it must be annotated @escaping, since the closure itself will be returned, and the closure contains a reference to f.

the memo dictionary will also be captured, but in Swift, it does not need any annotations to allow this to happen.

5 Likes
5

ChatGPT doesn’t do very well with code, but i asked GitHub Copilot and it gave me this:

AI GENERATED CONTENT

This Swift function memoize is a higher-order function that takes a function f as an argument and returns a memoized version of that function. Memoization is an optimization technique that stores the results of expensive function calls and returns the cached result when the same inputs occur again.

Here's a step-by-step explanation:

  1. Function Signature:

    func memoize <In: Hashable, Out> (_ f: @escaping (In) -> Out) -> (In) -> Out
    • In: Hashable: The input type In must conform to the Hashable protocol, which means it can be used as a key in a dictionary.
    • Out: The output type of the function.
    • _ f: @escaping (In) -> Out: The function f takes an input of type In and returns an output of type Out. The @escaping keyword indicates that the function f can outlive the scope in which it was passed.

  2. Memo Dictionary:

    var memo: [In: Out] = [:]
    • A dictionary memo is created to store the results of function calls. The keys are of type In and the values are of type Out.

  3. Return Closure:

    return {
  if let result = memo[$0] {
    return result
  } else {
    let result = f($0)
    memo[$0] = result
    return result
  }
}
    • The function returns a closure that takes an input of type In.
    • Inside the closure:
      • It first checks if the result for the given input is already in the memo dictionary.
      • If the result is found (if let result = memo[$0]), it returns the cached result.
      • If the result is not found, it calls the function f with the input, stores the result in the memo dictionary, and then returns the result.

This way, the memoize function ensures that the function f is only called once for each unique input, and subsequent calls with the same input return the cached result, improving performance for expensive or repetitive computations.

END OF AI GENERATED CONTENT

this is a pretty good breakdown for a newcomer so i would say GitHub Copilot does a decent job of understanding the Swift cache access pattern.

the AI doesn’t do a good job of explaining the closure literal syntax, which looks a lot like the structural braces you would find in places like if { } else { } when placed in the position right after the return, but at least it parsed it correctly.

3 Likes
6

Thank you, @rayx, @taylorswift

This is an excellent example showing how closures can be used to create stateful functions.

7

As functions/closures are reference type, one could write it like:

class Memoise<In: Hashable, Out> {
  
  private var memo: [In: Out] = [:]
  let f: (In) -> Out
    
  init(f: @escaping (In) -> Out) {
    self.f = f
  }
  
  func `do`(_ input: In) -> Out {
    if let result = memo[input] {
      return result
    } else {
      let result = f(input)
      memo[input] = result
      return result
    }
  }
}

You can just do it same way using func/closures and currying.

1 Like
8

Here's a variadic generic version of it that takes multiple parameters

2 Likes
9

You may want to check WWDC 2014 - session 404.
Here's an archived video for you:
https://devstreaming-cdn.apple.com/videos/wwdc/2014/404xxdxsstkaqjb/404/404_hd_advanced_swift.mov?dl=1, the related content starts around 35:40.

1 Like