Memoization with StateM🔗

Memoization is a technique used to speed up computer programs by storing the results of expensive function calls and returning the cached result when the same inputs occur again. A HashMap combined with a StateM monad is a powerful way to implement this in Lean. StateM allows us to carry the state of our cache (here HashMap) through our computations.

Below is an example of the Fibonacci sequence implemented with memoization.

/--
  Recursive Fibonacci with memoization.
  We use StateM Memo to carry the cache.
-/
abbrev FibState := HashMap Nat Nat
abbrev FibM := StateM FibState

def fib (n: Nat) : FibM Nat := do
  match n with
  | 0 => return 1
  | 1 => return 1
  | k + 2 => do
    let m ← get
    match m.get? n with -- check if we calculated it before
    | some v => return v
    | none => do
      let v1 ← fib k -- calculate at k and update the state
      let v2 ← fib (k + 1)
      let v := v1 + v2
      modify (fun m => m.insert n v)
      return v

/-
- `run` -> calculates given an initial state and 
  returns the result and the final state
- `run'` -> given an initial state returns the result
-/
10119911756749018713965376799211044556615579094364594923736162239653346274#eval fib 350 |>.run' {}
(20365011074,
 Std.HashMap.ofList [(2, 2),
  (3, 3),
  (4, 5),
  (5, 8),
  (6, 13),
  (7, 21),
  (8, 34),
  (9, 55),
  (10, 89),
  (11, 144),
  (12, 233),
  (13, 377),
  (14, 610),
  (15, 987),
  (16, 1597),
  (17, 2584),
  (18, 4181),
  (19, 6765),
  (20, 10946),
  (21, 17711),
  (22, 28657),
  (23, 46368),
  (24, 75025),
  (25, 121393),
  (26, 196418),
  (27, 317811),
  (28, 514229),
  (29, 832040),
  (30, 1346269),
  (31, 2178309),
  (32, 3524578),
  (33, 5702887),
  (34, 9227465),
  (35, 14930352),
  (36, 24157817),
  (37, 39088169),
  (38, 63245986),
  (39, 102334155),
  (40, 165580141),
  (41, 267914296),
  (42, 433494437),
  (43, 701408733),
  (44, 1134903170),
  (45, 1836311903),
  (46, 2971215073),
  (47, 4807526976),
  (48, 7778742049),
  (49, 12586269025),
  (50, 20365011074)])#eval fib 50 |>.run {}

Using memoization allows us to compute fib 350 almost instantly, whereas a naive recursive implementation would take a very long time.