Why is this memoized fibonacci so slow?

Stream: general

[James Sully]

import Std

partial def fibState (n : Nat) : Nat :=
  let rec loop : Nat → StateT (Std.HashMap Nat Nat) Id Nat
      | 0 => pure 0
      | 1 => pure 1
      | n => do
        let memo ← StateT.get
        match memo[n]? with
        | some result => pure result
        | none => do
          let result : Nat := (← loop (n-1)) + (← loop (n-2))
          StateT.set <| memo.insert n result
          pure result
  (loop n).run' ∅

#eval fibState 22

Once you start getting into the high teens low 20s it starts to take a noticeable amount of time to eval.

I would have assumed FBIP optimisations would work for something like this - am I not using the state linearly here? I can't remember or find it, but I seem to remember there was some kind of incantation for checking this.

Tbc, I'm not surprised it's slower than non-memoized fib for low n, but I am surprised it's this slow.

[James Sully]

It's true that there are two recursive calls, but because we're in a monad and there's a defined evaluation order, shouldn't the first recursive call run, and produce a new state value, which is then passed to the second recursive call? that still seems linear to me

[James Sully]

I just realised I got the order of the calls wrong, switching around

let result : Nat := (← loop (n-2)) + (← loop (n-1))

makes a big improvement. But it's still getting quite slow around n=40.

[Robin Arnez]

Yes, you get the cache, then run loop (n - 1) and loop (n - 2) which adds stuff to the cache and then discard all additions and add n, result to the original cache.

[James Sully]

doh

[James Sully]

thank you

[Robin Arnez]

To make sure you use the state linearly, always use modify for modifying operations:

import Std

partial def fibState (n : Nat) : Nat :=
  let rec loop : Nat → StateT (Std.HashMap Nat Nat) Id Nat
      | 0 => pure 0
      | 1 => pure 1
      | n => do
        let memo ← StateT.get
        match memo[n]? with
        | some result => pure result
        | none => do
          let result : Nat := (← loop (n-2)) + (← loop (n-1))
          modify fun memo => memo.insert n result
          pure result
  (loop n).run' ∅

[Robin Arnez]

set and get usually cause linearity problems because then, after the get, you have two reference to the value, the one from get, and the one which is saved in the state, so linearity gone

[James Sully]

thanks! I could only find https://leanprover-community.github.io/mathlib4_docs/Init/Control/State.html#StateT.modifyGet

[James Sully]

ahh, it's a monadstate method
https://leanprover-community.github.io/mathlib4_docs/Init/Prelude.html#modify

I forgot about that

[James Sully]

this was very helpful, thank you

[James Sully]

docs#dbgTraceIfShared was what I was looking for

[James Sully]

can this idea be made to work somehow?

instance : Nonempty (Thunk Nat) :=
  ⟨Thunk.mk (λ _ => 0)⟩

 partial def fibThunks (n : Nat) : Nat :=
  if h : n < 2 then
    n
  else
    let rec arr : Vector (Thunk Nat) n := Vector.ofFn
      λ | ⟨0,_⟩ => ↑(0 : Nat)
        | ⟨1,_⟩ => ↑(1 : Nat)
        | ⟨i+2,_⟩ => Thunk.mk λ () => arr[i].get + arr[i+1].get

    (dbgTraceIfShared "oops" arr)[n-2].get + arr[n-1].get

#eval fibThunks 25

[Aaron Liu]

This works

instance : Nonempty (Thunk Nat) :=
  ⟨Thunk.mk (λ _ => 0)⟩

partial def fibThunks (n : Nat) : Nat :=
  if h : n < 2 then
    n
  else
    let rec arr : Vector (Thunk Nat) n := Vector.ofFn
      λ | ⟨0,_⟩ => Thunk.pure 0
        | ⟨1,_⟩ => Thunk.pure 1
        | ⟨i+2,_⟩ => Thunk.mk λ () => arr[i].get + arr[i+1].get
    let arr := dbgTraceIfShared "oops" arr
    arr[n-2].get + arr[n-1].get

#eval fibThunks 25

[Ayhon]

Does it seem to be caching the values for you? When I execute fibThunks 31 it takes a while.

[James Sully]

I think it's not caching the values. I assumed the perf problem was due to sharing but I must have been mistaken