chore: add tracing for heartbeat usage of sym_n, refactor the way (…

…non-)effect hyps are tracked (leanprover#145)

### Description:

Add a new `Tactic.sym.heartbeat` trace option, which will print the
heartbeat usage of `sym_n`.

Also refactors the way (non-)effect theorems are tracked in the
`SymContext`, by storing a full `Simp.Context`, instead of only an array
of free variables. This makes it easy to do simp-based aggregation at
intermediate steps as well, if we want to do so.

### Testing:

`make all` ran locally

### License:

By submitting this pull request, I confirm that my contribution is
made under the terms of the Apache 2.0 license.

---------

Co-authored-by: Shilpi Goel <shigoel@gmail.com>

Tactics/Attr.lean

@@ -7,6 +7,9 @@ initialize
   -- enable tracing for `sym_n` tactic and related components
   registerTraceClass `Tactic.sym
 
+  -- enable tracing for heartbeat usage of `sym_n`
+  registerTraceClass `Tactic.sym.heartbeats
+
   -- enable extra checks for debugging `sym_n`,
   -- see `AxEffects.validate` for more detail on what is being type-checked
   registerOption `Tactic.sym.debug {

Tactics/Common.lean

@@ -259,3 +259,13 @@ def mkEqArmState (x y : Expr) : Expr :=
 /-- Return a proof of type `x = x`, where `x : ArmState` -/
 def mkEqReflArmState (x : Expr) : Expr :=
   mkApp2 (.const ``Eq.refl [1]) mkArmState x
+
+/-! ## Tracing helpers -/
+
+def traceHeartbeats (cls : Name) (header : Option String := none) :
+    MetaM Unit := do
+  let header := (header.map (· ++ ": ")).getD ""
+  let heartbeats ← IO.getNumHeartbeats
+  let percent ← heartbeatsPercent
+  trace cls fun _ =>
+    m!"{header}used {heartbeats} heartbeats ({percent}% of maximum)"

Tactics/Reflect/AxEffects.lean

@@ -565,65 +565,124 @@ by default):
 - `eff.programProof`, and
 - `eff.stackAlignmentProof?` (if the field is not `none`)
 
-Return a list of fvar ids of the newly added hypotheses -/
+Return an `AxEffect` where the expressions mentioned above have been replaced by
+`Expr.fvar <fvarId>`, with `fvarId` the id of the corresponding hypothesis
+that was just added to the local context -/
 def addHypothesesToLContext (eff : AxEffects) (hypPrefix : String := "h_")
     (mvar : Option MVarId := none) :
-    TacticM (List FVarId) :=
+    TacticM AxEffects :=
   let msg := m!"adding hypotheses to local context"
   withTraceNode `Tactic.sym (fun _ => pure msg) do
     eff.traceCurrentState
     let mut goal ← mvar.getDM getMainGoal
-    let mut fvars := []
 
-    for ⟨field, {proof, ..}⟩ in eff.fields do
-      let msg := m!"adding field {field}"
-      ⟨fvars, goal⟩ ← withTraceNode `Tactic.sym (fun _ => pure msg) <| goal.withContext do
-          trace[Tactic.sym] "raw proof: {proof}"
-          let name := Name.mkSimple (s!"{hypPrefix}{field}")
-          replaceOrNote fvars goal name proof
+    let fields ← do
+      let mut fields := []
+      for ⟨field, {value, proof}⟩ in eff.fields do
+        let msg := m!"adding field {field}"
+        let ⟨fvar, goal'⟩ ← withTraceNode `Tactic.sym (fun _ => pure msg) <| goal.withContext do
+            trace[Tactic.sym] "raw proof: {proof}"
+            let name := Name.mkSimple s!"{hypPrefix}{field}"
+            replaceOrNote goal name proof
+        goal := goal'
+        let fieldEff := FieldEffect.mk value (Expr.fvar fvar)
+        fields := (field, fieldEff) :: fields
+      pure (Std.HashMap.ofList fields)
 
     trace[Tactic.sym] "adding non-effects with {eff.nonEffectProof}"
-    ⟨fvars, goal⟩ ← goal.withContext do
+    let ⟨nonEffectProof, goal'⟩ ← goal.withContext do
       let name := .mkSimple s!"{hypPrefix}non_effects"
       let proof := eff.nonEffectProof
-      replaceOrNote fvars goal name proof
+      replaceOrNote goal name proof
+    let nonEffectProof := Expr.fvar nonEffectProof
+    goal := goal'
 
     trace[Tactic.sym] "adding memory effects with {eff.memoryEffectProof}"
-    ⟨fvars, goal⟩ ← goal.withContext do
+    let ⟨memoryEffectProof, goal'⟩ ← goal.withContext do
       let name := .mkSimple s!"{hypPrefix}memory_effects"
       let proof := eff.memoryEffectProof
-      replaceOrNote fvars goal name proof
+      replaceOrNote goal name proof
+    let memoryEffectProof := Expr.fvar memoryEffectProof
+    goal := goal'
 
     trace[Tactic.sym] "adding program hypothesis with {eff.programProof}"
-    ⟨fvars, goal⟩ ← goal.withContext do
+    let ⟨programProof, goal'⟩ ← goal.withContext do
       let name := .mkSimple s!"{hypPrefix}program"
       let proof := eff.programProof
-      replaceOrNote fvars goal name proof
+      replaceOrNote goal name proof
+    let programProof := Expr.fvar programProof
+    goal := goal'
 
     trace[Tactic.sym] "stackAlignmentProof? is {eff.stackAlignmentProof?}"
-    if let some stackAlignmentProof := eff.stackAlignmentProof? then
-      trace[Tactic.sym] "adding stackAlignment hypothesis"
-      ⟨fvars, goal⟩ ← goal.withContext do
-        let name := .mkSimple s!"{hypPrefix}sp_aligned"
-        replaceOrNote fvars goal name stackAlignmentProof
-    else
-      trace[Tactic.sym] "skipping stackAlignment hypothesis"
+    let (stackAlignmentProof?, goal') ← do
+      if let some stackAlignmentProof := eff.stackAlignmentProof? then
+        trace[Tactic.sym] "adding stackAlignment hypothesis"
+        let ⟨fvar, goal'⟩ ← goal.withContext do
+          let name := .mkSimple s!"{hypPrefix}sp_aligned"
+          replaceOrNote goal name stackAlignmentProof
+        pure (some (Expr.fvar fvar), goal')
+      else
+        trace[Tactic.sym] "skipping stackAlignment hypothesis"
+        pure (none, goal)
+    goal := goal'
 
     replaceMainGoal [goal]
-    return fvars
+    return {eff with
+      fields, nonEffectProof, memoryEffectProof, programProof,
+      stackAlignmentProof?
+    }
 where
-  replaceOrNote (fvars : List FVarId) (goal : MVarId)
+  replaceOrNote (goal : MVarId)
       (h : Name) (v : Expr) (t? : Option Expr := none) :
-      MetaM (List FVarId × MVarId) :=
+      MetaM (FVarId × MVarId) :=
     let msg := m!"adding {h} to the local context"
     withTraceNode `Tactic.sym (fun _ => pure msg) <| do
       trace[Tactic.sym] "with value {v} and type {t?}"
       if let some decl := (← getLCtx).findFromUserName? h then
         let ⟨fvar, goal, _⟩ ← goal.replace decl.fvarId v t?
-        return ⟨fvar :: fvars, goal⟩
+        return ⟨fvar, goal⟩
       else
         let ⟨fvar, goal⟩ ← goal.note h v t?
-        return ⟨fvar :: fvars, goal⟩
+        return ⟨fvar, goal⟩
+
+/-- Return a `SimpTheorems` with the proofs contained in the given `AxEffects`
+
+Note this adds *only* the (non-)effect proofs, to get a simp configuration with
+more default simp-lemmas, see `AxEffects.toSimp` -/
+def toSimpTheorems (eff : AxEffects) : MetaM SimpTheorems := do
+  let msg := m!"computing SimpTheorems for (non-)effect hypotheses"
+  withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+    let lctx ← getLCtx
+    let baseName? :=
+      if eff.currentState.isFVar then
+        (lctx.find? eff.currentState.fvarId!).map fun decl =>
+          Name.str decl.userName "AxEffects"
+      else
+        none
+    let baseName := baseName?.getD (Name.mkSimple "AxEffects")
+
+    let add (thms : SimpTheorems) (e : Expr) (name : String) := do
+      trace[Tactic.sym] "adding {e} with name {name}"
+      let origin : Origin :=
+        if e.isFVar then
+          .fvar e.fvarId!
+        else
+          .other <| Name.str baseName name
+      thms.add origin #[] e
+
+    let mut thms : SimpTheorems := {}
+
+    for ⟨field, {proof, ..}⟩ in eff.fields do
+      thms ← add thms proof s!"field_{field}"
+
+    thms ← add thms eff.nonEffectProof "nonEffectProof"
+    thms ← add thms eff.memoryEffectProof "memoryEffectProof"
+    thms ← add thms eff.programProof "programProof"
+    if let some stackAlignmentProof := eff.stackAlignmentProof? then
+      thms ← add thms stackAlignmentProof "stackAlignmentProof"
+
+    trace[Tactic.sym] "done"
 
+    pure thms
 
 end Tactic

Tactics/Sym.lean

@@ -23,6 +23,10 @@ private def evalTacticAndTrace (tactic : TSyntax `tactic) : TacticM Unit :=
     evalTactic tactic
     trace[Tactic.sym] "new goal state:\n{← getGoals}"
 
+private def Sym.traceHeartbeats (header : Option String := none) :=
+  _root_.traceHeartbeats `Tactic.sym.heartbeats header
+open Sym (traceHeartbeats)
+
 /-- `init_next_step h_run stepi_eq sn` splits the hypothesis
   `h_run: s_final = run (n+1) s`
 by adding a new state variable, `sn`, and two new hypotheses:
@@ -248,12 +252,19 @@ def explodeStep (c : SymContext) (hStep : Expr) : TacticM SymContext :=
                 #[eff.currentState, spEff.value, spEff.proof, hAligned]
             pure { eff with stackAlignmentProof? }
 
-    let axHyps ← withMainContext <| do
+    -- Add new (non-)effect hyps to the context
+    let simpThms ← withMainContext <| do
       if ←(getBoolOption `Tactic.sym.debug) then
         eff.validate
 
-      eff.addHypothesesToLContext s!"h_{c.next_state}_"
-    let c := { c with axHyps := axHyps ++ c.axHyps}
+      let eff ← eff.addHypothesesToLContext s!"h_{c.next_state}_"
+      withMainContext <| eff.toSimpTheorems
+
+    -- Add the new (non-)effect hyps to the aggregation simp context
+    let aggregateSimpCtx := { c.aggregateSimpCtx with
+      simpTheorems := c.aggregateSimpCtx.simpTheorems.push simpThms
+    }
+    let c := { c with aggregateSimpCtx}
 
     -- Attempt to reflect the new PC
     let nextPc ← eff.getField .PC
@@ -341,6 +352,7 @@ def sym1 (c : SymContext) (whileTac : TSyntax `tactic) : TacticM SymContext :=
       let goal ← goal.clear hStep.fvarId
       replaceMainGoal [goal]
 
+      traceHeartbeats
       return c
 
 /- used in `sym_n` tactic to specify an initial state -/
@@ -373,6 +385,8 @@ in which case a new goal of the appropriate type will be added.
 The other hypotheses *must* be present,
 since we infer required information from their types. -/
 elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
+  traceHeartbeats "initial heartbeats"
+
   let s ← s.mapM fun
     | `(sym_at|at $s:ident) => pure s.getId
     | _ => Lean.Elab.throwUnsupportedSyntax
@@ -433,6 +447,7 @@ Did you remember to generate step theorems with:
     for _ in List.range n do
       c ← sym1 c whileTac
 
+    traceHeartbeats "symbolic simulation total"
     -- Check if we can substitute the final state
     if c.runSteps? = some 0 then
       let msg := do
@@ -453,18 +468,16 @@ Did you remember to generate step theorems with:
 
         let goal ← subst goal hEqId
         trace[Tactic.sym] "performed subsitutition in:\n{goal}"
+        traceHeartbeats
 
         replaceMainGoal [goal]
 
     -- Rudimentary aggregation: we feed all the axiomatic effect hypotheses
     -- added while symbolically evaluating to `simp`
-    withMainContext <| do
-      let lctx ← getLCtx
-      let some axHyps := c.axHyps.toArray.mapM lctx.find?
-        | throwError "internal error: one of the following fvars could not \
-            be found:\n  {c.axHyps.map Expr.fvar}"
-      let (ctx, simprocs) ← LNSymSimpContext
-          (config := {decide := true, failIfUnchanged := false})
-          (decls := axHyps)
-      let goal? ← LNSymSimp (← getMainGoal) ctx simprocs
+    let msg := m!"aggregating (non-)effects"
+    withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext do
+      traceHeartbeats "pre"
+      let goal? ← LNSymSimp (← getMainGoal) c.aggregateSimpCtx c.aggregateSimprocs
       replaceMainGoal goal?.toList
+
+    traceHeartbeats "final usage"

Tactics/SymContext.lean

@@ -11,6 +11,7 @@ import Tactics.Common
 import Tactics.Attr
 import Tactics.Reflect.ProgramInfo
 import Tactics.Reflect.AxEffects
+import Tactics.Simp
 
 /-!
 This files defines the `SymContext` structure,
@@ -82,9 +83,14 @@ structure SymContext where
   `CheckSPAlignment state` -/
   h_sp?  : Option Name
 
-  /-- The list of all axiomatic-effect hypotheses added by `AxEffect`
-  throughout the current `sym_n` run -/
-  axHyps : List FVarId
+  /-- The `simp` context used for effect aggregation.
+  This collects references to all (non-)effect hypotheses of the intermediate
+  states, together with sensible default simp-lemmas used for
+  effect aggregation -/
+  aggregateSimpCtx : Simp.Context
+  /-- Simprocs used for aggregation. This is stored for performance benefits,
+  but should not be modified during the course of a `sym_n` call -/
+  aggregateSimprocs : Simp.SimprocsArray
 
   /-- `state_prefix` is used together with `curr_state_number`
   to determine the name of the next state variable that is added by `sym` -/
@@ -270,8 +276,11 @@ def fromLocalContext (state? : Option Name) : MetaM SymContext := do
           #generateStepEqTheorems {program}"
 
   -- Initialize the axiomatic hypotheses with hypotheses for the initial state
-  let axHyps := [h_program, h_pc] ++ h_err?.toList ++ h_sp?.toList
-  let axHyps := axHyps.map (·.fvarId)
+  let axHyps := #[h_program, h_pc] ++ h_err?.toArray ++ h_sp?.toArray
+  let (aggregateSimpCtx, aggregateSimprocs) ←
+    LNSymSimpContext
+      (config := {decide := true, failIfUnchanged := false})
+      (decls := axHyps)
 
   return inferStatePrefixAndNumber {
     state, finalState, runSteps?, program, pc,
@@ -281,7 +290,7 @@ def fromLocalContext (state? : Option Name) : MetaM SymContext := do
     h_err? := (·.userName) <$> h_err?,
     h_sp? := (·.userName) <$> h_sp?,
     programInfo,
-    axHyps
+    aggregateSimpCtx, aggregateSimprocs
   }
 where
   findLocalDeclOfType? (expectedType : Expr) : MetaM (Option LocalDecl) := do
@@ -402,19 +411,13 @@ def next (c : SymContext) (nextPc? : Option (BitVec 64) := none) :
     SymContext :=
   let curr_state_number := c.curr_state_number + 1
   let s := c.next_state
-  {
+  { c with
     state       := s
-    finalState  := c.finalState
-    h_run       := c.h_run
     h_program   := .mkSimple s!"h_{s}_program"
     h_pc        := .mkSimple s!"h_{s}_pc"
     h_err?      := c.h_err?.map (fun _ => .mkSimple s!"h_{s}_err")
     h_sp?       := c.h_sp?.map (fun _ => .mkSimple s!"h_{s}_sp_aligned")
     runSteps?   := (· - 1) <$> c.runSteps?
-    program     := c.program
-    programInfo := c.programInfo
     pc          := nextPc?.getD (c.pc + 4#64)
-    axHyps      := c.axHyps
     curr_state_number
-    state_prefix := c.state_prefix
   }