chore: improve ppState

We also must internalize non constant functions.

src/Lean/Meta/Tactic/Grind.lean

@@ -19,5 +19,7 @@ builtin_initialize registerTraceClass `grind
 builtin_initialize registerTraceClass `grind.eq
 builtin_initialize registerTraceClass `grind.issues
 builtin_initialize registerTraceClass `grind.add
+builtin_initialize registerTraceClass `grind.pre
+builtin_initialize registerTraceClass `grind.debug
 
 end Lean

src/Lean/Meta/Tactic/Grind/Core.lean

@@ -8,6 +8,71 @@ import Lean.Meta.Tactic.Grind.Types
 import Lean.Meta.LitValues
 
 namespace Lean.Meta.Grind
+/-- Helper function for pretty printing the state for debugging purposes. -/
+def ppENodeRef (e : Expr) : GoalM Format := do
+  let some n ← getENode? e | return "_"
+  return f!"#{n.idx}"
+
+/-- Returns expressions in the given expression equivalence class. -/
+partial def getEqc (e : Expr) : GoalM (List Expr) :=
+  go e e []
+where
+  go (first : Expr) (e : Expr) (acc : List Expr) : GoalM (List Expr) := do
+    let next ← getNext e
+    let acc := e :: acc
+    if isSameExpr first next then
+      return acc
+    else
+      go first next acc
+
+/-- Returns all equivalence classes in the current goal. -/
+partial def getEqcs : GoalM (List (List Expr)) := do
+  let mut r := []
+  for (_, node) in (← get).enodes do
+    if isSameExpr node.root node.self then
+      r := (← getEqc node.self) :: r
+  return r
+
+/-- Helper function for pretty printing the state for debugging purposes. -/
+def ppENodeDeclValue (e : Expr) : GoalM Format := do
+  if e.isApp && !(← isLitValue e) then
+    e.withApp fun f args => do
+      let r ← if f.isConst then
+        ppExpr f
+      else
+        ppENodeRef f
+      let mut r := r
+      for arg in args do
+        r := r ++ " " ++ (← ppENodeRef arg)
+      return r
+  else
+    ppExpr e
+
+/-- Helper function for pretty printing the state for debugging purposes. -/
+def ppENodeDecl (e : Expr) : GoalM Format := do
+  let mut r := f!"{← ppENodeRef e} := {← ppENodeDeclValue e}"
+  let n ← getENode e
+  unless isSameExpr e n.root do
+    r := r ++ f!" ↦ {← ppENodeRef n.root}"
+  if n.interpreted then
+    r := r ++ ", [val]"
+  if n.ctor then
+    r := r ++ ", [ctor]"
+  return r
+
+/-- Pretty print goal state for debugging purposes. -/
+def ppState : GoalM Format := do
+  let mut r := f!"Goal:"
+  let nodes := (← get).enodes.toArray.map (·.2)
+  let nodes := nodes.qsort fun a b => a.idx < b.idx
+  for node in nodes do
+    r := r ++ "\n" ++ (← ppENodeDecl node.self)
+  let eqcs ← getEqcs
+  for eqc in eqcs do
+    if eqc.length > 1 then
+      r := r ++ "\n" ++ "{" ++ (Format.joinSep (← eqc.mapM ppENodeRef) ", ") ++  "}"
+  return r
+
 /--
 Returns `true` if `e` is `True`, `False`, or a literal value.
 See `LitValues` for supported literals.
@@ -26,20 +91,6 @@ def mkENode (e : Expr) (generation : Nat) : GoalM Unit := do
   let interpreted ← isInterpreted e
   mkENodeCore e interpreted ctor generation
 
-/--
-Returns the root element in the equivalence class of `e`.
--/
-def getRoot (e : Expr) : GoalM Expr := do
-  let some n ← getENode? e | return e
-  return n.root
-
-/--
-Returns the next element in the equivalence class of `e`.
--/
-def getNext (e : Expr) : GoalM Expr := do
-  let some n ← getENode? e | return e
-  return n.next
-
 private def pushNewEqCore (lhs rhs proof : Expr) (isHEq : Bool) : GoalM Unit :=
   modify fun s => { s with newEqs := s.newEqs.push { lhs, rhs, proof, isHEq } }
 
@@ -70,22 +121,27 @@ partial def internalize (e : Expr) (generation : Nat) : GoalM Unit := do
   | .proj .. =>
     trace[grind.issues] "unexpected term during internalization{indentExpr e}"
     mkENodeCore e (ctor := false) (interpreted := false) (generation := generation)
-  | .app .. => e.withApp fun f args => do
-    let congrThm ← mkHCongrWithArity f args.size
-    let info ← getFunInfo f
-    let shouldInternalize (i : Nat) : GoalM Bool := do
-      if h : i < info.paramInfo.size then
-        let pinfo := info.paramInfo[i]
-        if pinfo.binderInfo.isInstImplicit || pinfo.isProp then
-          return false
-      return true
-    for h : i in [: args.size] do
-      let arg := args[i]
-      if (← shouldInternalize i) then
-        unless (← isTypeFormerType arg) do
-          internalize arg generation
-    mkENode e generation
-    addCongrTable e
+  | .app .. =>
+    if (← isLitValue e) then
+      -- We do not want to internalize the components of a literal value.
+      mkENode e generation
+    else e.withApp fun f args => do
+      unless f.isConst do
+        internalize f generation
+      let info ← getFunInfo f
+      let shouldInternalize (i : Nat) : GoalM Bool := do
+        if h : i < info.paramInfo.size then
+          let pinfo := info.paramInfo[i]
+          if pinfo.binderInfo.isInstImplicit || pinfo.isProp then
+            return false
+        return true
+      for h : i in [: args.size] do
+        let arg := args[i]
+        if (← shouldInternalize i) then
+          unless (← isTypeFormer arg) do
+            internalize arg generation
+      mkENode e generation
+      addCongrTable e
 
 /--
 The fields `target?` and `proof?` in `e`'s `ENode` are encoding a transitivity proof
@@ -98,7 +154,7 @@ private partial def invertTrans (e : Expr) : GoalM Unit := do
   go e false none none
 where
   go (e : Expr) (flippedNew : Bool) (targetNew? : Option Expr) (proofNew? : Option Expr) : GoalM Unit := do
-    let some node ← getENode? e | unreachable!
+    let node ← getENode e
     if let some target := node.target? then
       go target (!node.flipped) (some e) node.proof?
     setENode e { node with
@@ -114,20 +170,25 @@ def isInconsistent : GoalM Bool :=
   return (← get).inconsistent
 
 private partial def addEqStep (lhs rhs proof : Expr) (isHEq : Bool) : GoalM Unit := do
-  trace[grind.eq] "{lhs} {if isHEq then "=" else "≡"} {rhs}"
-  let some lhsNode ← getENode? lhs | return () -- `lhs` has not been internalized yet
-  let some rhsNode ← getENode? rhs | return () -- `rhs` has not been internalized yet
-  if isSameExpr lhsNode.root rhsNode.root then return () -- `lhs` and `rhs` are already in the same equivalence class.
-  let some lhsRoot ← getENode? lhsNode.root | unreachable!
-  let some rhsRoot ← getENode? rhsNode.root | unreachable!
+  trace[grind.eq] "{lhs} {if isHEq then "≡" else "="} {rhs}"
+  let lhsNode ← getENode lhs
+  let rhsNode ← getENode rhs
+  if isSameExpr lhsNode.root rhsNode.root then
+    -- `lhs` and `rhs` are already in the same equivalence class.
+    trace[grind.debug] "{← ppENodeRef lhs} and {← ppENodeRef rhs} are already in the same equivalence class"
+    return ()
+  let lhsRoot ← getENode lhsNode.root
+  let rhsRoot ← getENode rhsNode.root
   if    (lhsRoot.interpreted && !rhsRoot.interpreted)
      || (lhsRoot.ctor && !rhsRoot.ctor)
      || (lhsRoot.size > rhsRoot.size && !rhsRoot.interpreted && !rhsRoot.ctor) then
     go rhs lhs rhsNode lhsNode rhsRoot lhsRoot true
   else
     go lhs rhs lhsNode rhsNode lhsRoot rhsRoot false
+  trace[grind.debug] "after addEqStep, {← ppState}"
 where
   go (lhs rhs : Expr) (lhsNode rhsNode lhsRoot rhsRoot : ENode) (flipped : Bool) : GoalM Unit := do
+    trace[grind.debug] "adding {← ppENodeRef lhs} ↦ {← ppENodeRef rhs}"
     let mut valueInconsistency := false
     if lhsRoot.interpreted && rhsRoot.interpreted then
       if lhsNode.root.isTrue || rhsNode.root.isTrue then
@@ -152,9 +213,11 @@ where
     -- TODO: Remove parents from congruence table
     -- TODO: set propagateBool
     updateRoots lhs rhsNode.root true -- TODO
+    trace[grind.debug] "{← ppENodeRef lhs} new root {← ppENodeRef rhsNode.root}, {← ppENodeRef (← getRoot lhs)}"
     -- TODO: Reinsert parents into congruence table
     setENode lhsNode.root { lhsRoot with
       next := rhsRoot.next
+      root := rhsNode.root
     }
     setENode rhsNode.root { rhsRoot with
       next := lhsRoot.next
@@ -167,7 +230,7 @@ where
   updateRoots (lhs : Expr) (rootNew : Expr) (_propagateBool : Bool) : GoalM Unit := do
     let rec loop (e : Expr) : GoalM Unit := do
       -- TODO: propagateBool
-      let some n ← getENode? e | unreachable!
+      let n ← getENode e
       setENode e { n with root := rootNew }
       if isSameExpr lhs n.next then return ()
       loop n.next
@@ -216,7 +279,7 @@ where
       if isNeg then
         addEq p (← getFalseExpr) (← mkEqFalse proof)
       else
-        addEq p (← getFalseExpr) (← mkEqTrue proof)
+        addEq p (← getTrueExpr) (← mkEqTrue proof)
 
   goEq (p : Expr) (lhs rhs : Expr) (isNeg : Bool) (isHEq : Bool) : GoalM Unit := do
     if isNeg then
@@ -233,28 +296,4 @@ Adds a new hypothesis.
 def addHyp (fvarId : FVarId) (generation := 0) : GoalM Unit := do
   add (← fvarId.getType) (mkFVar fvarId) generation
 
-/--
-Returns expressions in the given expression equivalence class.
--/
-partial def getEqc (e : Expr) : GoalM (List Expr) :=
-  go e e []
-where
-  go (first : Expr) (e : Expr) (acc : List Expr) : GoalM (List Expr) := do
-    let next ← getNext e
-    let acc := e :: acc
-    if isSameExpr e next then
-      return acc
-    else
-      go first next acc
-
-/--
-Returns all equivalence classes in the current goal.
--/
-partial def getEqcs : GoalM (List (List Expr)) := do
-  let mut r := []
-  for (_, node) in (← get).enodes do
-    if isSameExpr node.root node.self then
-      r := (← getEqc node.self) :: r
-  return r
-
 end Lean.Meta.Grind

src/Lean/Meta/Tactic/Grind/Preprocessor.lean

@@ -145,8 +145,17 @@ partial def loop (goal : Goal) : PreM Unit := do
     else
       loop goal
 
+def ppGoals : PreM Format := do
+  let mut r := f!""
+  for goal in (← get).goals do
+    let (f, _) ← GoalM.run goal ppState
+    r := r ++ Format.line ++ f
+  return r
+
 def preprocess (mvarId : MVarId) : PreM State := do
   loop (← mkGoal mvarId)
+  if (← isTracingEnabledFor `grind.pre) then
+    trace[grind.pre] (← ppGoals)
   get
 
 end Preprocessor

src/Lean/Meta/Tactic/Grind/Types.lean

@@ -161,6 +161,7 @@ structure ENode where
   /-- Modification time -/
   mt : Nat := 0
   -- TODO: see Lean 3 implementation
+  deriving Inhabited, Repr
 
 structure NewEq where
   lhs   : Expr
@@ -199,6 +200,19 @@ Otherwise, returns `none`s.
 def getENode? (e : Expr) : GoalM (Option ENode) :=
   return (← get).enodes.find? (unsafe ptrAddrUnsafe e)
 
+/-- Returns node associated with `e`. It assumes `e` has already been internalized. -/
+def getENode (e : Expr) : GoalM ENode := do
+  let some n := (← get).enodes.find? (unsafe ptrAddrUnsafe e) | unreachable!
+  return n
+
+/-- Returns the root element in the equivalence class of `e`. -/
+def getRoot (e : Expr) : GoalM Expr :=
+  return (← getENode e).root
+
+/-- Returns the next element in the equivalence class of `e`. -/
+def getNext (e : Expr) : GoalM Expr :=
+  return (← getENode e).next
+
 /-- Returns `true` if `e` has already been internalized. -/
 def alreadyInternalized (e : Expr) : GoalM Bool :=
   return (← get).enodes.contains (unsafe ptrAddrUnsafe e)

tests/lean/run/grind_pre.lean

@@ -108,3 +108,10 @@ theorem ex3 (h : a₁ :: { x := a₂, y := a₃ : Point } :: as = b₁ :: { x :=
   grind_pre
   trace_state
   sorry
+
+def h (a : α) := a
+
+set_option trace.grind.pre true
+example (p : Prop) (a b c : Nat) : p → a = 0 → a = b → h a = h c → a = c ∧ c = a → a = b ∧ b = a → a ≠ c := by
+  grind_pre
+  sorry