chore: add tags to trace nodes

These tags show up in the profile (unlike the message itself)

Tactics/Sym.lean

@@ -90,7 +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
-  withTraceNode m!"unfoldRun (runSteps? := {c.runSteps?})" <|
+  withTraceNode m!"unfoldRun (runSteps? := {c.runSteps?})" (tag := "unfoldRun") <|
   match c.runSteps? with
     | some (_ + 1) => do
         trace[Tactic.sym] "runSteps is statically known to be non-zero, \
@@ -167,7 +167,7 @@ store an `AxEffects` object with the newly added variables in the monad state
 -/
 def explodeStep (hStep : Expr) : SymM Unit :=
   withMainContext' <|
-  withTraceNode m!"explodeStep {hStep}" <| do
+  withTraceNode m!"explodeStep {hStep}" (tag := "explodeStep") <| do
     let c ← getThe SymContext
     let mut eff ← AxEffects.fromEq hStep c.effects.stackAlignmentProof?
 
@@ -232,22 +232,24 @@ 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
-  withTraceNode m!"(sym1): simulating step {stateNumber}" <|
+  withTraceNode m!"(sym1): simulating step {stateNumber}" (tag:="sym1") <|
   withMainContext' do
-    withTraceNode "verbose context" <| do
+    withVerboseTraceNode "verbose context" (tag := "infoDump") <| do
       traceSymContext
       trace[Tactic.sym] "Goal state:\n {← getMainGoal}"
 
     let stepi_eq := Lean.mkIdent (.mkSimple s!"stepi_{← getCurrentStateName}")
     let h_step   := Lean.mkIdent (.mkSimple s!"h_step_{stateNumber + 1}")
 
     unfoldRun (fun _ => evalTacticAndTrace whileTac)
+
     -- Add new state to local context
-    let hRunId      := mkIdent <|← getHRunName
-    let nextStateId := mkIdent <|← getNextStateName
-    evalTacticAndTrace <|← `(tactic|
-      init_next_step $hRunId:ident $stepi_eq:ident $nextStateId:ident
-    )
+    withTraceNode "initNextStep" (tag := "initNextStep") <| do
+      let hRunId      := mkIdent <|← getHRunName
+      let nextStateId := mkIdent <|← getNextStateName
+      evalTacticAndTrace <|← `(tactic|
+        init_next_step $hRunId:ident $stepi_eq:ident $nextStateId:ident
+      )
 
     -- Apply relevant pre-generated `stepi` lemma
     withMainContext' <| do
@@ -262,15 +264,15 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
 
       -- If we know SP is aligned, `simp` with that fact
       if let some hSp := (← getThe AxEffects).stackAlignmentProof? then
-        trace[Tactic.sym] "simplifying {hStep.toExpr} \
-          with {hSp}"
-        let some goal ← do
-            let (ctx, simprocs) ← LNSymSimpContext
-              (config := {decide := false}) (exprs := #[hSp])
-            let goal ← getMainGoal
-            LNSymSimp goal ctx simprocs hStep.fvarId
-          | throwError "internal error: simp closed goal unexpectedly"
-        replaceMainGoal [goal]
+        let msg := m!"simplifying {hStep.toExpr} with {hSp}"
+        withTraceNode msg (tag := "simplifyHStep") <| do
+          let some goal ← do
+              let (ctx, simprocs) ← LNSymSimpContext
+                (config := {decide := false}) (exprs := #[hSp])
+              let goal ← getMainGoal
+              LNSymSimp goal ctx simprocs hStep.fvarId
+            | throwError "internal error: simp closed goal unexpectedly"
+          replaceMainGoal [goal]
       else
         trace[Tactic.sym] "we have no relevant local hypotheses, \
           skipping simplification step"
@@ -293,44 +295,46 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
 - log a warning and return `m`, if `runSteps? = some m` and `m < n`, or
 - return `n` unchanged, otherwise  -/
 def ensureAtMostRunSteps (n : Nat) : SymM Nat := do
-  let ctx ← getThe SymContext
-  match ctx.runSteps? with
-  | none => pure n
-  | some runSteps =>
-      if n ≤ runSteps then
-        pure n
-      else
-        withMainContext <| do
-          let hRun ← ctx.hRunDecl
-          logWarning m!"Symbolic simulation is limited to at most {runSteps} \
-            steps, because {hRun.toExpr} is of type:\n  {hRun.type}"
-          pure runSteps
-  return n
+  withVerboseTraceNode "" (tag := "ensureAtMostRunSteps") <| do
+    let ctx ← getThe SymContext
+    match ctx.runSteps? with
+    | none => pure n
+    | some runSteps =>
+        if n ≤ runSteps then
+          pure n
+        else
+          withMainContext <| do
+            let hRun ← ctx.hRunDecl
+            logWarning m!"Symbolic simulation is limited to at most {runSteps} \
+              steps, because {hRun.toExpr} is of type:\n  {hRun.type}"
+            pure runSteps
+    return n
 
 /-- Check that the step-thoerem corresponding to the current PC value exists,
 and throw a user-friendly error, pointing to `#genStepEqTheorems`,
 if it does not. -/
-def assertStepTheoremsGenerated : SymM Unit := do
-  let c ← getThe SymContext
-  let pc := c.pc.toHexWithoutLeadingZeroes
-  if !c.programInfo.instructions.contains c.pc then
-    let pcEff ← AxEffects.getFieldM .PC
-    throwError "\
-      Program {c.program} has no instruction at address {c.pc}.
-
-      We inferred this address as the program-counter from {pcEff.proof}, \
-      which has type:
-        {← inferType pcEff.proof}"
-
-  let step_thm := Name.str c.program ("stepi_eq_0x" ++ pc)
-  try
-    let _ ← getConstInfo step_thm
-  catch err =>
-    throwErrorAt err.getRef "{err.toMessageData}\n
-Did you remember to generate step theorems with:
-  #genStepEqTheorems {c.program}"
--- TODO: can we make this error ^^ into a `Try this:` suggestion that
---       automatically adds the right command just before the theorem?
+def assertStepTheoremsGenerated : SymM Unit :=
+  withVerboseTraceNode "" (tag := "assertStepTheoremsGenerated") <| do
+    let c ← getThe SymContext
+    let pc := c.pc.toHexWithoutLeadingZeroes
+    if !c.programInfo.instructions.contains c.pc then
+      let pcEff ← AxEffects.getFieldM .PC
+      throwError "\
+        Program {c.program} has no instruction at address {c.pc}.
+
+        We inferred this address as the program-counter from {pcEff.proof}, \
+        which has type:
+          {← inferType pcEff.proof}"
+
+    let step_thm := Name.str c.program ("stepi_eq_0x" ++ pc)
+    try
+      let _ ← getConstInfo step_thm
+    catch err =>
+      throwErrorAt err.getRef "{err.toMessageData}\n
+  Did you remember to generate step theorems with:
+    #genStepEqTheorems {c.program}"
+  -- TODO: can we make this error ^^ into a `Try this:` suggestion that
+  --       automatically adds the right command just before the theorem?
 
 /- used in `sym_n` tactic to specify an initial state -/
 syntax sym_at := "at" ident
@@ -390,32 +394,35 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
         sym1 whileTac
 
     traceHeartbeats "symbolic simulation total"
-    let c ← getThe SymContext
-    -- Check if we can substitute the final state
-    if c.runSteps? = some 0 then
-      let msg := do
-        let hRun ← userNameToMessageData c.h_run
-        pure m!"runSteps := 0, substituting along {hRun}"
-      withTraceNode `Tactic.sym (fun _ => msg) <| withMainContext' do
-        let sfEq ← mkEq (← getCurrentState) c.finalState
-
-        let goal ← getMainGoal
-        trace[Tactic.sym] "original goal:\n{goal}"
-        let ⟨hEqId, goal⟩ ← do
-          let hRun ← SymContext.findFromUserName c.h_run
-          goal.note `this (← mkEqSymm hRun.toExpr) sfEq
-        goal.withContext <| do
-          trace[Tactic.sym] "added {← userNameToMessageData `this} of type \
-            {sfEq} in:\n{goal}"
-
-        let goal ← subst goal hEqId
-        trace[Tactic.sym] "performed subsitutition in:\n{goal}"
-        replaceMainGoal [goal]
-
-    -- Rudimentary aggregation: we feed all the axiomatic effect hypotheses
-    -- added while symbolically evaluating to `simp`
-    withTraceNode m!"aggregating (non-)effects" <| withMainContext' do
-      let goal? ← LNSymSimp (← getMainGoal) c.aggregateSimpCtx c.aggregateSimprocs
-      replaceMainGoal goal?.toList
+    withTraceNode "Post processing" (tag := "postProccessing") <| do
+      let c ← getThe SymContext
+      -- Check if we can substitute the final state
+      if c.runSteps? = some 0 then
+        let msg := do
+          let hRun ← userNameToMessageData c.h_run
+          pure m!"runSteps := 0, substituting along {hRun}"
+        withMainContext' <|
+        withTraceNode `Tactic.sym (fun _ => msg) <| do
+          let sfEq ← mkEq (← getCurrentState) c.finalState
+
+          let goal ← getMainGoal
+          trace[Tactic.sym] "original goal:\n{goal}"
+          let ⟨hEqId, goal⟩ ← do
+            let hRun ← SymContext.findFromUserName c.h_run
+            goal.note `this (← mkEqSymm hRun.toExpr) sfEq
+          goal.withContext <| do
+            trace[Tactic.sym] "added {← userNameToMessageData `this} of type \
+              {sfEq} in:\n{goal}"
+
+          let goal ← subst goal hEqId
+          trace[Tactic.sym] "performed subsitutition in:\n{goal}"
+          replaceMainGoal [goal]
+
+      -- Rudimentary aggregation: we feed all the axiomatic effect hypotheses
+      -- added while symbolically evaluating to `simp`
+      withMainContext' <|
+      withTraceNode m!"aggregating (non-)effects" (tag := "aggregateEffects") <| do
+        let goal? ← LNSymSimp (← getMainGoal) c.aggregateSimpCtx c.aggregateSimprocs
+        replaceMainGoal goal?.toList
 
     traceHeartbeats "final usage"

Tactics/Sym/Context.lean

@@ -146,6 +146,26 @@ end SymM
 
 namespace SymContext
 
+/-! ## Trace Nodes -/
+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]
+
+def withTraceNode (msg : MessageData) (k : m α)
+    (collapsed : Bool := true)
+    (tag : String := "")
+    : m α := do
+  Lean.withTraceNode `Tactic.sym (fun _ => pure msg) k collapsed tag
+
+def withVerboseTraceNode (msg : MessageData) (k : m α)
+    (collapsed : Bool := true)
+    (tag : String := "")
+    : m α := do
+  Lean.withTraceNode `Tactic.sym.verbose (fun _ => pure msg) k collapsed tag
+
+end Tracing
+
 /-! ## Simple projections -/
 section
 open Lean (Ident mkIdent)
@@ -156,10 +176,11 @@ def program : Name := c.programInfo.name
 
 /-- Find the local declaration that corresponds to a given name,
 or throw an error if no local variable of that name exists -/
-def findFromUserName (name : Name) : MetaM LocalDecl := do
-  let some decl := (← getLCtx).findFromUserName? name
-    | throwError "Unknown local variable `{name}`"
-  return decl
+def findFromUserName (name : Name) : MetaM LocalDecl :=
+  withVerboseTraceNode m!"[findFromUserName] {name}" <| do
+    let some decl := (← getLCtx).findFromUserName? name
+      | throwError "Unknown local variable `{name}`"
+    return decl
 
 /-- Find the local declaration that corresponds to `c.h_run`,
 or throw an error if no local variable of that name exists -/
@@ -204,22 +225,12 @@ 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]
-
-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 m!"SymContext: " <| do
     let m ← (← getThe SymContext).toMessageData
     trace[Tactic.sym] m
 
-end Tracing
+
 
 /-! ## Adding new simp theorems for aggregation -/
 
@@ -437,7 +448,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 msg <| withMainContext' do
+  withTraceNode msg (tag := "fromMainContext") <| withMainContext' do
   trace[Tactic.Sym] "state? := {state?}"
   let lctx ← getLCtx
 
@@ -470,16 +481,17 @@ evaluation:
   * the `currentStateNumber` is incremented
 -/
 def prepareForNextStep : SymM Unit := do
-  let pc ← do
-    let { value, ..} ← AxEffects.getFieldM .PC
-    try
-      reflectBitVecLiteral 64 value
-    catch err =>
-      trace[Tactic.sym] "failed to reflect PC: {err.toMessageData}"
-      pure <| (← getThe SymContext).pc + 4
-
-  modifyThe SymContext (fun c => { c with
-    pc
-    runSteps?   := (· - 1) <$> c.runSteps?
-    currentStateNumber := c.currentStateNumber + 1
-  })
+  withVerboseTraceNode "prepareForNextStep" (tag := "prepareForNextStep") <| do
+    let pc ← do
+      let { value, ..} ← AxEffects.getFieldM .PC
+      try
+        reflectBitVecLiteral 64 value
+      catch err =>
+        trace[Tactic.sym] "failed to reflect PC: {err.toMessageData}"
+        pure <| (← getThe SymContext).pc + 4
+
+    modifyThe SymContext (fun c => { c with
+      pc
+      runSteps?   := (· - 1) <$> c.runSteps?
+      currentStateNumber := c.currentStateNumber + 1
+    })