From a781f9858c4ab2dcbb8947e95bcd3e9e49370fcf Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sun, 29 Dec 2024 01:33:25 +0100
Subject: [PATCH] feat: missing data for `grind` e-match (#6469)

This PR adds support code for implementing e-match in the (WIP) `grind`
tactic.
---
 src/Lean/Meta/Tactic/Grind/Core.lean        | 14 ++++++++++++++
 src/Lean/Meta/Tactic/Grind/Internalize.lean | 10 ++++++++++
 src/Lean/Meta/Tactic/Grind/Types.lean       |  7 +++++++
 3 files changed, 31 insertions(+)

diff --git a/src/Lean/Meta/Tactic/Grind/Core.lean b/src/Lean/Meta/Tactic/Grind/Core.lean
index 5c4755345e..06c69ba3b5 100644
--- a/src/Lean/Meta/Tactic/Grind/Core.lean
+++ b/src/Lean/Meta/Tactic/Grind/Core.lean
@@ -69,6 +69,18 @@ private def closeGoalWithValuesEq (lhs rhs : Expr) : GoalM Unit := do
   let falseProof ← mkEqMP pEqFalse hp
   closeGoal falseProof
 
+/--
+Updates the modification time to `gmt` for the parents of `root`.
+The modification time is used to decide which terms are considered during e-matching.
+-/
+private partial def updateMT (root : Expr) : GoalM Unit := do
+  let gmt := (← get).gmt
+  for parent in (← getParents root) do
+    let node ← getENode parent
+    if node.mt < gmt then
+      setENode parent { node with mt := gmt }
+      updateMT parent
+
 private partial def addEqStep (lhs rhs proof : Expr) (isHEq : Bool) : GoalM Unit := do
   trace[grind.eq] "{lhs} {if isHEq then "≡" else "="} {rhs}"
   let lhsNode ← getENode lhs
@@ -137,6 +149,8 @@ where
     unless (← isInconsistent) do
       for parent in parents do
         propagateUp parent
+    unless (← isInconsistent) do
+      updateMT rhsRoot.self
 
   updateRoots (lhs : Expr) (rootNew : Expr) : GoalM Unit := do
     let rec loop (e : Expr) : GoalM Unit := do
diff --git a/src/Lean/Meta/Tactic/Grind/Internalize.lean b/src/Lean/Meta/Tactic/Grind/Internalize.lean
index 5f38815ec0..32d1de928a 100644
--- a/src/Lean/Meta/Tactic/Grind/Internalize.lean
+++ b/src/Lean/Meta/Tactic/Grind/Internalize.lean
@@ -28,6 +28,15 @@ def addCongrTable (e : Expr) : GoalM Unit := do
   else
     modify fun s => { s with congrTable := s.congrTable.insert { e } }
 
+private def updateAppMap (e : Expr) : GoalM Unit := do
+  let key := e.toHeadIndex
+  modify fun s => { s with
+    appMap := if let some es := s.appMap.find? key then
+      s.appMap.insert key (e :: es)
+    else
+      s.appMap.insert key [e]
+  }
+
 partial def internalize (e : Expr) (generation : Nat) : GoalM Unit := do
   if (← alreadyInternalized e) then return ()
   match e with
@@ -63,6 +72,7 @@ partial def internalize (e : Expr) (generation : Nat) : GoalM Unit := do
           registerParent e arg
       mkENode e generation
       addCongrTable e
+      updateAppMap e
       propagateUp e
 
 end Lean.Meta.Grind
diff --git a/src/Lean/Meta/Tactic/Grind/Types.lean b/src/Lean/Meta/Tactic/Grind/Types.lean
index 3cdbb09176..edc96c5db5 100644
--- a/src/Lean/Meta/Tactic/Grind/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Types.lean
@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Lean.Util.ShareCommon
+import Lean.HeadIndex
 import Lean.Meta.Basic
 import Lean.Meta.CongrTheorems
 import Lean.Meta.AbstractNestedProofs
@@ -258,6 +259,12 @@ structure Goal where
   enodes       : ENodes := {}
   parents      : ParentMap := {}
   congrTable   : CongrTable enodes := {}
+  /--
+  A mapping from each function application index (`HeadIndex`) to a list of applications with that index.
+  Recall that the `HeadIndex` for a constant is its constant name, and for a free variable,
+  it is its unique id.
+  -/
+  appMap       : PHashMap HeadIndex (List Expr) := {}
   /-- Equations to be processed. -/
   newEqs       : Array NewEq := #[]
   /-- `inconsistent := true` if `ENode`s for `True` and `False` are in the same equivalence class. -/