From 7b59c38963846d3a00486700eb4a0c1f3ba2656e Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sat, 28 Dec 2024 13:16:16 -0800
Subject: [PATCH] feat: add `appMap` field to `Goal`

---
 src/Lean/Meta/Tactic/Grind/Internalize.lean | 10 ++++++++++
 src/Lean/Meta/Tactic/Grind/Types.lean       |  7 +++++++
 2 files changed, 17 insertions(+)

diff --git a/src/Lean/Meta/Tactic/Grind/Internalize.lean b/src/Lean/Meta/Tactic/Grind/Internalize.lean
index 5f38815ec011..32d1de928af2 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 3cdbb09176bd..edc96c5db537 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. -/