From eb5ca196111a23ffa80516789ed6c5c33199ba43 Mon Sep 17 00:00:00 2001
From: Mikael Mayer <mimayere@amazon.com>
Date: Mon, 30 Mar 2026 13:03:51 +0000
Subject: [PATCH] feat: define CachedAnalyses invariant, update PrecondElim to
 maintain call graph (#702)

PrecondElim generates WF procedures but did not update the cached call
graph, which prevented FilterProcedures from being used after PrecondElim
in the selective verification pipeline.

- Add CallGraph.addLeafNode for inserting nodes with no callees
- Update PrecondElim to register generated WF procedures in the call graph
- Move the second FilterProcedures.run after PrecondElim in the Verifier
- Document the CachedAnalyses call graph invariant
---
 Strata/Languages/Core/CallGraph.lean |  4 ++++
 Strata/Languages/Core/Verifier.lean  |  2 +-
 Strata/Transform/CoreTransform.lean  |  8 +++++++-
 Strata/Transform/PrecondElim.lean    | 16 ++++++++++++++--
 4 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/Strata/Languages/Core/CallGraph.lean b/Strata/Languages/Core/CallGraph.lean
index 2a544ca7e..46f58dab7 100644
--- a/Strata/Languages/Core/CallGraph.lean
+++ b/Strata/Languages/Core/CallGraph.lean
@@ -64,6 +64,10 @@ def CallGraph.decrementEdge (cg : CallGraph) (caller : String) (callee : String)
     callers := new_callers
   }
 
+/-- Add a node with no callees to the call graph. -/
+def CallGraph.addLeafNode (cg : CallGraph) (name : String) : CallGraph :=
+  { cg with callees := cg.callees.insert name {} }
+
 /-- Compute transitive closure of callees; the result does not contain `name`. -/
 partial def CallGraph.getCalleesClosure (cg : CallGraph) (name : String) : List String :=
   let rec go (visited : List String) (toVisit : List String) : List String :=
diff --git a/Strata/Languages/Core/Verifier.lean b/Strata/Languages/Core/Verifier.lean
index 732e29b1a..d48346f86 100644
--- a/Strata/Languages/Core/Verifier.lean
+++ b/Strata/Languages/Core/Verifier.lean
@@ -715,8 +715,8 @@ def verify (program : Program)
       let passes := fun prog => do
         let prog ← FilterProcedures.run prog procs
         let (_changed,prog) ← CallElim.callElim' prog
-        let prog ← FilterProcedures.run prog procs
         let prog ← runPrecondElim prog
+        let prog ← FilterProcedures.run prog procs
         return prog
       let res := Transform.run program passes
       match res with
diff --git a/Strata/Transform/CoreTransform.lean b/Strata/Transform/CoreTransform.lean
index bc3cafcbe..0df61784c 100644
--- a/Strata/Transform/CoreTransform.lean
+++ b/Strata/Transform/CoreTransform.lean
@@ -85,7 +85,13 @@ def isGlobalVar (p : Program) (ident : Expression.Ident) : Bool :=
 
 
 /-- Cached results of program analyses that are helpful for program
-    transformation. -/
+    transformation.
+
+    Invariant: When `callGraph` is `some cg`, `cg` must reflect the current
+    program's procedure call structure. Every procedure in the program must
+    have an entry in `cg.callees`, and every call edge must be accounted for.
+    Transforms that add, remove, or modify procedures must update the call
+    graph accordingly (or set it to `.none` to invalidate it). -/
 structure CachedAnalyses where
   callGraph: Option CallGraph := .none
 
diff --git a/Strata/Transform/PrecondElim.lean b/Strata/Transform/PrecondElim.lean
index afdd200dc..ab7d7f1b6 100644
--- a/Strata/Transform/PrecondElim.lean
+++ b/Strata/Transform/PrecondElim.lean
@@ -273,6 +273,14 @@ end
 
 /-! ## Main transformation -/
 
+/-- Add a WF procedure as a leaf node in the cached call graph.
+WF procedures contain only assert/assume statements and make no procedure calls. -/
+private def addWFProcToCallGraph (name : String) : CoreTransformM Unit :=
+  modify fun σ => match σ.cachedAnalyses.callGraph with
+  | .some cg => { σ with cachedAnalyses := { σ.cachedAnalyses with
+      callGraph := .some (cg.addLeafNode name) } }
+  | .none => σ
+
 /--
 Transform an entire program:
 1. For each procedure, transform its body and if needed generate a WF procedure
@@ -301,7 +309,9 @@ where
         let procDecl := Decl.proc proc' md
         let (changed', rest') ← transformDecls rest
         match mkContractWFProc F proc with
-        | some wfDecl => return (true, wfDecl :: procDecl :: rest')
+        | some wfDecl => do
+          addWFProcToCallGraph (wfProcName (CoreIdent.toPretty proc.header.name))
+          return (true, wfDecl :: procDecl :: rest')
         | none => return (changed || changed', procDecl :: rest')
       | .func func md => do
         let F ← getFactory
@@ -312,7 +322,9 @@ where
         let hasPreconds := !func.preconditions.isEmpty
         let (changed, rest') ← transformDecls rest
         match mkFuncWFProc F' func with
-        | some wfDecl => return (true, wfDecl :: funcDecl :: rest')
+        | some wfDecl => do
+          addWFProcToCallGraph (wfProcName (CoreIdent.toPretty func.name))
+          return (true, wfDecl :: funcDecl :: rest')
         | none => return (changed || hasPreconds, funcDecl :: rest')
       | .type (.data block) _ => do
         let F ← getFactory