From 2f18de6c83130b8fd309d9fc390a627743fd92e0 Mon Sep 17 00:00:00 2001
From: Nada Amin <namin@alum.mit.edu>
Date: Mon, 30 Mar 2026 21:41:52 -0400
Subject: [PATCH] Bound range forIn step condition to in-range indices

---
 Loom/MonadAlgebras/WP/Basic.lean | 36 +++++++++++++++++++-------------
 Loom/MonadAlgebras/WP/Gen.lean   | 11 +++++-----
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/Loom/MonadAlgebras/WP/Basic.lean b/Loom/MonadAlgebras/WP/Basic.lean
index 674de8d..1ac6b72 100644
--- a/Loom/MonadAlgebras/WP/Basic.lean
+++ b/Loom/MonadAlgebras/WP/Basic.lean
@@ -298,7 +298,8 @@ theorem triple_forIn_list {α β}
 theorem triple_forIn_range'_aux {β}
   {xstart : ℕ} {step : ℕ} (n : ℕ) (init : β) {f : ℕ → β → m (ForInStep β)}
   (inv : ℕ → β → l)
-  (hstep : ∀ i b,
+  (hstep_pos : step > 0)
+  (hstep : ∀ i b, xstart ≤ i → i < xstart + n * step →
     triple
       (inv i b)
       (f i b)
@@ -308,40 +309,47 @@ theorem triple_forIn_range'_aux {β}
   unhygienic induction i generalizing init
   { simp [triple_pure] }
   intro le
-  simp [List.range']; apply triple_bind; apply hstep
-  rintro (_|_) <;> simp [triple_pure]
-  simp [Nat.add_assoc, <-Nat.add_one_mul]
-  have : (n - (n_1 + 1) + 1) = n - n_1 := by omega
-  rw [this]; apply a; omega
+  simp [List.range']; apply triple_bind
+  · have hlo : xstart ≤ xstart + (n - (n_1 + 1)) * step := Nat.le_add_right _ _
+    have hhi : xstart + (n - (n_1 + 1)) * step < xstart + n * step := by
+      apply Nat.add_lt_add_left; apply Nat.mul_lt_mul_of_pos_right _ hstep_pos; omega
+    exact hstep _ _ hlo hhi
+  · rintro (_|_) <;> simp [triple_pure]
+    simp [Nat.add_assoc, <-Nat.add_one_mul]
+    have : (n - (n_1 + 1) + 1) = n - n_1 := by omega
+    rw [this]; apply a; omega
 
 
 theorem triple_forIn_range' {β}
   {xstart : ℕ} {step : ℕ} (n : ℕ) {init : β} {f : ℕ → β → m (ForInStep β)}
   (inv : ℕ → β → l)
-  (hstep : ∀ i b,
+  (hstep_pos : step > 0)
+  (hstep : ∀ i b, xstart ≤ i → i < xstart + n * step →
     triple
       (inv i b)
       (f i b)
       (fun | .yield b' => inv (i + step) b' | .done b' => inv (xstart + n * step) b')) :
   triple (inv xstart init) (forIn (List.range' xstart n step) init f) (inv (xstart + n * step)) := by
-    have := triple_forIn_range'_aux (i := n) (m := m) n init inv hstep (by omega)
+    have := triple_forIn_range'_aux (i := n) (m := m) n init inv hstep_pos hstep (by omega)
     simp_all only [Nat.sub_self, Nat.zero_mul, Nat.add_zero]
 
 theorem triple_forIn_range {β}
   (xs : Std.Range) (init : β) (f : ℕ → β → m (ForInStep β))
   (inv : ℕ → β → l)
-  (hstep : ∀ i b,
+  (hstep_pos : xs.step > 0)
+  (hstep : ∀ i b, xs.start ≤ i → i < xs.start + ((xs.stop - xs.start + xs.step - 1) / xs.step) * xs.step →
     triple
       (inv i b)
       (f i b)
       (fun | .yield b' => inv (i + xs.step) b' | .done b' => inv (xs.start + ((xs.stop - xs.start + xs.step - 1) / xs.step) * xs.step) b')) :
   triple (inv xs.start init) (forIn xs init f) (inv (xs.start + ((xs.stop - xs.start + xs.step - 1) / xs.step) * xs.step)) := by
-  simp; apply triple_forIn_range'; apply hstep
+  simp; apply triple_forIn_range' (hstep_pos := hstep_pos)
+  intro i b hi hlt; exact hstep i b hi hlt
 
 theorem triple_forIn_range_step1 {β}
   {xs : Std.Range} {init : β} {f : ℕ → β → m (ForInStep β)}
   (inv : ℕ → β → l)
-  (hstep : ∀ i b,
+  (hstep : ∀ i b, xs.start ≤ i → i < xs.stop →
     triple
       (inv i b)
       (f i b)
@@ -349,10 +357,10 @@ theorem triple_forIn_range_step1 {β}
   xs.step = 1 ->
   xs.start <= xs.stop ->
   triple (inv xs.start init) (forIn xs init f) (inv xs.stop) := by
-  have := triple_forIn_range (m := m)  xs init f inv
-  intro h eq;
+  intro h eq
+  have := triple_forIn_range (m := m) xs init f inv (by omega : xs.step > 0)
   simp [h] at *
-  rw [Nat.add_sub_of_le eq] at this;
+  rw [Nat.add_sub_of_le eq] at this
   solve_by_elim
 
 end Loops
diff --git a/Loom/MonadAlgebras/WP/Gen.lean b/Loom/MonadAlgebras/WP/Gen.lean
index 2d26445..98b8961 100644
--- a/Loom/MonadAlgebras/WP/Gen.lean
+++ b/Loom/MonadAlgebras/WP/Gen.lean
@@ -209,7 +209,7 @@ theorem triple_forIn_deacreasing {β} {measure : β -> ℕ}
   apply le_trans'; apply wp_cons; rotate_left 2; apply le_trans; rotate_left 1
   apply triple_forIn_range_step1 (inv := fun i b => ⌜ measure b + i <= measure init ⌝ ⊓ inv b) <;>
     try solve | aesop
-  { simp; intro i b
+  { simp; intro i b _hlt
     by_cases h : measure b + i ≤ measure init <;> simp [h, triple]
     apply le_trans; apply hstep; omega
     apply wp_cons; rintro (b'|b') <;> simp
@@ -228,7 +228,7 @@ noncomputable
 def WPGen.forWithInvariant {xs : Std.Range} {init : β} {f : ℕ → β → m (ForInStep β)}
   (inv : ℕ → β → List l) (wpg : ∀ i b, WPGen (f i b)) (xs1 : xs.step = 1) (xs_le : xs.start <= xs.stop := by omega) :
     WPGen (forIn xs init (fun i b => do invariantGadget (inv i b); (f i b))) where
-    get := ⌜∀ i b, invariantSeq (inv i b) <= (wpg i b).get fun
+    get := ⌜∀ i b, xs.start ≤ i → i < xs.stop → invariantSeq (inv i b) <= (wpg i b).get fun
       | .yield b' => invariantSeq <| inv (i + 1) b'
       | .done b'  => invariantSeq <| inv xs.stop b'⌝
       ⊓ spec
@@ -239,9 +239,10 @@ def WPGen.forWithInvariant {xs : Std.Range} {init : β} {f : ℕ → β → m (F
       apply (triple_spec ..).mpr
       simp [invariantGadget]
       apply triple_forIn_range_step1 (fun i b => (inv i b).foldr (·⊓·) ⊤)
-      simp [invariantSeq, <-xs1] at h
-      intro i b; apply (wpg i b).intro
-      all_goals solve_by_elim
+      · simp [invariantSeq, <-xs1] at h
+        intro i b hi hlt; apply (wpg i b).intro
+        exact h i b hi hlt
+      all_goals assumption
 
 -- noncomputable
 -- def WPGen.forWithInvariantDecreasing {β} {measure : β -> ℕ}