From 7c05adadff93bdb71f60e903b2569c6c73bc7dfb Mon Sep 17 00:00:00 2001 From: Danila Fedorin Date: Tue, 29 Sep 2026 20:56:01 -0500 Subject: [PATCH] Restore constant folding and expression correctness without trace experiments --- lean/Spa.lean | 1 + lean/Spa/Transformation/Constant.lean | 131 ++++++++++++++++++++++++++ 2 files changed, 132 insertions(+) create mode 100644 lean/Spa/Transformation/Constant.lean diff --git a/lean/Spa.lean b/lean/Spa.lean index 3bc5711..2d605ac 100644 --- a/lean/Spa.lean +++ b/lean/Spa.lean @@ -21,3 +21,4 @@ import Spa.Analysis.Sign import Spa.Analysis.Constant import Spa.Analysis.Reaching import Spa.Transformation.Licm +import Spa.Transformation.Constant diff --git a/lean/Spa/Transformation/Constant.lean b/lean/Spa/Transformation/Constant.lean new file mode 100644 index 0000000..d1d8dcb --- /dev/null +++ b/lean/Spa/Transformation/Constant.lean @@ -0,0 +1,131 @@ +import Spa.Language.Base +import Spa.Language.Program +import Spa.Lattice.FiniteMap +import Spa.Analysis.Constant +import Spa.Analysis.Forward + +/-! +# Constant folding + +Rewrites each assignment's right-hand side to a literal wherever the constant +analysis (`Spa/Analysis/Constant.lean`) pins its value down. + +The traversal recurses over the plain `Stmt`, threading a `GGraph.Embed` of the +current subtree's CFG into the program's (`Program.rootEmbed`, then one +`Embed.trans` per descent). At an assignment, `Embed.singletonIndex` gives its +CFG state, and the facts to fold with are `Forward.joinForKey` at that state — +the join over predecessors, i.e. the values *entering* the node, which is what +the right-hand side reads. (`Forward.variablesAt` would be the values *leaving* +it, which already include this assignment's own effect.) +-/ + +namespace Spa + +namespace ConstantTransform + +open GGraph Forward + +variable (prog : Program) + +abbrev result := Forward.result ConstLattice prog + +/-- Replace an expression by a literal when the analysis pins its value down, +recursing into its subexpressions otherwise. Whole subexpressions are tried +first, so `(x + 1) - x` folds outright when `x` is known, rather than only in its +leaves. -/ +def foldExpr (vs : VariableValues ConstLattice prog) : Expr → Expr + | .num n => .num n + | .var k => + match ConstAnalysis.eval prog (.var k) vs with + | .mk z => .num z + | _ => .var k + | .add a b => + match ConstAnalysis.eval prog (.add a b) vs with + | .mk z => .num z + | _ => .add (foldExpr vs a) (foldExpr vs b) + | .sub a b => + match ConstAnalysis.eval prog (.sub a b) vs with + | .mk z => .num z + | _ => .sub (foldExpr vs a) (foldExpr vs b) + +/-- Constant-fold every assignment in a statement. + +`sv` is the analysis result, taken as a parameter rather than read from `result` +at each node: it is a fixpoint computation, so recomputing it per assignment +would make folding quadratic in the analysis. + +Guards of `ifElse`/`whileLoop` are deliberately left alone. `Stmt.cfg` gives a +conditional's guard no node at all (`ifElse` overlays the two branches), so there +is no state whose entry facts describe where it is evaluated. A `whileLoop`'s +guard does have a candidate — the loop header `GGraph.loopIn` — but tying the +guard's evaluation environment to that node needs a lemma that does not exist +yet, so folding it here would be an unproven soundness claim. -/ +def foldStmt (sv : StateVariables ConstLattice prog) : + (s : Stmt) → Embed s.cfg prog.cfg → Stmt + | .basic .noop, _ => .basic .noop + | .basic (.assign k v), e => + .basic (.assign k (foldExpr prog (joinForKey e.singletonIndex sv) v)) + | .andThen s₁ s₂, e => + .andThen (foldStmt sv s₁ ((Embed.sequenceLeft s₁.cfg s₂.cfg).trans e)) + (foldStmt sv s₂ ((Embed.sequenceRight s₁.cfg s₂.cfg).trans e)) + | .ifElse cond s₁ s₂, e => + .ifElse cond (foldStmt sv s₁ ((Embed.overlayLeft s₁.cfg s₂.cfg).trans e)) + (foldStmt sv s₂ ((Embed.overlayRight s₁.cfg s₂.cfg).trans e)) + | .whileLoop cond body, e => + .whileLoop cond (foldStmt sv body ((Embed.loop body.cfg).trans e)) + +/-- Constant-fold a whole program, running the analysis once. -/ +def foldProgram : Stmt := foldStmt prog (result prog) prog.rootStmt prog.rootEmbed + +/-! ## Correctness + +Folding preserves meaning *provided the facts folded with actually hold of the +environment folded in*. That proviso is the whole content: `foldExpr` is sound +against any `vs` that over-approximates `ρ`, and it is the analysis engine's job +(`Forward.analyze_correct_at`) to supply such a `vs` at each program point. -/ + +variable {prog} + +/-- If the analysis pins an expression to a constant and its facts hold of `ρ`, +then the expression really does evaluate to that constant. This is +`ValidExprEvaluator` specialised to the `.mk` case, where `interpConst` says +exactly `v = .int z`. -/ +lemma eq_int_of_eval_mk {vs : VariableValues ConstLattice prog} {ρ : Env} + {e : Expr} {v : Value} {z : ℤ} + (hev : EvalExpr ρ e v) (hvs : ⟦vs⟧ ρ) (hz : ConstAnalysis.eval prog e vs = .mk z) : + v = .int z := by + have h := ValidExprEvaluator.valid (L := ConstLattice) (prog := prog) hev hvs + rw [show ExprEvaluator.eval e vs = ConstAnalysis.eval prog e vs from rfl, hz] at h + exact h + +/-- **Expression folding is meaning-preserving.** Whenever `vs` over-approximates +`ρ`, the folded expression evaluates in `ρ` to whatever the original did. -/ +theorem foldExpr_eval {vs : VariableValues ConstLattice prog} {ρ : Env} (hvs : ⟦vs⟧ ρ) : + ∀ {e : Expr} {v : Value}, EvalExpr ρ e v → EvalExpr ρ (foldExpr prog vs e) v := by + intro e + induction e with + | num n => intro v hev; simpa [foldExpr] using hev + | var k => + intro v hev + simp only [foldExpr] + split + · case h_1 z hz => rw [eq_int_of_eval_mk hev hvs hz]; exact EvalExpr.num ρ z + · exact hev + | add a b iha ihb => + intro v hev + simp only [foldExpr] + split + · case h_1 z hz => rw [eq_int_of_eval_mk hev hvs hz]; exact EvalExpr.num ρ z + · cases hev with + | add _ _ z₁ z₂ h₁ h₂ => exact EvalExpr.add ρ _ _ z₁ z₂ (iha h₁) (ihb h₂) + | sub a b iha ihb => + intro v hev + simp only [foldExpr] + split + · case h_1 z hz => rw [eq_int_of_eval_mk hev hvs hz]; exact EvalExpr.num ρ z + · cases hev with + | sub _ _ z₁ z₂ h₁ h₂ => exact EvalExpr.sub ρ _ _ z₁ z₂ (iha h₁) (ihb h₂) + +end ConstantTransform + +end Spa