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 foldr
specialized 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