Prove that analysis results apply to all states, not just the final one

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

lean/Main.lean

@@ -26,7 +26,7 @@ def testCodeCond₂ : Stmt := [obj_stmt|
   if var { x := 1 } else { noop }
 ]
 
-def testProgram : Program := ⟨testCode⟩
+def testProgram : Program := { rootStmt := testCode }
 
 end Spa
 

lean/Spa/Analysis/Constant.lean

@@ -137,6 +137,13 @@ theorem analyze_correct {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
     ⟦ variablesAt prog.finalState (result ConstLattice prog) ⟧ ρ () :=
   Forward.analyze_correct ConstLattice prog hrun
 
+theorem analyze_correct_at {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
+    {s : prog.State} {ρin ρout : Env} {stin stout : PUnit}
+    (hr : Reaches (L := ConstLattice) (prog.trace hrun) PUnit.unit s ρin ρout stin stout) :
+    ⟦ joinForKey s (result ConstLattice prog) ⟧ ρin stin
+      ∧ ⟦ variablesAt s (result ConstLattice prog) ⟧ ρout stout :=
+  Forward.analyze_correct_at ConstLattice prog hrun hr
+
 end ConstAnalysis
 
 end Spa

lean/Spa/Analysis/Forward.lean

@@ -99,6 +99,28 @@ noncomputable def stepTraceState :
   | s₁, _, _, _, .edge hnode _ subtr, st =>
       stepTraceState subtr (stepNode s₁ hnode st)
 
+/-- `Reaches tr st₁ s ρin ρout stin stout` witnesses that, when the trace `tr`
+    (starting at state `st₁`) is executed, node `s` is visited at some point: `ρin`
+    and `ρout` are the concrete environments just before and after `s`'s basic block
+    runs, and `stin`/`stout` are the corresponding abstract execution states. A node
+    inside a loop is reached once per iteration, each with its own environments. -/
+inductive Reaches : {s₁ s₂ : prog.State} → {ρ₁ ρ₂ : Env} →
+    Trace prog.cfg s₁ s₂ ρ₁ ρ₂ → S.St ρ₁ →
+    (s : prog.State) → (ρin ρout : Env) → S.St ρin → S.St ρout → Prop
+  | single_here {s₁ : prog.State} {ρ₁ ρ₂ : Env} {st₁ : S.St ρ₁}
+      (hnode : EvalBasicStmtOpt ρ₁ (prog.code s₁) ρ₂) :
+      Reaches (.single hnode) st₁ s₁ ρ₁ ρ₂ st₁ (stepNode s₁ hnode st₁)
+  | edge_here {s₁ s₂ s₃ : prog.State} {ρ₁ ρ₂ ρ₃ : Env} {st₁ : S.St ρ₁}
+      (hnode : EvalBasicStmtOpt ρ₁ (prog.code s₁) ρ₂)
+      (hedge : (s₁, s₂) ∈ prog.cfg.edges) (rest : Trace prog.cfg s₂ s₃ ρ₂ ρ₃) :
+      Reaches (.edge hnode hedge rest) st₁ s₁ ρ₁ ρ₂ st₁ (stepNode s₁ hnode st₁)
+  | edge_there {s₁ s₂ s₃ : prog.State} {ρ₁ ρ₂ ρ₃ : Env} {st₁ : S.St ρ₁}
+      (hnode : EvalBasicStmtOpt ρ₁ (prog.code s₁) ρ₂)
+      (hedge : (s₁, s₂) ∈ prog.cfg.edges) (rest : Trace prog.cfg s₂ s₃ ρ₂ ρ₃)
+      {s : prog.State} {ρin ρout : Env} {stin : S.St ρin} {stout : S.St ρout} :
+      Reaches rest (stepNode s₁ hnode st₁) s ρin ρout stin stout →
+      Reaches (.edge hnode hedge rest) st₁ s ρin ρout stin stout
+
 omit [DecidableEq L] in
 lemma evalStmtOrNone_valid {s : prog.State} {ρ₁ ρ₂ : Env} {st : S.St ρ₁}
     {vs : VariableValues L prog} (o : Option BasicStmt) (hco : prog.code s = o)
@@ -126,20 +148,46 @@ lemma stepTrace {s₁ : prog.State} {ρ₁ ρ₂ : Env} {st : S.St ρ₁}
   rw [variablesAt_joinAll]
   exact hjoin
 
+/-- Soundness at *every* visited node: if the analysis result over-approximates the
+    incoming environment at the start of the trace, then at each node reached along the
+    way it over-approximates both the environment entering that node (via `joinForKey`)
+    and the environment leaving it (via `variablesAt`). The intermediate `variablesAt`
+    evidence used to be computed and discarded inside `walkTrace`; here it is returned. -/
+lemma walkTrace_reaches {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env} {st₁ : S.St ρ₁}
+    {s : prog.State} {ρin ρout : Env} {stin : S.St ρin} {stout : S.St ρout}
+    {tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂}
+    (hr : Reaches tr st₁ s ρin ρout stin stout)
+    (hjoin : ⟦ joinForKey s₁ (result L prog) ⟧ ρ₁ st₁) :
+    ⟦ joinForKey s (result L prog) ⟧ ρin stin
+      ∧ ⟦ variablesAt s (result L prog) ⟧ ρout stout := by
+  induction hr with
+  | single_here hnode => exact ⟨hjoin, stepTrace hjoin hnode⟩
+  | edge_here hnode hedge rest => exact ⟨hjoin, stepTrace hjoin hnode⟩
+  | edge_there hnode hedge rest hr' ih =>
+    have hstep := stepTrace hjoin hnode
+    have hmem := FiniteMap.mem_valuesAt prog.states_nodup
+      (prog.mem_incoming_of_edge hedge) (variablesAt_mem _ (result L prog))
+    exact ih (interp_foldr hstep hmem)
+
+omit [DecidableEq L] in
+/-- The final node of a trace is always reached, with the environment/state the trace
+    ends in. Used to recover the final-state soundness theorem from `walkTrace_reaches`. -/
+lemma reaches_final {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env} (st₁ : S.St ρ₁)
+    (tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂) :
+    ∃ ρin, ∃ stin : S.St ρin,
+      Reaches tr st₁ s₂ ρin ρ₂ stin (stepTraceState tr st₁) := by
+  induction tr with
+  | single hnode => exact ⟨_, _, .single_here hnode⟩
+  | edge hnode hedge rest ih =>
+    obtain ⟨ρin, stin, hr⟩ := ih (stepNode _ hnode st₁)
+    exact ⟨ρin, stin, .edge_there hnode hedge rest hr⟩
+
 lemma walkTrace {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env} {st₁ : S.St ρ₁}
     (hjoin : ⟦ joinForKey s₁ (result L prog) ⟧ ρ₁ st₁)
     (tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂) :
     ⟦ variablesAt s₂ (result L prog) ⟧ ρ₂ (stepTraceState tr st₁) := by
-  induction tr with
+  obtain ⟨_, _, hr⟩ := reaches_final st₁ tr
-  | single hnode => exact stepTrace hjoin hnode
+  exact (walkTrace_reaches hr hjoin).2
-  | @edge _ ρ' _ i₁ i₂ _ hnode hedge _ ih =>
-    have hstep : ⟦ variablesAt i₁ (result L prog) ⟧ ρ' (stepNode i₁ hnode st₁) :=
-      stepTrace hjoin hnode
-    have hmem : variablesAt i₁ (result L prog)
-        ∈ (result L prog).valuesAt (prog.incoming i₂) :=
-      FiniteMap.mem_valuesAt prog.states_nodup
-        (prog.mem_incoming_of_edge hedge) (variablesAt_mem i₁ (result L prog))
-    exact ih (interp_foldr hstep hmem)
 
 variable (L prog) in
 theorem analyze_correct_state {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
@@ -149,6 +197,20 @@ theorem analyze_correct_state {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
   rw [joinForKey_initialState]
   exact ValidStateEvaluator.botV_init
 
+variable (L prog) in
+/-- Soundness at every program point reached during execution: for any node `s` visited
+    by the run `hrun` (witnessed by `hr`), the analysis result over-approximates both the
+    environment entering `s` and the one leaving it. The final-state theorem
+    `analyze_correct_state` is the special case where `s` is `prog.finalState`. -/
+theorem analyze_correct_at {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
+    {s : prog.State} {ρin ρout : Env} {stin : S.St ρin} {stout : S.St ρout}
+    (hr : Reaches (prog.trace hrun) S.init s ρin ρout stin stout) :
+    ⟦ joinForKey s (result L prog) ⟧ ρin stin
+      ∧ ⟦ variablesAt s (result L prog) ⟧ ρout stout := by
+  refine walkTrace_reaches hr ?_
+  rw [joinForKey_initialState]
+  exact ValidStateEvaluator.botV_init
+
 end
 
 variable (L prog) in

lean/Spa/Analysis/Reaching.lean

@@ -96,6 +96,13 @@ theorem analyze_correct {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
       (stepTraceState (prog.trace hrun) (stateInterp prog).init) :=
   Forward.analyze_correct_state (DefSet prog) prog hrun
 
+theorem analyze_correct_at {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
+    {s : prog.State} {ρin ρout : Env} {stin : Run prog} {stout : Run prog}
+    (hr : Reaches (prog.trace hrun) (stateInterp prog).init s ρin ρout stin stout) :
+    ⟦ joinForKey s (result (DefSet prog) prog) ⟧ ρin stin
+      ∧ ⟦ variablesAt s (result (DefSet prog) prog) ⟧ ρout stout :=
+  Forward.analyze_correct_at (DefSet prog) prog hrun hr
+
 end ReachingAnalysis
 
 end Spa

lean/Spa/Analysis/Sign.lean

@@ -213,6 +213,13 @@ theorem analyze_correct {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
     ⟦ variablesAt prog.finalState (result SignLattice prog) ⟧ ρ () :=
   Forward.analyze_correct SignLattice prog hrun
 
+theorem analyze_correct_at {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
+    {s : prog.State} {ρin ρout : Env} {stin stout : PUnit}
+    (hr : Reaches (L := SignLattice) (prog.trace hrun) PUnit.unit s ρin ρout stin stout) :
+    ⟦ joinForKey s (result SignLattice prog) ⟧ ρin stin
+      ∧ ⟦ variablesAt s (result SignLattice prog) ⟧ ρout stout :=
+  Forward.analyze_correct_at SignLattice prog hrun hr
+
 end SignAnalysis
 
 end Spa