fix(realize): preserve captured iterate environments

Complexity/Tactic/Realize.lean

@@ -31,7 +31,8 @@ the client syntax, unfolding definitions, with one rule per form:
 * `match` on an optional word or word pair: the client equals the branch on its presence tag
   (`Realizes.congr`, proved by cases).
 * `iterate step D n a` whose step begins with a query: `Realizes.iterate₀`, after
-  `iterate_comap` moves `D` into the step.
+  adapting `D` and any captured live words to the whole environment. The source and reached-state
+  domain are transported back to the client's original loop.
 * `boundedWhile G step D n a` whose step begins with a query: `Realizes.while₀`. The guard
   `G D r` may be `a = b`, `a < b`, `a > b` or `a ≠ 0` (tested as `0 < a`), with any decidability
   instance, over the live words and the state's words; an operand independent of the state, such
@@ -953,13 +954,14 @@ private partial def realizeIterate (ctx : Ctx) (env : Env) (d : Expr) (atoms : A
     | none => encodingOf σ
   let l ← encodingWidth out
   let envB ← extend env σ out l
+  let captured := step.containsFVar d.fvarId!
   -- the body over the whole environment
-  let body ← if D == d then pure step else
+  let body ← if D == d && !captured then pure step else
     withLocalDeclD `x envB.ty fun x => do
       let x1 ← mkAppM ``Prod.fst #[x]
       let x2 ← mkAppM ``Prod.snd #[x]
       let Dx := (← mkLambdaFVars #[d] D).beta #[x1]
-      mkLambdaFVars #[x] (mkApp step (← mkAppM ``Prod.mk #[Dx, x2]))
+      mkLambdaFVars #[x] (mkApp (step.replaceFVar d x1) (← mkAppM ``Prod.mk #[Dx, x2]))
   let hB ← realizeProg ctx envB body (some out)
   let hq ← proveQuery body envB.ty
   let hLoop ← mkAppM ``Realizes.iterate₀ #[env.input, out, hB.proof, hq]
@@ -989,10 +991,71 @@ private partial def realizeIterate (ctx : Ctx) (env : Env) (d : Expr) (atoms : A
       mkLambdaFVars #[w, d, s] (mkAnd (hA.safe.beta #[w, d, s]) (mkAnd sn sl))
   let progR ← mkLambdaFVars #[d] (← mkAppM ``iterate #[body, d, n, a])
   let proof ← certify ctx env proof progR out safe
-  if D == d then return { proof, output := out, safe }
+  if D == d && !captured then return { proof, output := out, safe }
   let g ← mkLambdaFVars #[d] D
-  let eq ← mkLambdaFVars #[d] (← mkEqSymm (← mkAppM ``iterate_comap #[g, step, d, n, a]))
-  return { proof := ← mkAppM ``Realizes.congr #[proof, eq], output := out, safe }
+  let eq ← if captured then do
+    -- The adapted step agrees at this fixed environment, not at every possible environment.
+    let ty ← withLocalDeclD `n (mkConst ``Nat) fun n => withLocalDeclD `a σ fun a => do
+      mkForallFVars #[n, a] (← mkEq (← mkAppM ``iterate #[body, d, n, a])
+        (← mkAppM ``iterate #[step, D, n, a]))
+    let h ← Term.withSynthesize <| Term.elabTermEnsuringType (← `(by
+      intro n a
+      induction n generalizing a with
+      | zero => rfl
+      | succ n ih =>
+        simp only [iterate]
+        congr 1
+        funext b
+        exact ih b)) (some ty)
+    mkLambdaFVars #[d] (mkApp2 (← instantiateMVars h) n a)
+  else
+    mkLambdaFVars #[d] (← mkEqSymm (← mkAppM ``iterate_comap #[g, step, d, n, a]))
+  let proof ← mkAppM ``Realizes.congr #[proof, eq]
+  if !captured then return { proof, output := out, safe }
+  -- State the reached-state domain in the original loop's terms, as in `realizeWhile`.
+  let nat := mkConst ``Nat
+  let state := mkConst ``Memory.State
+  let pairTy ← mkAppM ``Prod #[← inferType D, σ]
+  let clientSafe ← withLocalDeclD `w nat fun w =>
+    withLocalDeclD `x pairTy fun x =>
+      withLocalDeclD `s state fun s => do
+        let pair ← mkAppM ``Prod.mk #[d, ← mkAppM ``Prod.snd #[x]]
+        mkLambdaFVars #[w, x, s] (← reducePairs (hB.safe.beta #[w, pair, s]))
+  let safe' ← withLocalDeclD `w nat fun w => withLocalDeclD `s state fun s => do
+    let pair ← mkAppM ``Prod.mk #[d, a]
+    let sn ← reducePairs (hN.safe.beta #[w, pair, s])
+    let sl ← mkAppM ``iterationSafe #[step, clientSafe, w, D, n, a, s]
+    mkLambdaFVars #[w, d, s] (mkAnd (hA.safe.beta #[w, d, s]) (mkAnd sn sl))
+  let pointwise ← withLocalDeclD `w nat fun w => withLocalDeclD `s state fun s => do
+    let ty ← withLocalDeclD `n nat fun n => withLocalDeclD `a σ fun a =>
+      withLocalDeclD `s state fun s => do
+        mkForallFVars #[n, a, s] (← mkEq
+          (← mkAppM ``iterationSafe #[body, hB.safe, w, d, n, a, s])
+          (← mkAppM ``iterationSafe #[step, clientSafe, w, D, n, a, s]))
+    let h ← Term.withSynthesize <| Term.elabTermEnsuringType (← `(by
+      intro n a s
+      induction n generalizing a s with
+      | zero => rfl
+      | succ n ih =>
+        simp only [iterationSafe]
+        congr 1
+        exact ih _ _)) (some ty)
+    let pair ← mkAppM ``Prod.mk #[d, a]
+    let sn ← reducePairs (hN.safe.beta #[w, pair, s])
+    let wrap ← withLocalDeclD `p (mkSort Level.zero) fun p =>
+      mkLambdaFVars #[p] (mkAnd (hA.safe.beta #[w, d, s]) (mkAnd sn p))
+    mkLambdaFVars #[w, d, s] (← mkAppM ``congrArg
+      #[wrap, mkApp3 (← instantiateMVars h) n a s])
+  let safeEq ← withLocalDeclD `w nat fun w => do
+    let hd ← withLocalDeclD `e env.ty fun e => do
+      mkLambdaFVars #[e] (← mkAppM ``funext #[mkApp2 pointwise w e])
+    mkLambdaFVars #[w] (← mkAppM ``funext #[hd])
+  let safeEq ← mkAppM ``funext #[safeEq]
+  let proofTy ← inferType proof
+  let wrap ← withLocalDeclD `safe (← inferType safe) fun s =>
+    mkLambdaFVars #[s] (mkAppN proofTy.getAppFn (proofTy.getAppArgs.set! 8 s))
+  let proof ← mkAppM ``Eq.mp #[← mkAppM ``congrArg #[wrap, safeEq], proof]
+  return { proof, output := out, safe := safe' }
 
 /-- Certify `boundedWhile G step D n a` for the environment variable `d`.
 

ComplexityTests/WordRAMRealizeIterateCapture.lean

@@ -0,0 +1,134 @@
+/-
+Copyright (c) 2026 Samuel Schlesinger. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Samuel Schlesinger
+-/
+
+module
+-- upstream: local-only (captured iterate environment regression tests)
+
+public import Complexity.Tactic.Realize
+
+/-!
+# Iteration steps retain runtime captures
+
+The explicit loop environment can be `Unit` while the step closes over other live runtime words.
+The same source, input encoding, per-step memory query and zero local-work grading are retained.
+-/
+
+public section
+
+namespace Cslib.WordRAM.Realization.Tests.IterateCapture
+
+open Cslib.DataStructure
+
+/-- The read address is captured; neither the initial state nor the count selects native code. -/
+theorem captured_read : Realizes .exact
+    (fun d : Fin 3 → ℕ => iterate (fun _ : Unit × ℕ => Memory.read (d 0)) () (d 1) (d 2))
+    (.refl _) (WordEncoding.encode (α := ℕ)) (fun _ _ _ => True) (fun _ _ => 0) := by
+  realize
+
+/-- Captures are also adapted when the explicit loop argument is the whole environment. -/
+theorem captured_read_whole : Realizes .exact
+    (fun d : Fin 3 → ℕ =>
+      iterate (fun _ : (Fin 3 → ℕ) × ℕ => Memory.read (d 0)) d (d 1) (d 2))
+    (.refl _) (WordEncoding.encode (α := ℕ)) (fun _ _ _ => True) (fun _ _ => 0) := by
+  realize
+
+/-- A separately certified call keeps its domain at every reached memory, in the original
+captured source's terms rather than an auxiliary environment chosen by the frontend. -/
+theorem captured_call (m : ℕ × ℕ → Memory.Prog ℕ)
+    (P : ℕ → ℕ × ℕ → Memory.State → Prop)
+    (hm : Realizes .exact m (WordEncoding.encode (α := ℕ × ℕ))
+      (WordEncoding.encode (α := ℕ)) P (fun _ _ => 0)) :
+    Realizes .exact
+      (fun d : Fin 3 → ℕ => iterate (fun _ : Unit × ℕ => do
+        let v ← Memory.read (d 0)
+        m (d 0, v)) () (d 1) (d 2))
+      (.refl _) (WordEncoding.encode (α := ℕ))
+      (fun w d s => iterationSafe
+        (fun _ : Unit × ℕ => do let v ← Memory.read (d 0); m (d 0, v))
+        (fun w _ s => P w (d 0, s.mem (d 0)) s) w () (d 1) (d 2) s)
+      (fun _ _ => 0) := by
+  realize
+  exact h
+
+/-- The same emitted body handles distinct captured addresses and zero/two iterations.
+Both executions retain every initial memory cell; the query count remains charged. -/
+theorem same_body_distinct_environments : ∃ r C, ∃ code : Code .exact r,
+    ∃ result : Fin r, ∃ x₀ y₀ x₁ y₁ : Data r,
+      ¬code.ReadsWidth ∧ x₀.Bounded 3 ∧ x₁.Bounded 3 ∧
+      ExecWithin 3 code x₀ y₀ C ∧ y₀.regs result = 5 ∧ y₀.mem = x₀.mem ∧
+      ExecWithin 3 code x₁ y₁ (3 * C) ∧ y₁.regs result = 7 ∧ y₁.mem = x₁.mem := by
+  obtain ⟨n, C, emit, he⟩ := captured_read
+  let r := 4 + (1 + n + 1)
+  let frontier : Fin r := ⟨0, by omega⟩
+  let src : Fin 3 → Fin r := fun i => ⟨1 + i.val, by omega⟩
+  let tmp : Fin (1 + n + 1) → Fin r := fun i => ⟨4 + i.val, by omega⟩
+  let result := tmp 0
+  have hi : Function.Injective tmp := by
+    intro i j h
+    apply Fin.ext
+    have := congrArg Fin.val h
+    dsimp only [tmp] at this
+    omega
+  have hd (i : Fin 3) (j : Fin (1 + n + 1)) : src i ≠ tmp j := by
+    intro h
+    have := congrArg Fin.val h
+    dsimp only [src, tmp] at this
+    have := i.isLt
+    omega
+  have hf (j : Fin (1 + n + 1)) : tmp j ≠ frontier := by
+    intro h
+    have := congrArg Fin.val h
+    dsimp only [tmp, frontier] at this
+    omega
+  have hs (i : Fin 3) : src i ≠ frontier := by
+    intro h
+    have := congrArg Fin.val h
+    dsimp only [src, frontier] at this
+    omega
+  let d₀ : Fin 3 → ℕ := ![1, 0, 5]
+  let d₁ : Fin 3 → ℕ := ![2, 2, 0]
+  let x₀ : Data r := ⟨fun i => if i.val = 0 then 3 else
+    if i.val = 1 then 1 else if i.val = 3 then 5 else 0, fun _ => 7⟩
+  let x₁ : Data r := ⟨fun i => if i.val = 0 then 3 else
+    if i.val = 1 ∨ i.val = 2 then 2 else 0, fun _ => 7⟩
+  have hx₀ : x₀.Bounded 3 := by
+    constructor
+    · intro i
+      change Fits 3 (if i.val = 0 then 3 else
+        if i.val = 1 then 1 else if i.val = 3 then 5 else 0)
+      split_ifs <;> decide
+    · intro a
+      change Fits 3 7
+      decide
+  have hx₁ : x₁.Bounded 3 := by
+    constructor
+    · intro i
+      change Fits 3 (if i.val = 0 then 3 else if i.val = 1 ∨ i.val = 2 then 2 else 0)
+      split_ifs <;> decide
+    · intro a
+      change Fits 3 7
+      decide
+  have hin₀ : x₀.regs ∘ src = d₀ := by funext i; fin_cases i <;> rfl
+  have hin₁ : x₁.regs ∘ src = d₁ := by funext i; fin_cases i <;> rfl
+  obtain ⟨hnw, hrun⟩ := he frontier src tmp hi hd hf hs
+  obtain ⟨y₀, hy₀, hm₀, ho₀, _⟩ := hrun 3 d₀ x₀ hx₀ hin₀ trivial (by trivial)
+  obtain ⟨y₁, hy₁, hm₁, ho₁, _⟩ := hrun 3 d₁ x₁ hx₁ hin₁ trivial (by
+    change 2 < 3 ∧ 2 < 3 ∧ True
+    decide)
+  have hc₀ : memoryCost (iterate (fun _ : Unit × ℕ => Memory.read (d₀ 0)) () (d₀ 1)
+      (d₀ 2)) (memoryState frontier x₀) = 0 := rfl
+  have hc₁ : memoryCost (iterate (fun _ : Unit × ℕ => Memory.read (d₁ 0)) () (d₁ 1)
+      (d₁ 2)) (memoryState frontier x₁) = 2 := rfl
+  have hy₀' : ExecWithin 3 (emit frontier src tmp) x₀ y₀ C := by
+    simpa only [hc₀, Nat.zero_add, Nat.add_zero, Nat.mul_one]
+      using hy₀
+  have hy₁' : ExecWithin 3 (emit frontier src tmp) x₁ y₁ (3 * C) := by
+    simpa only [hc₁, Nat.add_zero, Nat.reduceAdd, Nat.mul_comm]
+      using hy₁
+  exact ⟨r, C, emit frontier src tmp, result, x₀, y₀, x₁, y₁, hnw, hx₀, hx₁, hy₀',
+    ho₀ 0, congrArg Memory.State.mem hm₀, hy₁', ho₁ 0, congrArg Memory.State.mem hm₁⟩
+
+end Cslib.WordRAM.Realization.Tests.IterateCapture