feat: ensure trace messages at MetaM save Environment, MetavarContext, and LocalContext

src/Init/Lean/Meta/Basic.lean

@@ -303,6 +303,12 @@ instance tracer : SimpleMonadTracerAdapter MetaM :=
   getTraceState    := getTraceState,
   modifyTraceState := fun f => modify $ fun s => { traceState := f s.traceState, .. s } }
 
+@[inline] def trace (cls : Name) (msg : Unit → MessageData) : MetaM Unit :=
+whenM (MonadTracerAdapter.isTracingEnabledFor cls) $ do
+  ctx ← read;
+  s   ← get;
+  MonadTracerAdapter.addTrace cls (MessageData.context s.env s.mctx ctx.lctx (msg ()))
+
 def getConstAux (constName : Name) (exception? : Bool) : MetaM (Option ConstantInfo) :=
 do env ← getEnv;
    match env.find constName with

src/Init/Lean/Meta/Check.lean

@@ -85,7 +85,7 @@ def isTypeCorrect (e : Expr) : MetaM Bool :=
 catch
   (traceCtx `Meta.check $ do checkAux e; pure true)
   (fun ex => do
-    trace! `Meta.typeError ex.toMessageData;
+    trace `Meta.typeError $ fun _ => ex.toMessageData;
     pure false)
 
 end Meta

src/Init/Lean/Meta/ExprDefEq.lean

@@ -438,22 +438,18 @@ end CheckAssignment
 private def checkAssignmentFailure (mvarId : Name) (fvars : Array Expr) (v : Expr) (ex : CheckAssignment.Exception) : MetaM (Option Expr) :=
 match ex with
 | CheckAssignment.Exception.occursCheck => do
-  trace! `Meta.isDefEq.assign.occursCheck
-    (mkMVar mvarId ++ fvars ++ " := " ++ v);
+  trace `Meta.isDefEq.assign.occursCheck $ fun _ => mkMVar mvarId ++ fvars ++ " := " ++ v;
   pure none
 | CheckAssignment.Exception.useFOApprox =>
   pure none
 | CheckAssignment.Exception.outOfScopeFVar fvarId => do
-  trace! `Meta.isDefEq.assign.outOfScopeFVar
-    (mkFVar fvarId ++ " @ " ++ mkMVar mvarId ++ fvars ++ " := " ++ v);
+  trace `Meta.isDefEq.assign.outOfScopeFVar $ fun _ => mkFVar fvarId ++ " @ " ++ mkMVar mvarId ++ fvars ++ " := " ++ v;
   pure none
 | CheckAssignment.Exception.readOnlyMVarWithBiggerLCtx nestedMVarId => do
-  trace! `Meta.isDefEq.assign.readOnlyMVarWithBiggerLCtx
-    (mkMVar nestedMVarId ++ " @ " ++ mkMVar mvarId ++ fvars ++ " := " ++ v);
+  trace `Meta.isDefEq.assign.readOnlyMVarWithBiggerLCtx $ fun _ => mkMVar nestedMVarId ++ " @ " ++ mkMVar mvarId ++ fvars ++ " := " ++ v;
   pure none
 | CheckAssignment.Exception.mvarTypeNotWellFormedInSmallerLCtx nestedMVarId => do
-  trace! `Meta.isDefEq.assign.mvarTypeNotWellFormedInSmallerLCtx
-    (mkMVar nestedMVarId ++ " @ " ++ mkMVar mvarId ++ fvars ++ " := " ++ v);
+  trace `Meta.isDefEq.assign.mvarTypeNotWellFormedInSmallerLCtx $ fun _ => mkMVar nestedMVarId ++ " @ " ++ mkMVar mvarId ++ fvars ++ " := " ++ v;
   pure none
 | CheckAssignment.Exception.unknownExprMVar mvarId =>
   -- This case can only happen if the MetaM API is being misused
@@ -608,7 +604,7 @@ else
 -/
 private partial def processAssignmentFOApprox (mvar : Expr) (args : Array Expr) : Expr → MetaM Bool
 | v => do
-  trace! `Meta.isDefEq.foApprox (mvar ++ " " ++ args ++ " := " ++ v);
+  trace `Meta.isDefEq.foApprox $ fun _ => mvar ++ " " ++ args ++ " := " ++ v;
   condM (try $ processAssignmentFOApproxAux mvar args v)
     (pure true)
     (do v? ← unfoldDefinition v;
@@ -675,7 +671,7 @@ private partial def processAssignmentAux (mvar : Expr) (mvarDecl : MetavarDecl)
            vType    ← inferType v;
            condM (usingTransparency TransparencyMode.default $ isExprDefEqAux mvarType vType)
              (do assignExprMVar mvarId v; pure true)
-             (do trace! `Meta.isDefEq.assign.typeMismatch (mvar ++ " : " ++ mvarType ++ " := " ++ v ++ " : " ++ vType);
+             (do trace `Meta.isDefEq.assign.typeMismatch $ fun _ => mvar ++ " : " ++ mvarType ++ " := " ++ v ++ " : " ++ vType;
                  pure false)
         };
         if args.any (fun arg => mvarDecl.lctx.containsFVar arg) then
@@ -683,7 +679,7 @@ private partial def processAssignmentAux (mvar : Expr) (mvarDecl : MetavarDecl)
              a type incorrect term. See discussion at A2 -/
           condM (isTypeCorrect v)
             (finalize ())
-            (do trace! `Meta.isDefEq.assign.typeError (mvar ++ " := " ++ v);
+            (do trace `Meta.isDefEq.assign.typeError $ fun _ => mvar ++ " := " ++ v;
                 useFOApprox ())
         else
           finalize ()
@@ -692,7 +688,7 @@ private partial def processAssignmentAux (mvar : Expr) (mvarDecl : MetavarDecl)
     It assumes `?m` is unassigned. -/
 private def processAssignment (mvarApp : Expr) (v : Expr) : MetaM Bool :=
 traceCtx `Meta.isDefEq.assign $ do
-  trace! `Meta.isDefEq.assign (mvarApp ++ " := " ++ v);
+  trace `Meta.isDefEq.assign $ fun _ => mvarApp ++ " := " ++ v;
   let mvar := mvarApp.getAppFn;
   mvarDecl ← getMVarDecl mvar.mvarId!;
   processAssignmentAux mvar mvarDecl v 0 mvarApp.getAppArgs
@@ -708,17 +704,17 @@ toLBoolM $ isListLevelDefEqAux us vs
 
 /-- Auxiliary method for isDefEqDelta -/
 private def isDefEqLeft (fn : Name) (t s : Expr) : MetaM LBool :=
-do trace! `Meta.isDefEq.delta.unfoldLeft fn;
+do trace `Meta.isDefEq.delta.unfoldLeft $ fun _ => fn;
    toLBoolM $ isExprDefEqAux t s
 
 /-- Auxiliary method for isDefEqDelta -/
 private def isDefEqRight (fn : Name) (t s : Expr) : MetaM LBool :=
-do trace! `Meta.isDefEq.delta.unfoldRight fn;
+do trace `Meta.isDefEq.delta.unfoldRight $ fun _ => fn;
    toLBoolM $ isExprDefEqAux t s
 
 /-- Auxiliary method for isDefEqDelta -/
 private def isDefEqLeftRight (fn : Name) (t s : Expr) : MetaM LBool :=
-do trace! `Meta.isDefEq.delta.unfoldLeftRight fn;
+do trace `Meta.isDefEq.delta.unfoldLeftRight $ fun _ => fn;
    toLBoolM $ isExprDefEqAux t s
 
 /-- Try to solve `f a₁ ... aₙ =?= f b₁ ... bₙ` by solving `a₁ =?= b₁, ..., aₙ =?= bₙ`.
@@ -732,7 +728,7 @@ traceCtx `Meta.isDefEq.delta $
     b ← isDefEqArgs tFn t.getAppArgs s.getAppArgs
         <&&>
         isListLevelDefEqAux tFn.constLevels! sFn.constLevels!;
-    unless b $ trace! `Meta.isDefEq.delta ("heuristic failed " ++ t ++ " =?= " ++ s);
+    unless b $ trace `Meta.isDefEq.delta $ fun _ => "heuristic failed " ++ t ++ " =?= " ++ s;
     pure b
 
 /-- Auxiliary method for isDefEqDelta -/
@@ -984,7 +980,7 @@ pure false
 
 partial def isExprDefEqAuxImpl : Expr → Expr → MetaM Bool
 | t, s => do
-  trace! `Meta.isDefEq.step (t ++ " =?= " ++ s);
+  trace `Meta.isDefEq.step $ fun _ => t ++ " =?= " ++ s;
   tryL (isDefEqQuick t s) $
   tryL (isDefEqProofIrrel t s) $
   isDefEqWHNF t s $ fun t s => do

src/Init/Lean/Meta/LevelDefEq.lean

@@ -58,7 +58,7 @@ partial def isLevelDefEqAux : Level → Level → MetaM Bool
   if lhs == rhs then
     pure true
   else do
-    trace! `Meta.isLevelDefEq.step (lhs ++ " =?= " ++ rhs);
+    trace `Meta.isLevelDefEq.step $ fun _ => lhs ++ " =?= " ++ rhs;
     lhs' ← instantiateLevelMVars lhs;
     let lhs' := lhs'.normalize;
     rhs' ← instantiateLevelMVars rhs;
@@ -104,7 +104,7 @@ do s ← get;
    pure ps
 
 private def processPostponedStep : MetaM Bool :=
-traceCtx `type_context.level_is_def_eq.postponed_step $ do
+traceCtx `Meta.isLevelDefEq.postponed.step $ do
   ps ← getResetPostponed;
   ps.foldlM
     (fun (r : Bool) (p : PostponedEntry) =>
@@ -120,7 +120,7 @@ private partial def processPostponedAux : Unit → MetaM Bool
   if numPostponed == 0 then
     pure true
   else do
-    trace! `type_context.level_is_def_eq ("processing #" ++ toString numPostponed ++ " postponed is-def-eq level constraints");
+    trace `Meta.isLevelDefEq.postponed $ fun _ => "processing #" ++ toString numPostponed ++ " postponed is-def-eq level constraints";
     r ← processPostponedStep;
     if !r then
       pure r
@@ -131,13 +131,13 @@ private partial def processPostponedAux : Unit → MetaM Bool
       else if numPostponed' < numPostponed then
         processPostponedAux ()
       else do
-        trace! `type_context.level_is_def_eq ("no progress solving pending is-def-eq level constraints");
+        trace `Meta.isLevelDefEq.postponed $ fun _ => format "no progress solving pending is-def-eq level constraints";
         pure false
 
 private def processPostponed : MetaM Bool :=
 do numPostponed ← getNumPostponed;
    if numPostponed == 0 then pure true
-   else traceCtx `type_context.level_is_def_eq.postponed $ processPostponedAux ()
+   else traceCtx `Meta.isLevelDefEq.postponed $ processPostponedAux ()
 
 
 private def restore (env : Environment) (mctx : MetavarContext) (postponed : PersistentArray PostponedEntry) : MetaM Unit :=
@@ -168,13 +168,13 @@ do s ← get;
 def isLevelDefEq (u v : Level) : MetaM Bool :=
 traceCtx `Meta.isLevelDefEq $ do
   b ← try $ isLevelDefEqAux u v;
-  trace! `Meta.isLevelDefEq (u ++ " =?= " ++ v ++ " ... " ++ if b then "success" else "failure");
+  trace `Meta.isLevelDefEq $ fun _ => u ++ " =?= " ++ v ++ " ... " ++ if b then "success" else "failure";
   pure b
 
 def isExprDefEq (t s : Expr) : MetaM Bool :=
 traceCtx `Meta.isDefEq $ do
   b ← try $ isExprDefEqAux t s;
-  trace! `Meta.isDefEq (t ++ " =?= " ++ s ++ " ... " ++ if b then "success" else "failure");
+  trace `Meta.isDefEq $ fun _ => t ++ " =?= " ++ s ++ " ... " ++ if b then "success" else "failure";
   pure b
 
 end Meta

src/Init/Lean/Trace.lean

@@ -36,7 +36,7 @@ do oldTraces ← getTraces;
    modifyTraces $ fun _ => #[];
    pure oldTraces
 
-private def addTrace (cls : Name) (msg : MessageData) : m Unit :=
+def addTrace (cls : Name) (msg : MessageData) : m Unit :=
 modifyTraces $ fun traces => traces.push (MessageData.tagged cls msg)
 
 @[inline] protected def trace (cls : Name) (msg : Unit → MessageData) : m Unit :=