Remove 'prog.code s = some bs' argument to eval

lean/Spa/Analysis/Reaching.lean

@@ -19,23 +19,23 @@ namespace ReachingAnalysis
 
 variable (prog : Program)
 
-def genSet (s : prog.State) {bs : BasicStmt} (h : prog.code s = some bs) :
-    DefSet prog :=
-  {prog.nodeIdOfNonempty s h}
+def genSet (s : prog.State) : DefSet prog := (prog.nodeIdOf s).elim {} (fun x => {x})
 
-def eval (s : prog.State) :
-    (bs : BasicStmt) → prog.code s = some bs →
-    VariableValues (DefSet prog) prog → VariableValues (DefSet prog) prog
-  | .assign k _, h, vs =>
-    FiniteMap.generalizedUpdate id (fun _ _ => genSet prog s h) [k] vs
-  | .noop, _, vs => vs
+def eval (s : prog.State) (vs : VariableValues (DefSet prog) prog) : VariableValues (DefSet prog) prog :=
+  match prog.code s with
+  | none => vs
+  | some bs =>
+    match bs with
+    | .assign k _ => FiniteMap.generalizedUpdate id (fun _ _ => genSet prog s) [k] vs
+    | .noop => vs
 
-lemma eval_mono (s : prog.State) (bs : BasicStmt) (h : prog.code s = some bs) :
-    Monotone (eval prog s bs h) := by
-  cases bs with
-  | assign k e =>
-    exact FiniteMap.generalizedUpdate_monotone monotone_id (fun _ => monotone_const)
-  | noop => exact monotone_id
+lemma eval_mono (s : prog.State) :
+    Monotone (eval prog s) := by
+  intros vs₁ vs₂ hle
+  unfold eval; split <;> try simpa
+  split <;> try simpa
+  apply FiniteMap.generalizedUpdate_monotone monotone_id (fun _ => monotone_const)
+  assumption
 
 instance stmtEvaluator : StmtEvaluator (DefSet prog) prog :=
   ⟨eval prog, eval_mono prog⟩
@@ -45,18 +45,18 @@ def output : String :=
 
 inductive Run (prog : Program) where
   | nil : Run prog
-  | cons (s : prog.State) (bs : BasicStmt) (hc : prog.code s = some bs)
+  | cons (s : prog.State) (bs : BasicStmt)
       (rest : Run prog) : Run prog
 
 @[aesop unsafe cases]
 inductive LastAssign (prog : Program) (x : String) : Run prog → prog.NodeId → Prop
   | here (s : prog.State) (e : Expr) (hc : prog.code s = some (.assign x e))
       (rest : Run prog) :
-      LastAssign prog x (Run.cons s (.assign x e) hc rest) (prog.nodeIdOfNonempty s hc)
+      LastAssign prog x (Run.cons s (.assign x e) rest) (prog.nodeIdOfNonempty s hc)
   | there (s : prog.State) (bs : BasicStmt) (hc : prog.code s = some bs)
       (rest : Run prog) {n : prog.NodeId} :
       (∀ e, bs ≠ .assign x e) → LastAssign prog x rest n →
-      LastAssign prog x (Run.cons s bs hc rest) n
+      LastAssign prog x (Run.cons s bs rest) n
 
 instance stateInterp : StateInterpretation (DefSet prog) prog where
   St := fun _ => Run prog
@@ -72,22 +72,27 @@ instance stateInterp : StateInterpretation (DefSet prog) prog where
     obtain ⟨a₁, a₂, rfl, h₁, h₂⟩ := FiniteMap.mem_inf hmem
     aesop (add simp Finset.mem_inter)
 
+private def stepAt (s : prog.State) (obs : Option BasicStmt) { ρ₁ ρ₂ : Env} : EvalBasicStmtOpt ρ₁ obs ρ₂ → Run prog → Run prog
+  | .none, rest => rest
+  | .some (bs := bs) _, rest => Run.cons s bs rest
+
 instance validStateEvaluator : ValidStateEvaluator (DefSet prog) prog where
-  step := by intro s _ _ bs hcode _ rest; exact Run.cons s bs hcode rest
+  step := fun s ρ₁ ρ₂ => stepAt prog s (prog.code s)
   valid := by
-    intro s ρ₁ ρ₂ bs vs st hcode hbs hvs
+    simp [StmtEvaluator.eval, eval];
+    intro s ρ₁ ρ₂ vs; generalize prog.code s = obs; intro hst hbs hvs
+    rcases hbs with _ | @⟨_, bs, hbs⟩; try (simpa [stepAt])
     cases hbs with
     | noop => intro x assigners hmem n hla; aesop
     | assign x e v hev =>
-      intro k assigners hmem n hla
-      have hmem2 : (k, assigners) ∈
-          FiniteMap.generalizedUpdate id (fun _ _ => genSet prog s hcode) [x] vs := hmem
+      simp; intro k assigners hmem n hla
       by_cases hx : k = x
       · subst hx
-        have hd := FiniteMap.generalizedUpdate_mem_eq (List.mem_singleton.mpr rfl) hmem2
-        aesop (add simp [genSet, Finset.mem_singleton])
+        have hd := FiniteMap.generalizedUpdate_mem_eq (List.mem_singleton.mpr rfl) hmem
+        rcases hla
+          <;> simp [Program.nodeIdOfNonempty, hd, genSet, Option.get] <;> aesop
       · have hmem' := FiniteMap.generalizedUpdate_not_mem_backward
-          (fun hc => hx (List.mem_singleton.mp hc)) hmem2
+          (fun hc => hx (List.mem_singleton.mp hc)) hmem
         aesop
   botV_init := by intro x assigners _ n hla; cases hla