Zulip Chat Archive

Stream: general

Topic: Why would hand-written recursion be slower than foldr?


Paro (Dec 11 2024 at 09:15):

I am doing Advent of Code with Lean 4. When solving the problem of today, I wrote the following code:

def blink : Nat -> List Nat
| 0 => [1]
| x =>
  let s := Nat.toDigits 10 x
  if s.length % 2 == 0
  then
    let (l, r) := s.splitAt (s.length / 2)
    [l.asString.toNat!, r.asString.toNat!]
  else [x * 2024]

--@[specialize] def foldr (f : α → β → β) (init : β) : List α → β
--  | []     => init
--  | a :: l => f a (foldr f init l)

def blinkAll : List Nat -> List Nat
| []      => []
| x :: xs => blink x ++ blinkAll xs

def solvePart1 (stones : List Nat) (n : Nat := 25) :=
  n.repeat (·.foldr (blink · ++ ·) []) stones

def parseInput (input : String) :=
  input.splitOn.map String.toNat!

def input := parseInput "5688 62084 2 3248809 179 79 0 172169"

#eval solvePart1 input 25 |>.length

As you can see, the blinkAll function is essentially foldrspecialized with f := (blink · ++ ·) and init := []. However, the solvePart1 implemented using foldr can successfully give an answer with n := 25, but if (·.foldr (blink · ++ ·) []) is replaced with blinkAll, the language server would crash. (Actually, it crashes when n := 20.)

What did Lean 4 do to make foldr mysteriously faster than hand-written recursion?

Markus Himmel (Dec 11 2024 at 09:25):

Lean defines an optimized version of List.foldr based on Array and tells the compiler to use that instead of the naive foldr implementation. You can see this redefinition here: https://github.com/leanprover/lean4/blob/ffac974dba799956a97d63ffcb13a774f700149c/src/Init/Data/List/Impl.lean#L90-L94 The @[csimp] attribute tells the compiler to use List.foldrTR instead when encountering a call to List.foldr. Digging a little deeper, you will find that List.foldrTR eventually calls docs#Array.foldrM, which has an implemented_by attribute. implemented_by is basically the unsafe version of csimp, and so the compiler will actually invoke docs#Array.foldrMUnsafe, which is even more optimized.

Jannis Limperg (Dec 11 2024 at 09:48):

Wait, so List.foldr actually copies the whole list into an array, then folds over that? I feel like this should be mentioned more prominently, ideally in the foldr docstring.

Kim Morrison (Dec 11 2024 at 11:26):

Depends what you mean by "actually". :-)

Sebastian Ullrich (Dec 11 2024 at 12:29):

But also, do compile your code into a binary if you care about performance


Last updated: May 02 2025 at 03:31 UTC