chore: add more trace nodes, for better profiling data

Tactics/Attr.lean

@@ -9,8 +9,11 @@ open Lean
 initialize
   -- CSE tactic's non-verbose summary logging.
   registerTraceClass `Tactic.cse.summary
+
   -- enable tracing for `sym_n` tactic and related components
   registerTraceClass `Tactic.sym
+  -- enable verbose tracing
+  registerTraceClass `Tactic.sym.debug
 
   -- enable tracing for heartbeat usage of `sym_n`
   registerTraceClass `Tactic.sym.heartbeats

Tactics/Sym.lean

@@ -11,15 +11,15 @@ import Tactics.Sym.Context
 import Lean
 
 open BitVec
-open Lean Meta
-open Lean.Elab.Tactic
+open Lean
+open Lean.Meta Lean.Elab.Tactic
 
 open AxEffects SymContext
 
 /-- A wrapper around `evalTactic` that traces the passed tactic script,
 executes those tactics, and then traces the new goal state -/
 private def evalTacticAndTrace (tactic : TSyntax `tactic) : TacticM Unit :=
-  withTraceNode `Tactic.sym (fun _ => pure m!"running: {tactic}") <| do
+  withTraceNode m!"running: {tactic}" <| do
     evalTactic tactic
     trace[Tactic.sym] "new goal state:\n{← getGoals}"
 
@@ -50,7 +50,8 @@ to add a new local hypothesis in terms of `w` and `write_mem`
   `h_step : ?s' = w _ _ (w _ _ (... ?s))`
 -/
 def stepiTac (stepiEq : Expr) (hStep : Name) : SymReaderM Unit := fun ctx =>
-  withMainContext' do
+  withMainContext' <|
+  withVerboseTraceNode m!"stepiTac: {stepiEq}" <| do
     let pc := (Nat.toDigits 16 ctx.pc.toNat).asString
     --  ^^ The PC in hex
     let stepLemma := Name.str ctx.program s!"stepi_eq_0x{pc}"
@@ -89,8 +90,7 @@ Finally, we use this proof to change the type of `h_run` accordingly.
 -/
 def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := do
   let c ← readThe SymContext
-  let msg := m!"unfoldRun (runSteps? := {c.runSteps?})"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <|
+  withTraceNode m!"unfoldRun (runSteps? := {c.runSteps?})" <|
   match c.runSteps? with
     | some (_ + 1) => do
         trace[Tactic.sym] "runSteps is statically known to be non-zero, \
@@ -124,9 +124,9 @@ def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := do
         let subGoal ← mkFreshMVarId
         let _ ← mkFreshExprMVarWithId subGoal subGoalTy
 
-        let msg := m!"runSteps is not statically known, so attempt to prove:\
-          {subGoal}"
-        withTraceNode `Tactic.sym (fun _ => pure msg) <| subGoal.withContext <| do
+        withTraceNode m!"runSteps is not statically known, so attempt to prove:\
+          {subGoal}" <|
+        subGoal.withContext <| do
           setGoals [subGoal]
           whileTac () -- run `whileTac` to attempt to close `subGoal`
 
@@ -166,29 +166,32 @@ add the relevant hypotheses to the local context, and
 store an `AxEffects` object with the newly added variables in the monad state
 -/
 def explodeStep (hStep : Expr) : SymM Unit :=
-  withMainContext' do
+  withMainContext' <|
+  withTraceNode m!"explodeStep {hStep}" <| do
     let c ← getThe SymContext
     let oldEff ← getThe AxEffects
 
     let mut eff ← oldEff.updateWithEq hStep
-    if let some hSp := c.effects.stackAlignmentProof? then
-      for subGoal in eff.sideConditions do
-        trace[Tactic.sym] "attempting to discharge side-condition:\n  {subGoal}"
-        let subGoal? ← do
-          let (ctx, simprocs) ←
-            LNSymSimpContext
-              (config := {failIfUnchanged := false, decide := true})
-              (exprs := #[hSp])
-          LNSymSimp subGoal ctx simprocs
-
-        if let some subGoal := subGoal? then
-          trace[Tactic.sym] "subgoal got simplified to:\n{subGoal}"
-          subGoal.setTag (.mkSimple s!"h_{← getNextStateName}_sp_aligned")
-          appendGoals [subGoal]
-        else
-          trace[Tactic.sym] "subgoal got closed by simplification"
-    else
-      appendGoals eff.sideConditions
+
+    withVerboseTraceNode m!"discharging side condiitions" <| do
+      if let some hSp := c.effects.stackAlignmentProof? then
+        for subGoal in eff.sideConditions do
+          trace[Tactic.sym] "attempting to discharge side-condition:\n  {subGoal}"
+          let subGoal? ← do
+            let (ctx, simprocs) ←
+              LNSymSimpContext
+                (config := {failIfUnchanged := false, decide := true})
+                (exprs := #[hSp])
+            LNSymSimp subGoal ctx simprocs
+
+          if let some subGoal := subGoal? then
+            trace[Tactic.sym] "subgoal got simplified to:\n{subGoal}"
+            subGoal.setTag (.mkSimple s!"h_{← getNextStateName}_sp_aligned")
+            appendGoals [subGoal]
+          else
+            trace[Tactic.sym] "subgoal got closed by simplification"
+      else
+        appendGoals eff.sideConditions
     eff := { eff with sideConditions := [] }
 
     if ←(getBoolOption `Tactic.sym.debug) then
@@ -214,9 +217,9 @@ Symbolically simulate a single step, according the the symbolic simulation
 context `c`, returning the context for the next step in simulation. -/
 def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
   let stateNumber ← getCurrentStateNumber
-  let msg := m!"(sym1): simulating step {stateNumber}"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext' do
-    withTraceNode `Tactic.sym (fun _ => pure "verbose context") <| do
+  withTraceNode m!"(sym1): simulating step {stateNumber}" <|
+  withMainContext' do
+    withTraceNode "verbose context" <| do
       traceSymContext
       trace[Tactic.sym] "Goal state:\n {← getMainGoal}"
 
@@ -383,17 +386,17 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
 
         let goal ← subst goal hEqId
         trace[Tactic.sym] "performed subsitutition in:\n{goal}"
-
         replaceMainGoal [goal]
 
-    let _eff ← withMainContext' do
-      let hypPrefix := s!"h_{←getCurrentStateName}_"
-      c.effects.addHypothesesToLContext (hypPrefix := hypPrefix)
+    let _eff ←
+      withTraceNode m!"adding (non-)effect hypotheses to local context" <|
+      withMainContext' do
+        let hypPrefix := s!"h_{←getCurrentStateName}_"
+        c.effects.addHypothesesToLContext (hypPrefix := hypPrefix)
 
     -- Rudimentary aggregation: we feed all the axiomatic effect hypotheses
     -- added while symbolically evaluating to `simp`
-    let msg := m!"aggregating (non-)effects"
-    withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext' do
+    withTraceNode m!"aggregating (non-)effects" <| withMainContext' do
       let goal? ← LNSymSimp (← getMainGoal) c.aggregateSimpCtx c.aggregateSimprocs
       replaceMainGoal goal?.toList
 

Tactics/Sym/AxEffects.lean

@@ -799,7 +799,9 @@ def toSimpTheoremArray (eff : AxEffects) : MetaM (Array SimpTheorem) := do
     pure thms
 
 def toSimpTheorems (eff : AxEffects) : MetaM SimpTheorems := do
-  let thms ← eff.toSimpTheoremArray
-  return thms.foldl addSimpTheoremEntry {}
+  let msg := m!"building DiscrTree for (non-)effect hypotheses"
+  withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+    let thms ← eff.toSimpTheoremArray
+    return thms.foldl addSimpTheoremEntry {}
 
 end Tactic

Tactics/Sym/Context.lean

@@ -210,17 +210,23 @@ def toMessageData (c : SymContext) : MetaM MessageData := do
   curr_state_number := {c.currentStateNumber},
   effects := {c.effects} }"
 
+section Tracing
 variable {α : Type} {m : Type → Type} [Monad m] [MonadTrace m] [MonadLiftT IO m]
   [MonadRef m] [AddMessageContext m] [MonadOptions m] {ε : Type}
-  [MonadAlwaysExcept ε m] [MonadLiftT BaseIO m] in
-def withSymTraceNode (msg : MessageData) (k : m α) : m α := do
-  withTraceNode `Tactic.sym (fun _ => pure msg) k
+  [MonadAlwaysExcept ε m] [MonadLiftT BaseIO m]
+
+def withTraceNode (msg : MessageData) (k : m α) : m α := do
+  Lean.withTraceNode `Tactic.sym (fun _ => pure msg) k
+def withVerboseTraceNode (msg : MessageData) (k : m α) : m α := do
+  Lean.withTraceNode `Tactic.sym.verbose (fun _ => pure msg) k
 
 def traceSymContext : SymM Unit :=
-  withTraceNode `Tactic.sym (fun _ => pure m!"SymContext: ") <| do
+  withTraceNode m!"SymContext: " <| do
     let m ← (← getThe SymContext).toMessageData
     trace[Tactic.sym] m
 
+end Tracing
+
 /-! ## Adding new simp theorems for aggregation -/
 
 /-- Add a set of new simp-theorems to the simp-theorems used
@@ -439,7 +445,7 @@ we create a new subgoal of this type.
 -/
 def fromMainContext (state? : Option Name) : TacticM SymContext := do
   let msg := m!"Building a `SymContext` from the local context"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext' do
+  withTraceNode msg <| withMainContext' do
   trace[Tactic.Sym] "state? := {state?}"
   let lctx ← getLCtx