diff --git a/src/Init/Data/Range.lean b/src/Init/Data/Range.lean index 5073bc3eba..c9791906fc 100644 --- a/src/Init/Data/Range.lean +++ b/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