perf: allocation-free for i in [n:m] do

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