chore: add missing withTraceNode

Motivation: improve `trace.profiler`

src/Lean/Elab/PreDefinition/Basic.lean

@@ -53,20 +53,21 @@ private def getLevelParamsPreDecls (preDefs : Array PreDefinition) (scopeLevelNa
   | Except.error msg      => throwError msg
   | Except.ok levelParams => pure levelParams
 
-def fixLevelParams (preDefs : Array PreDefinition) (scopeLevelNames allUserLevelNames : List Name) : TermElabM (Array PreDefinition) :=
-do profileitM Exception s!"fix level params" (← getOptions) do
-  -- We used to use `shareCommon` here, but is was a bottleneck
-  let levelParams ← getLevelParamsPreDecls preDefs scopeLevelNames allUserLevelNames
-  let us := levelParams.map mkLevelParam
-  let fixExpr (e : Expr) : Expr :=
-    e.replace fun c => match c with
-      | Expr.const declName _ => if preDefs.any fun preDef => preDef.declName == declName then some $ Lean.mkConst declName us else none
-      | _ => none
-  return preDefs.map fun preDef =>
-    { preDef with
-      type        := fixExpr preDef.type,
-      value       := fixExpr preDef.value,
-      levelParams := levelParams }
+def fixLevelParams (preDefs : Array PreDefinition) (scopeLevelNames allUserLevelNames : List Name) : TermElabM (Array PreDefinition) := do
+  profileitM Exception s!"fix level params" (← getOptions) do
+    withTraceNode `Elab.def.fixLevelParams (fun _ => return m!"fix level params") do
+      -- We used to use `shareCommon` here, but is was a bottleneck
+      let levelParams ← getLevelParamsPreDecls preDefs scopeLevelNames allUserLevelNames
+      let us := levelParams.map mkLevelParam
+      let fixExpr (e : Expr) : Expr :=
+        e.replace fun c => match c with
+          | Expr.const declName _ => if preDefs.any fun preDef => preDef.declName == declName then some $ Lean.mkConst declName us else none
+          | _ => none
+      return preDefs.map fun preDef =>
+        { preDef with
+          type        := fixExpr preDef.type,
+          value       := fixExpr preDef.value,
+          levelParams := levelParams }
 
 def applyAttributesOf (preDefs : Array PreDefinition) (applicationTime : AttributeApplicationTime) : TermElabM Unit := do
   for preDef in preDefs do
@@ -212,15 +213,17 @@ def checkCodomainsLevel (preDefs : Array PreDefinition) : MetaM Unit := do
             m!"for `{preDefs[0]!.declName}` is{indentExpr type₀} : {← inferType type₀}\n" ++
             m!"and for `{preDefs[i]!.declName}` is{indentExpr typeᵢ} : {← inferType typeᵢ}"
 
-def shareCommonPreDefs (preDefs : Array PreDefinition) : CoreM (Array PreDefinition) := do profileitM Exception "share common exprs" (← getOptions) do
-  let mut es := #[]
-  for preDef in preDefs do
-    es := es.push preDef.type |>.push preDef.value
-  es := ShareCommon.shareCommon' es
-  let mut result := #[]
-  for h : i in [:preDefs.size] do
-    let preDef := preDefs[i]
-    result := result.push { preDef with type := es[2*i]!, value := es[2*i+1]! }
-  return result
+def shareCommonPreDefs (preDefs : Array PreDefinition) : CoreM (Array PreDefinition) := do
+  profileitM Exception "share common exprs" (← getOptions) do
+    withTraceNode `Elab.def.maxSharing (fun _ => return m!"share common exprs") do
+      let mut es := #[]
+      for preDef in preDefs do
+        es := es.push preDef.type |>.push preDef.value
+      es := ShareCommon.shareCommon' es
+      let mut result := #[]
+      for h : i in [:preDefs.size] do
+        let preDef := preDefs[i]
+        result := result.push { preDef with type := es[2*i]!, value := es[2*i+1]! }
+      return result
 
 end Lean.Elab

src/Lean/Elab/PreDefinition/Main.lean

@@ -152,69 +152,71 @@ def shouldUseWF (preDefs : Array PreDefinition) : Bool :=
     preDef.termination.decreasingBy?.isSome
 
 
-def addPreDefinitions (preDefs : Array PreDefinition) : TermElabM Unit := withLCtx {} {} do profileitM Exception "add pre-definitions" (← getOptions) do
-  for preDef in preDefs do
-    trace[Elab.definition.body] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
-  let preDefs ← preDefs.mapM ensureNoUnassignedMVarsAtPreDef
-  let preDefs ← betaReduceLetRecApps preDefs
-  let cliques := partitionPreDefs preDefs
-  for preDefs in cliques do
-    trace[Elab.definition.scc] "{preDefs.map (·.declName)}"
-    if preDefs.size == 1 && isNonRecursive preDefs[0]! then
-      /-
-      We must erase `recApp` annotations even when `preDef` is not recursive
-      because it may use another recursive declaration in the same mutual block.
-      See issue #2321
-      -/
-      let preDef ← eraseRecAppSyntax preDefs[0]!
-      ensureEqnReservedNamesAvailable preDef.declName
-      if preDef.modifiers.isNoncomputable then
-        addNonRec preDef
-      else
-        addAndCompileNonRec preDef
-      preDef.termination.ensureNone "not recursive"
-    else if preDefs.any (·.modifiers.isUnsafe) then
-      addAndCompileUnsafe preDefs
-      preDefs.forM (·.termination.ensureNone "unsafe")
-    else if preDefs.any (·.modifiers.isPartial) then
+def addPreDefinitions (preDefs : Array PreDefinition) : TermElabM Unit := withLCtx {} {} do
+  profileitM Exception "process pre-definitions" (← getOptions) do
+    withTraceNode `Elab.def.processPreDef (fun _ => return m!"process pre-definitions") do
       for preDef in preDefs do
-        if preDef.modifiers.isPartial && !(← whnfD preDef.type).isForall then
-          withRef preDef.ref <| throwError "invalid use of 'partial', '{preDef.declName}' is not a function{indentExpr preDef.type}"
-      addAndCompilePartial preDefs
-      preDefs.forM (·.termination.ensureNone "partial")
-    else
-      ensureFunIndReservedNamesAvailable preDefs
-      try
-        checkCodomainsLevel preDefs
-        checkTerminationByHints preDefs
-        let termArg?s ← elabTerminationByHints preDefs
-        if shouldUseStructural preDefs then
-          structuralRecursion preDefs termArg?s
-        else if shouldUseWF preDefs then
-          wfRecursion preDefs termArg?s
+        trace[Elab.definition.body] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
+      let preDefs ← preDefs.mapM ensureNoUnassignedMVarsAtPreDef
+      let preDefs ← betaReduceLetRecApps preDefs
+      let cliques := partitionPreDefs preDefs
+      for preDefs in cliques do
+        trace[Elab.definition.scc] "{preDefs.map (·.declName)}"
+        if preDefs.size == 1 && isNonRecursive preDefs[0]! then
+          /-
+          We must erase `recApp` annotations even when `preDef` is not recursive
+          because it may use another recursive declaration in the same mutual block.
+          See issue #2321
+          -/
+          let preDef ← eraseRecAppSyntax preDefs[0]!
+          ensureEqnReservedNamesAvailable preDef.declName
+          if preDef.modifiers.isNoncomputable then
+            addNonRec preDef
+          else
+            addAndCompileNonRec preDef
+          preDef.termination.ensureNone "not recursive"
+        else if preDefs.any (·.modifiers.isUnsafe) then
+          addAndCompileUnsafe preDefs
+          preDefs.forM (·.termination.ensureNone "unsafe")
+        else if preDefs.any (·.modifiers.isPartial) then
+          for preDef in preDefs do
+            if preDef.modifiers.isPartial && !(← whnfD preDef.type).isForall then
+              withRef preDef.ref <| throwError "invalid use of 'partial', '{preDef.declName}' is not a function{indentExpr preDef.type}"
+          addAndCompilePartial preDefs
+          preDefs.forM (·.termination.ensureNone "partial")
         else
-          withRef (preDefs[0]!.ref) <| mapError
-            (orelseMergeErrors
-              (structuralRecursion preDefs termArg?s)
-              (wfRecursion preDefs termArg?s))
-            (fun msg =>
-              let preDefMsgs := preDefs.toList.map (MessageData.ofExpr $ mkConst ·.declName)
-              m!"fail to show termination for{indentD (MessageData.joinSep preDefMsgs Format.line)}\nwith errors\n{msg}")
-      catch ex =>
-        logException ex
-        let s ← saveState
-        try
-          if preDefs.all fun preDef => preDef.kind == DefKind.def || preDefs.all fun preDef => preDef.kind == DefKind.abbrev then
-            -- try to add as partial definition
+          ensureFunIndReservedNamesAvailable preDefs
+          try
+            checkCodomainsLevel preDefs
+            checkTerminationByHints preDefs
+            let termArg?s ← elabTerminationByHints preDefs
+            if shouldUseStructural preDefs then
+              structuralRecursion preDefs termArg?s
+            else if shouldUseWF preDefs then
+              wfRecursion preDefs termArg?s
+            else
+              withRef (preDefs[0]!.ref) <| mapError
+                (orelseMergeErrors
+                  (structuralRecursion preDefs termArg?s)
+                  (wfRecursion preDefs termArg?s))
+                (fun msg =>
+                  let preDefMsgs := preDefs.toList.map (MessageData.ofExpr $ mkConst ·.declName)
+                  m!"fail to show termination for{indentD (MessageData.joinSep preDefMsgs Format.line)}\nwith errors\n{msg}")
+          catch ex =>
+            logException ex
+            let s ← saveState
             try
-              addAndCompilePartial preDefs (useSorry := true)
-            catch _ =>
-              -- Compilation failed try again just as axiom
-              s.restore
-              addAsAxioms preDefs
-          else if preDefs.all fun preDef => preDef.kind == DefKind.theorem then
-            addAsAxioms preDefs
-        catch _ => s.restore
+              if preDefs.all fun preDef => preDef.kind == DefKind.def || preDefs.all fun preDef => preDef.kind == DefKind.abbrev then
+                -- try to add as partial definition
+                try
+                  addAndCompilePartial preDefs (useSorry := true)
+                catch _ =>
+                  -- Compilation failed try again just as axiom
+                  s.restore
+                  addAsAxioms preDefs
+              else if preDefs.all fun preDef => preDef.kind == DefKind.theorem then
+                addAsAxioms preDefs
+            catch _ => s.restore
 
 builtin_initialize
   registerTraceClass `Elab.definition.body