From cde35bcc0d5de9253c7ec8877c675fef8c76ae8d Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Tue, 24 Dec 2024 05:33:05 +0100
Subject: [PATCH] test: `grind` (#6440)

This PR adds additional tests for `grind` and fixed minor issues.
---
 src/Lean/Meta/Tactic/Grind/Types.lean         |  3 +-
 ...CanonInsts.lean => grind_canon_insts.lean} |  2 ++
 ...CanonTypes.lean => grind_canon_types.lean} |  0
 tests/lean/run/grind_many_eqs.lean            | 34 +++++++++++++++++++
 tests/lean/run/grind_nested_proofs.lean       | 26 ++++++++++++++
 ...LevelIssue.lean => grind_norm_levels.lean} |  0
 6 files changed, 64 insertions(+), 1 deletion(-)
 rename tests/lean/run/{grindCanonInsts.lean => grind_canon_insts.lean} (98%)
 rename tests/lean/run/{grindCanonTypes.lean => grind_canon_types.lean} (100%)
 create mode 100644 tests/lean/run/grind_many_eqs.lean
 rename tests/lean/run/{grindNormLevelIssue.lean => grind_norm_levels.lean} (100%)

diff --git a/src/Lean/Meta/Tactic/Grind/Types.lean b/src/Lean/Meta/Tactic/Grind/Types.lean
index a6754eb60646..56a3bc4f5c2b 100644
--- a/src/Lean/Meta/Tactic/Grind/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Types.lean
@@ -257,7 +257,8 @@ def getENode? (e : Expr) : GoalM (Option ENode) :=
 
 /-- 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!
+  let some n := (← get).enodes.find? (unsafe ptrAddrUnsafe e)
+    | throwError "internal `grind` error, term has not been internalized{indentExpr e}"
   return n
 
 /-- Returns `true` is the root of its equivalence class. -/
diff --git a/tests/lean/run/grindCanonInsts.lean b/tests/lean/run/grind_canon_insts.lean
similarity index 98%
rename from tests/lean/run/grindCanonInsts.lean
rename to tests/lean/run/grind_canon_insts.lean
index 59b5b481f48b..ddbadb62f5e9 100644
--- a/tests/lean/run/grindCanonInsts.lean
+++ b/tests/lean/run/grind_canon_insts.lean
@@ -7,6 +7,8 @@ elab "grind_test" : tactic => withMainContext do
     let nodes ← filterENodes fun e => return e.self.isAppOf ``HMul.hMul
     logInfo (nodes.toList.map (·.self))
 
+set_option grind.debug true
+
 class Semigroup (α : Type u) extends Mul α where
   mul_assoc (a b c : α) : a * b * c = a * (b * c)
 
diff --git a/tests/lean/run/grindCanonTypes.lean b/tests/lean/run/grind_canon_types.lean
similarity index 100%
rename from tests/lean/run/grindCanonTypes.lean
rename to tests/lean/run/grind_canon_types.lean
diff --git a/tests/lean/run/grind_many_eqs.lean b/tests/lean/run/grind_many_eqs.lean
new file mode 100644
index 000000000000..5f9825eadab2
--- /dev/null
+++ b/tests/lean/run/grind_many_eqs.lean
@@ -0,0 +1,34 @@
+import Lean
+
+def f (a : Nat) := a + a + a
+def g (a : Nat) := a + a
+def h (n : Nat) : Prop :=
+  match n with
+  | 0   => f 0 = f 1
+  | n+1 => f (n+1) = f n ∧ g (2*n + 1) = g (2*n) ∧ h n
+
+-- Prints the equivalence class containing a `f` application
+open Lean Meta Elab Tactic Grind in
+elab "grind_test" n:num : tactic => withMainContext do
+  let n := n.getNat
+  let declName := (← Term.getDeclName?).getD `_main
+  Meta.Grind.preprocessAndProbe (← getMainGoal) declName do
+    let f0 ← Grind.shareCommon (mkApp (mkConst ``f) (mkNatLit 0))
+    -- The `f 0` equivalence class contains `n+1` elements
+    assert! (← getEqc f0).length == n + 1
+    forEachENode fun node => do
+      if node.self.isAppOf ``g then
+        -- Any equivalence class containing a `g`-application contains 2 elements
+        assert! (← getEqc (← getRoot node.self)).length == 2
+
+set_option grind.debug true in
+example : h 5 → False := by
+  simp [h]
+  grind_test 5
+  sorry
+
+set_option maxRecDepth 2048
+example : h 100 → False := by
+  simp [h]
+  grind_test 100
+  sorry
diff --git a/tests/lean/run/grind_nested_proofs.lean b/tests/lean/run/grind_nested_proofs.lean
index b7989373bd97..b01d80730cfe 100644
--- a/tests/lean/run/grind_nested_proofs.lean
+++ b/tests/lean/run/grind_nested_proofs.lean
@@ -8,17 +8,43 @@ elab "grind_test" : tactic => withMainContext do
   Meta.Grind.preprocessAndProbe (← getMainGoal) declName do
     let nodes ← filterENodes fun e => return e.self.isAppOf ``Lean.Grind.nestedProof
     logInfo (nodes.toList.map (·.self))
+    let nodes ← filterENodes fun e => return e.self.isAppOf ``GetElem.getElem
+    let [_, n, _] := nodes.toList | unreachable!
+    logInfo (← getEqc n.self)
 
 set_option grind.debug true
 
+/-
+Recall that array access terms, such as `a[i]`, have nested proofs.
+The following test relies on `grind` `nestedProof` wrapper to
+detect equalities between array access terms.
+-/
+
 /--
 info: [Lean.Grind.nestedProof (i < a.toList.length) (_example.proof_1 i j a b h1 h2),
  Lean.Grind.nestedProof (j < a.toList.length) h1,
  Lean.Grind.nestedProof (j < b.toList.length) h]
 ---
+info: [a[i], b[j], a[j]]
+---
 warning: declaration uses 'sorry'
 -/
 #guard_msgs in
 example (i j : Nat) (a b : Array Nat) (h1 : j < a.size) (h : j < b.size) (h2 : i ≤ j) : a[i] < a[j] + b[j] → i = j → a = b → False := by
   grind_test
   sorry
+
+
+/--
+info: [Lean.Grind.nestedProof (i < a.toList.length) (_example.proof_1 i j a b h1 h2),
+ Lean.Grind.nestedProof (j < a.toList.length) h1,
+ Lean.Grind.nestedProof (j < b.toList.length) h]
+---
+info: [a[i], a[j]]
+---
+warning: declaration uses 'sorry'
+-/
+#guard_msgs in
+example (i j : Nat) (a b : Array Nat) (h1 : j < a.size) (h : j < b.size) (h2 : i ≤ j) : a[i] < a[j] + b[j] → i = j → False := by
+  grind_test
+  sorry
diff --git a/tests/lean/run/grindNormLevelIssue.lean b/tests/lean/run/grind_norm_levels.lean
similarity index 100%
rename from tests/lean/run/grindNormLevelIssue.lean
rename to tests/lean/run/grind_norm_levels.lean