From 51ef8704d1e0320cdaeffc80994fb58c0d8d83c5 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Wed, 1 Jan 2025 15:48:33 -0800
Subject: [PATCH] feat: `forall` propagator

---
 src/Init/Grind/Lemmas.lean                  |  8 ++++++
 src/Lean/Meta/Tactic/Grind/ForallProp.lean  | 31 +++++++++++++++++++++
 src/Lean/Meta/Tactic/Grind/Internalize.lean |  8 +++++-
 src/Lean/Meta/Tactic/Grind/Run.lean         |  2 ++
 tests/lean/run/grind_prop_arrow.lean        |  6 ++++
 5 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 src/Lean/Meta/Tactic/Grind/ForallProp.lean
 create mode 100644 tests/lean/run/grind_prop_arrow.lean

diff --git a/src/Init/Grind/Lemmas.lean b/src/Init/Grind/Lemmas.lean
index 33d32fbe5b07..4372e8c91848 100644
--- a/src/Init/Grind/Lemmas.lean
+++ b/src/Init/Grind/Lemmas.lean
@@ -8,6 +8,7 @@ import Init.Core
 import Init.SimpLemmas
 import Init.Classical
 import Init.ByCases
+import Init.Grind.Util
 
 namespace Lean.Grind
 
@@ -50,4 +51,11 @@ theorem false_of_not_eq_self {a : Prop} (h : (Not a) = a) : False := by
 theorem eq_eq_of_eq_true_left {a b : Prop} (h : a = True) : (a = b) = b := by simp [h]
 theorem eq_eq_of_eq_true_right {a b : Prop} (h : b = True) : (a = b) = a := by simp [h]
 
+/-! Forall -/
+
+theorem forall_propagator (p : Prop) (q : p → Prop) (q' : Prop) (h₁ : p = True) (h₂ : q (of_eq_true h₁) = q') : (∀ hp : p, q hp) = q' := by
+  apply propext; apply Iff.intro
+  · intro h'; exact Eq.mp h₂ (h' (of_eq_true h₁))
+  · intro h'; intros; exact Eq.mpr h₂ h'
+
 end Lean.Grind
diff --git a/src/Lean/Meta/Tactic/Grind/ForallProp.lean b/src/Lean/Meta/Tactic/Grind/ForallProp.lean
new file mode 100644
index 000000000000..d6c42053ea06
--- /dev/null
+++ b/src/Lean/Meta/Tactic/Grind/ForallProp.lean
@@ -0,0 +1,31 @@
+/-
+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 Init.Grind.Lemmas
+import Lean.Meta.Tactic.Grind.Types
+import Lean.Meta.Tactic.Grind.Internalize
+import Lean.Meta.Tactic.Grind.Simp
+
+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 propagateForallProp (parent : Expr) : GoalM Unit := do
+  let .forallE n p q bi := parent | return ()
+  unless (← isEqTrue p) do return ()
+  let h₁ ← mkEqTrueProof p
+  let qh₁ := q.instantiate1 (mkApp2 (mkConst ``of_eq_true) p h₁)
+  let r ← pre qh₁
+  let q := mkLambda n bi p q
+  let q' := r.expr
+  internalize q' (← getGeneration parent)
+  let h₂ ← r.getProof
+  let h := mkApp5 (mkConst ``Lean.Grind.forall_propagator) p q q' h₁ h₂
+  pushEq parent q' h
+
+end Lean.Meta.Grind
diff --git a/src/Lean/Meta/Tactic/Grind/Internalize.lean b/src/Lean/Meta/Tactic/Grind/Internalize.lean
index 75d650ea2ad8..fb4986404eb9 100644
--- a/src/Lean/Meta/Tactic/Grind/Internalize.lean
+++ b/src/Lean/Meta/Tactic/Grind/Internalize.lean
@@ -75,8 +75,14 @@ partial def internalize (e : Expr) (generation : Nat) : GoalM Unit := do
   match e with
   | .bvar .. => unreachable!
   | .sort .. => return ()
-  | .fvar .. | .letE .. | .lam .. | .forallE .. =>
+  | .fvar .. | .letE .. | .lam .. =>
     mkENodeCore e (ctor := false) (interpreted := false) (generation := generation)
+  | .forallE _ d _ _ =>
+    mkENodeCore e (ctor := false) (interpreted := false) (generation := generation)
+    if (← isProp d <&&> isProp e) then
+      internalize d generation
+      registerParent e d
+      propagateUp e
   | .lit .. | .const .. =>
     mkENode e generation
   | .mvar ..
diff --git a/src/Lean/Meta/Tactic/Grind/Run.lean b/src/Lean/Meta/Tactic/Grind/Run.lean
index 2f723dd52c17..8626591cad8a 100644
--- a/src/Lean/Meta/Tactic/Grind/Run.lean
+++ b/src/Lean/Meta/Tactic/Grind/Run.lean
@@ -8,6 +8,7 @@ import Init.Grind.Lemmas
 import Lean.Meta.Tactic.Grind.Types
 import Lean.Meta.Tactic.Grind.PropagatorAttr
 import Lean.Meta.Tactic.Grind.Proj
+import Lean.Meta.Tactic.Grind.ForallProp
 
 namespace Lean.Meta.Grind
 
@@ -15,6 +16,7 @@ def mkMethods : CoreM Methods := do
   let builtinPropagators ← builtinPropagatorsRef.get
   return {
     propagateUp := fun e => do
+     propagateForallProp e
      let .const declName _ := e.getAppFn | return ()
      propagateProjEq e
      if let some prop := builtinPropagators.up[declName]? then
diff --git a/tests/lean/run/grind_prop_arrow.lean b/tests/lean/run/grind_prop_arrow.lean
new file mode 100644
index 000000000000..2f29830ae212
--- /dev/null
+++ b/tests/lean/run/grind_prop_arrow.lean
@@ -0,0 +1,6 @@
+opaque f (a : Array Bool) (i : Nat) (h : i < a.size) : Bool
+
+set_option trace.grind.eqc true
+
+example : (p ∨ ∀ h : i < a.size, f a i h) → (hb : i < b.size) → a = b → ¬p → f b i hb := by
+  grind