From 2273527a41fe6f5456ba125149021bfbbf8cc9f3 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Fri, 27 Dec 2024 15:36:19 -0800
Subject: [PATCH] feat: add `propagateProjEq`

---
 src/Lean/Meta/Tactic/Grind/Proj.lean | 35 ++++++++++++++++++++++++++++
 src/Lean/Meta/Tactic/Grind/Run.lean  |  2 ++
 tests/lean/run/grind_t1.lean         | 20 ++++++++++++++++
 3 files changed, 57 insertions(+)
 create mode 100644 src/Lean/Meta/Tactic/Grind/Proj.lean

diff --git a/src/Lean/Meta/Tactic/Grind/Proj.lean b/src/Lean/Meta/Tactic/Grind/Proj.lean
new file mode 100644
index 000000000000..b742642b2f5f
--- /dev/null
+++ b/src/Lean/Meta/Tactic/Grind/Proj.lean
@@ -0,0 +1,35 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Lean.ProjFns
+import Lean.Meta.Tactic.Grind.Types
+import Lean.Meta.Tactic.Grind.Internalize
+
+namespace Lean.Meta.Grind
+
+/--
+If `parent` is a projection-application `proj_i c`,
+check whether the root of the equivalence class containing `c` is a constructor-application `ctor ... a_i ...`.
+If so, internalize the term `proj_i (ctor ... a_i ...)` and add the equality `proj_i (ctor ... a_i ...) = a_i`.
+-/
+def propagateProjEq (parent : Expr) : GoalM Unit := do
+  let .const declName _ := parent.getAppFn | return ()
+  let some info ← getProjectionFnInfo? declName | return ()
+  unless info.numParams + 1 == parent.getAppNumArgs do return ()
+  let arg := parent.appArg!
+  let ctor ← getRoot arg
+  unless ctor.isAppOf info.ctorName do return ()
+  if isSameExpr arg ctor then
+    let idx := info.numParams + info.i
+    unless idx < ctor.getAppNumArgs do return ()
+    let v := ctor.getArg! idx
+    pushEq parent v (← mkEqRefl v)
+  else
+    let newProj := mkApp parent.appFn! ctor
+    let newProj ← shareCommon newProj
+    internalize newProj (← getGeneration parent)
+
+end Lean.Meta.Grind
diff --git a/src/Lean/Meta/Tactic/Grind/Run.lean b/src/Lean/Meta/Tactic/Grind/Run.lean
index 52e5c9af8b13..f2b534b4557e 100644
--- a/src/Lean/Meta/Tactic/Grind/Run.lean
+++ b/src/Lean/Meta/Tactic/Grind/Run.lean
@@ -7,6 +7,7 @@ prelude
 import Init.Grind.Lemmas
 import Lean.Meta.Tactic.Grind.Types
 import Lean.Meta.Tactic.Grind.PropagatorAttr
+import Lean.Meta.Tactic.Grind.Proj
 
 namespace Lean.Meta.Grind
 
@@ -15,6 +16,7 @@ def mkMethods : CoreM Methods := do
   return {
     propagateUp := fun e => do
      let .const declName _ := e.getAppFn | return ()
+     propagateProjEq e
      if let some prop := builtinPropagators.up[declName]? then
        prop e
     propagateDown := fun e => do
diff --git a/tests/lean/run/grind_t1.lean b/tests/lean/run/grind_t1.lean
index 0d4b9ada9813..59f8749a8775 100644
--- a/tests/lean/run/grind_t1.lean
+++ b/tests/lean/run/grind_t1.lean
@@ -55,3 +55,23 @@ example (a b c : BitVec 32) : a = c → a = 1#32 → c = 2#32 → c = b → Fals
 
 example (a b c : UInt32) : a = c → a = 1 → c = 200 → c = b → False := by
   grind
+
+structure Boo (α : Type) where
+  a : α
+  b : α
+  c : α
+
+example (a b d : Nat) (f : Nat → Boo Nat) : (f d).1 ≠ a → f d = ⟨b, v₁, v₂⟩ → b = a → False := by
+  grind
+
+def ex (a b c d : Nat) (f : Nat → Boo Nat) : (f d).2 ≠ a → f d = ⟨b, c, v₂⟩ → c = a → False := by
+  grind
+
+example (a b c : Nat) (f : Nat → Nat) : { a := f b, c, b := 4 : Boo Nat }.1 ≠ f a → f b = f c → a = c → False := by
+  grind
+
+example (a b c : Nat) (f : Nat → Nat) : p = { a := f b, c, b := 4 : Boo Nat } → p.1 ≠ f a → f b = f c → a = c → False := by
+  grind
+
+example (a b c : Nat) (f : Nat → Nat) : p.1 ≠ f a → p = { a := f b, c, b := 4 : Boo Nat } → f b = f c → a = c → False := by
+  grind