From bbfcb1cfb2bb00bd8a24649fd8ea1474728b6019 Mon Sep 17 00:00:00 2001 From: Sebastian Ullrich Date: Sat, 25 Dec 2021 17:09:41 +0100 Subject: [PATCH] perf: allocation-free `for i in [n:m] do` --- src/Init/Data/Range.lean | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/Init/Data/Range.lean b/src/Init/Data/Range.lean index 804cda9e78..5073bc3eba 100644 --- a/src/Init/Data/Range.lean +++ b/src/Init/Data/Range.lean @@ -18,27 +18,28 @@ namespace Range universe u v @[inline] protected def forIn {β : Type u} {m : Type u → Type v} [Monad m] (range : Range) (init : β) (f : Nat → β → m (ForInStep β)) : m β := - let rec @[specialize] loop (i : Nat) (j : Nat) (b : β) : m β := do - if j ≥ range.stop then + -- pass `stop` and `step` separately so the `range` object can be eliminated through inlining + let rec @[specialize] loop (fuel i stop step : Nat) (b : β) : m β := do + if i ≥ stop then pure b - else match i with + else match fuel with | 0 => pure b - | i+1 => match (← f j b) with + | fuel+1 => match (← f i b) with | ForInStep.done b => pure b - | ForInStep.yield b => loop i (j + range.step) b - loop range.stop range.start init + | ForInStep.yield b => loop fuel (i + step) stop step b + loop range.stop range.start range.stop range.step init instance : ForIn m Range Nat where forIn := Range.forIn @[inline] protected def forM {m : Type u → Type v} [Monad m] (range : Range) (f : Nat → m PUnit) : m PUnit := - let rec @[specialize] loop (i : Nat) (j : Nat) : m PUnit := do - if j ≥ range.stop then + let rec @[specialize] loop (fuel i stop step : Nat) : m PUnit := do + if i ≥ stop then pure ⟨⟩ - else match i with + else match fuel with | 0 => pure ⟨⟩ - | i+1 => f j; loop i (j + range.step) - loop range.stop range.start + | fuel+1 => f i; loop fuel (i + step) stop step + loop range.stop range.start range.stop range.step instance : ForM m Range Nat where forM := Range.forM