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