feat: add ForIn' instance for Range

src/Init/Data/Range.lean

@@ -14,6 +14,9 @@ structure Range where
   stop  : Nat
   step  : Nat := 1
 
+instance : Membership Nat Range where
+  mem i r := r.start ≤ i ∧ i < r.stop
+
 namespace Range
 universe u v
 
@@ -21,7 +24,7 @@ universe u v
   -- 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
+      return b
     else match fuel with
      | 0   => pure b
      | fuel+1 => match (← f i b) with
@@ -32,6 +35,21 @@ universe u v
 instance : ForIn m Range Nat where
   forIn := Range.forIn
 
+@[inline] protected def forIn' {β : Type u} {m : Type u → Type v} [Monad m] (range : Range) (init : β) (f : (i : Nat) → i ∈ range → β → m (ForInStep β)) : m β :=
+  let rec @[specialize] loop (start stop step : Nat) (f : (i : Nat) → start ≤ i ∧ i < stop → β → m (ForInStep β)) (fuel i : Nat) (hl : start ≤ i) (b : β) : m β := do
+    if hu : i < stop then
+      match fuel with
+      | 0   => pure b
+      | fuel+1 => match (← f i ⟨hl, hu⟩ b) with
+         | ForInStep.done b  => pure b
+         | ForInStep.yield b => loop start stop step f fuel (i + step) (Nat.le_trans hl (Nat.le_add_right ..)) b
+    else
+      return b
+  loop range.start range.stop range.step f range.stop range.start (Nat.le_refl ..) init
+
+instance : ForIn' m Range Nat inferInstance 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 (fuel i stop step : Nat) : m PUnit := do
     if i ≥ stop then
@@ -57,3 +75,11 @@ macro_rules
 
 end Range
 end Std
+
+theorem Membership.mem.upper {i : Nat} {r : Std.Range} (h : i ∈ r) : i < r.stop := by
+  simp [Membership.mem] at h
+  exact h.2
+
+theorem Membership.mem.lower {i : Nat} {r : Std.Range} (h : i ∈ r) : r.start ≤ i := by
+  simp [Membership.mem] at h
+  exact h.1