Improve tracing for rewrites

Lean/Egg/Tactic/Basic.lean

@@ -45,10 +45,23 @@ where
     else
       throwError "expected goal to be of type '=' or '↔', but found:\n{← ppExpr goalType}"
 
+private def traceRewrites
+    (basic : Rewrites) (stx : Array Syntax) (tc : Rewrites) (cfg : Config.Gen) : TacticM Unit := do
+  let cls := `egg.rewrites
+  withTraceNode cls (fun _ => return m!"Rewrites") do
+    withTraceNode cls (fun _ => return m!"Basic ({basic.size})") do basic.trace stx cls
+    withTraceNode cls (fun _ => return m!"Generated ({tc.size})") do tc.trace #[] cls
+    withTraceNode cls (fun _ => return "Definitional") do
+      if cfg.genBetaRw    then Lean.trace cls fun _ => "β-Reduction"
+      if cfg.genEtaRw     then Lean.trace cls fun _ => "η-Reduction"
+      if cfg.genNatLitRws then Lean.trace cls fun _ => "Natural Number Literals"
+
 private partial def genRewrites
     (goal : Goal) (rws : TSyntax `egg_rws) (guides : Guides) (cfg : Config) : TacticM Rewrites := do
-  let rws ← Rewrites.parse cfg.betaReduceRws cfg.etaReduceRws rws
-  return rws ++ (← genTcRws rws)
+  let (rws, stx) ← Rewrites.parse cfg.betaReduceRws cfg.etaReduceRws rws
+  let tcRws ← genTcRws rws
+  traceRewrites rws stx tcRws cfg.toGen
+  return rws ++ tcRws
 where
   genTcRws (rws : Rewrites) : TacticM Rewrites := do
     let mut projTodo := #[]

Lean/Egg/Tactic/Rewrites.lean

@@ -78,8 +78,8 @@ def star (beta eta : Bool) : MetaM Rewrites := do
     then result := result.push rw
   return result
 
-def parse (beta eta : Bool) : (TSyntax `egg_rws) → TacticM Rewrites
-  | `(egg_rws|)          => return {}
+def parse (beta eta : Bool) : (TSyntax `egg_rws) → TacticM (Rewrites × Array Syntax)
+  | `(egg_rws|)          => return (#[], #[])
   | `(egg_rws|[$args,*]) => do
     let mut result : Rewrites := #[]
     let mut noStar := true
@@ -93,5 +93,5 @@ def parse (beta eta : Bool) : (TSyntax `egg_rws) → TacticM Rewrites
         result := result ++ (← star beta eta)
       | _ =>
         throwUnsupportedSyntax
-    return result
+    return (result, args)
   | _ => throwUnsupportedSyntax

Lean/Egg/Tactic/Trace.lean

@@ -1,24 +1,58 @@
 import Egg.Core.Request
 import Lean
-open Lean Elab Tactic
+open Lean Meta Elab Tactic Std Format
 
 initialize registerTraceClass `egg
+initialize registerTraceClass `egg.rewrites (inherited := true)
 initialize registerTraceClass `egg.frontend (inherited := true)
 initialize registerTraceClass `egg.reconstruction (inherited := true)
 
 namespace Egg
 
--- TODO: Replace this with a `MessageData` instance.
+#check MessageData
+
+nonrec def MVars.format (mvars : MVars) : MetaM Format := do
+  let expr := format <| ← mvars.expr.toList.mapM (ppExpr <| Expr.mvar ·)
+  let tc   := format <| ← mvars.tc.toList.mapM   (ppExpr <| Expr.mvar ·)
+  let lvl  := format <|   mvars.lvl.toList.map   (Level.mvar ·)
+  return "expr: " ++ expr ++ "\n" ++ "class: " ++ tc ++ "\n" ++ "level: " ++ lvl
+
+def Directions.format : Directions → Format
+  | .none     => "∅"
+  | .forward  => "⇒"
+  | .backward => "⇐"
+  | .both     => "⇔"
+
+def Congr.Rel.format : Congr.Rel → Format
+  | .eq  => "="
+  | .iff => "↔"
+
+def Congr.format (cgr : Congr) : MetaM Format :=
+  return (← ppExpr cgr.lhs) ++ " " ++ cgr.rel.format ++ " " ++ (← ppExpr cgr.rhs)
+
+def Rewrite.trace (rw : Rewrite) (stx? : Option Syntax) (cls : Name) : TacticM Unit := do
+  let mut header := m!"{rw.src.description}({rw.validDirs.format})"
+  if let some stx := stx? then header := m!"{header}: {stx}"
+  withTraceNode cls (fun _ => return header) do
+    Lean.traceM cls fun _ => return m!"{← rw.toCongr.format}"
+    Lean.traceM cls fun _ => return m!"LHS MVars\n{← rw.lhsMVars.format}"
+    Lean.traceM cls fun _ => return m!"RHS MVars\n{← rw.rhsMVars.format}"
+
+def Rewrites.trace (rws : Rewrites) (stx : Array Syntax) (cls : Name) : TacticM Unit := do
+  for rw in rws, idx in [:rws.size] do rw.trace stx[idx]? cls
+
+def Rewrite.Encoded.trace (rw : Rewrite.Encoded) (cls : Name) : TacticM Unit := do
+  let header := m!"{rw.name}({rw.dirs.format})"
+  withTraceNode cls (fun _ => return header) do
+    Lean.trace cls fun _ => m!"LHS: {rw.lhs}"
+    Lean.trace cls fun _ => m!"RHS: {rw.rhs}"
+
 def Request.trace (req : Request) : TacticM Unit := do
   withTraceNode `egg.frontend (fun _ => return "Request") do
     withTraceNode `egg.frontend (fun _ => return "Goal") (collapsed := false) do
       withTraceNode `egg.frontend (fun _ => return "LHS") do trace[egg.frontend] req.lhs
       withTraceNode `egg.frontend (fun _ => return "RHS") do trace[egg.frontend] req.rhs
     let rwsTitle := (if req.rws.isEmpty && !req.cfg.genNatLitRws then "No " else "") ++ "Rewrites"
     withTraceNode `egg.frontend (fun _ => return rwsTitle) (collapsed := false) do
-      for rw in req.rws do
-        withTraceNode `egg.frontend (fun _ => return m!"{rw.name}") do
-          trace[egg.frontend] "Directions: {rw.dirs}"
-          withTraceNode `egg.frontend (fun _ => return "LHS") (collapsed := false) do trace[egg.frontend] rw.lhs
-          withTraceNode `egg.frontend (fun _ => return "RHS") (collapsed := false) do trace[egg.frontend] rw.rhs
+      for rw in req.rws do rw.trace `egg.frontend
       if req.cfg.genNatLitRws then trace[egg.frontend] "Nat Literal Conversions"