refactor: polymorphic setTraceState, getTraceState, etc

src/Lean/CoreM.lean

@@ -65,12 +65,6 @@ def Context.incCurrRecDepth (ctx : Context) : Context :=
 @[inline] def withIncRecDepth {α} (x : CoreM α) : CoreM α := do
 checkRecDepth; adaptReader Context.incCurrRecDepth x
 
-def getTraceState : CoreM TraceState := do
-s ← get; pure s.traceState
-
-def setTraceState (traceState : TraceState) : CoreM Unit := do
-modify fun s => { s with traceState := traceState }
-
 def getNGen : CoreM NameGenerator := do
 s ← get; pure s.ngen
 
@@ -89,22 +83,12 @@ def Context.replaceRef (ref : Syntax) (ctx : Context) : Context :=
 @[inline] def withRef {α} (ref : Syntax) (x : CoreM α) : CoreM α := do
 adaptReader (Context.replaceRef ref) x
 
-@[inline] private def getTraceState : CoreM TraceState := do
-s ← get; pure s.traceState
-
 instance tracer : SimpleMonadTracerAdapter (CoreM) :=
 { getOptions       := getOptions,
-  getTraceState    := getTraceState,
-  addContext       := fun msg => Prod.snd <$> addContext Syntax.missing msg,
+  getTraceState    := do s ← get; pure s.traceState,
+  addTraceContext  := fun msg => Prod.snd <$> addContext Syntax.missing msg,
   modifyTraceState := fun f => modify $ fun s => { s with traceState := f s.traceState } }
 
-def printTraces : CoreM Unit := do
-traceState ← getTraceState;
-traceState.traces.forM $ fun m => liftIO $ IO.println $ format m
-
-def resetTraceState : CoreM Unit :=
-modify fun s => { s with traceState := {} }
-
 @[inline] def CoreM.run {α} (x : CoreM α) (ctx : Context) (s : State) : EIO Exception (α × State) :=
 (x.run ctx).run s
 

src/Lean/Elab/Term.lean

@@ -131,12 +131,6 @@ liftM $ x
 instance : MonadIO TermElabM :=
 { liftIO := fun α x => liftMetaMCore $ liftIO x }
 
-private def getTraceState : TermElabM TraceState :=
-liftMetaMCore Meta.getTraceState
-
-private def setTraceState (s : TraceState) : TermElabM Unit :=
-liftMetaMCore $ Meta.setTraceState s
-
 private def saveTraceAsMessages (traceState : TraceState) : TermElabM Unit :=
 unless traceState.traces.isEmpty do
   ref ← getRef;

src/Lean/Meta/Basic.lean

@@ -152,6 +152,12 @@ instance meta.monadError : MonadError MetaM :=
 { getRef     := liftM (getRef : CoreM Syntax),
   addContext := fun ref msg => do msg ← addContext msg; pure (ref, msg) }
 
+instance meta.simpleMonadTracerAdapter : SimpleMonadTracerAdapter MetaM :=
+{ getOptions       := liftM $ (getOptions : CoreM _),
+  getTraceState    := liftM $ (getTraceState : CoreM _),
+  modifyTraceState := fun f => liftM (modifyTraceState f : CoreM _),
+  addTraceContext  := addContext }
+
 def throwIsDefEqStuck {α} : MetaM α :=
 throw $ Exception.internal isDefEqStuckExceptionId
 
@@ -185,12 +191,6 @@ liftM Core.getNGen
 def setNGen (ngen : NameGenerator) : MetaM Unit :=
 liftM $ Core.setNGen ngen
 
-def getTraceState  : MetaM TraceState :=
-liftM $ Core.getTraceState
-
-def setTraceState (traceState : TraceState) : MetaM Unit :=
-liftM $ Core.setTraceState traceState
-
 def mkWHNFRef : IO (IO.Ref (Expr → MetaM Expr)) :=
 IO.mkRef $ fun _ => throwError "whnf implementation was not set"
 

src/Lean/Util/Trace.lean

@@ -26,7 +26,7 @@ inferInstanceAs (MonadTracer (ReaderT _ _))
 
 class MonadTracerAdapter (m : Type → Type) :=
 (isTracingEnabledFor : Name → m Bool)
-(addContext          : MessageData → m MessageData)
+(addTraceContext     : MessageData → m MessageData)
 (enableTracing       : Bool → m Bool)
 (getTraces           : m (PersistentArray MessageData))
 (modifyTraces        : (PersistentArray MessageData → PersistentArray MessageData) → m Unit)
@@ -57,7 +57,7 @@ modifyTraces $ fun _ => {};
 pure oldTraces
 
 def addTrace (cls : Name) (msg : MessageData) : m Unit := do
-msg ← addContext msg;
+msg ← addTraceContext msg;
 modifyTraces $ fun traces => traces.push (MessageData.tagged cls msg)
 
 @[inline] protected def trace (cls : Name) (msg : Unit → MessageData) : m Unit :=
@@ -111,10 +111,10 @@ instance monadTracerAdapterExcept {ε : Type} {m : Type → Type} [Monad m] [Mon
 
 instance liftMonadTracerAdapter {m n : Type → Type} [MonadTracerAdapter n] [HasMonadLift n m] : MonadTracerAdapter m :=
 { isTracingEnabledFor := fun cls => liftM (MonadTracerAdapter.isTracingEnabledFor cls : n _),
-  addContext := fun msg => liftM (MonadTracerAdapter.addContext msg : n _),
-  enableTracing := fun b => liftM (MonadTracerAdapter.enableTracing b : n _),
-  getTraces := liftM (MonadTracerAdapter.getTraces : n _),
-  modifyTraces := fun f => liftM (MonadTracerAdapter.modifyTraces f : n _) }
+  addTraceContext     := fun msg => liftM (MonadTracerAdapter.addTraceContext msg : n _),
+  enableTracing       := fun b => liftM (MonadTracerAdapter.enableTracing b : n _),
+  getTraces           := liftM (MonadTracerAdapter.getTraces : n _),
+  modifyTraces        := fun f => liftM (MonadTracerAdapter.modifyTraces f : n _) }
 
 structure TraceState :=
 (enabled : Bool := true)
@@ -141,7 +141,7 @@ class SimpleMonadTracerAdapter (m : Type → Type) :=
 (getOptions        : m Options)
 (modifyTraceState  : (TraceState → TraceState) → m Unit)
 (getTraceState     : m TraceState)
-(addContext        : MessageData → m MessageData)
+(addTraceContext   : MessageData → m MessageData)
 
 namespace SimpleMonadTracerAdapter
 variables {m : Type → Type} [Monad m] [SimpleMonadTracerAdapter m]
@@ -179,7 +179,7 @@ instance simpleMonadTracerAdapter {m : Type → Type} [SimpleMonadTracerAdapter
 { isTracingEnabledFor := @SimpleMonadTracerAdapter.isTracingEnabledFor _ _ _,
   enableTracing       := @SimpleMonadTracerAdapter.enableTracing _ _ _,
   getTraces           := @SimpleMonadTracerAdapter.getTraces _ _ _,
-  addContext          := @SimpleMonadTracerAdapter.addContext _ _,
+  addTraceContext     := @SimpleMonadTracerAdapter.addTraceContext _ _,
   modifyTraces        := @SimpleMonadTracerAdapter.modifyTraces _ _ _ }
 
 export MonadTracer (traceCtx trace traceM)
@@ -210,4 +210,27 @@ Recipe for adding tracing support for a monad `M`.
 def registerTraceClass (traceClassName : Name) : IO Unit :=
 registerOption (`trace ++ traceClassName) { group := "trace", defValue := false, descr := "enable/disable tracing for the given module and submodules" }
 
+export SimpleMonadTracerAdapter (getTraceState modifyTraceState)
+
+def setTraceState {m} [SimpleMonadTracerAdapter m] (s : TraceState) : m Unit :=
+modifyTraceState (fun _ => s)
+
+def printTraces {m} [Monad m] [SimpleMonadTracerAdapter m] [MonadIO m] : m Unit := do
+traceState ← getTraceState;
+traceState.traces.forM $ fun m => liftIO $ IO.println $ format m
+
+def resetTraceState {m} [SimpleMonadTracerAdapter m] : m Unit :=
+modifyTraceState (fun _ => {})
+
+/- We currently cannot mark the following definition as an instance since it increases the search space too much -/
+def simpleMonadTracerAdapterLift (m : Type → Type) {n : Type → Type} [SimpleMonadTracerAdapter m] [HasMonadLiftT m n] : SimpleMonadTracerAdapter n :=
+{ getOptions       := liftM (SimpleMonadTracerAdapter.getOptions : m _),
+  modifyTraceState := fun f => liftM (SimpleMonadTracerAdapter.modifyTraceState f : m _),
+  getTraceState    := liftM (SimpleMonadTracerAdapter.getTraceState : m _),
+  addTraceContext  := fun msg => liftM (SimpleMonadTracerAdapter.addTraceContext msg : m _) }
+
+instance ReaderT.tracer {m ρ} [Monad m] [SimpleMonadTracerAdapter m] : SimpleMonadTracerAdapter (ReaderT ρ m) := simpleMonadTracerAdapterLift m
+instance StateRefT.tracer {ω m σ} [SimpleMonadTracerAdapter m] : SimpleMonadTracerAdapter (StateRefT' ω σ m) := simpleMonadTracerAdapterLift m
+instance OptionT.tracer {m} [Monad m] [SimpleMonadTracerAdapter m] : SimpleMonadTracerAdapter (OptionT m) := simpleMonadTracerAdapterLift m
+
 end Lean