From 59a09fb4e7321c7848c6559eab448a2e501f26a2 Mon Sep 17 00:00:00 2001
From: Joachim Breitner <mail@joachim-breitner.de>
Date: Mon, 17 Jun 2024 20:22:28 +0200
Subject: [PATCH] feat: use priorities to ensure simp applies eqational lemmas
 in order (#4434)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This assigns priorities to the equational lemmas so that more specific
ones
are tried first before a possible catch-all with possible
side-conditions.

We assign very low priorities to match the simplifiers behavior when
unfolding
a definition, which happens in `simpLoop`’ `visitPreContinue` after
applying
rewrite rules.

Definitions with more than 100 equational theorems will use priority 1
for all
but the last (a heuristic, not perfect).

fixes #4173, to some extent.
---
 src/Init/Tactics.lean                       |  1 +
 src/Lean/Meta/Tactic/Simp/SimpTheorems.lean | 21 ++++++++++-
 tests/lean/run/eqnsPrio.lean                | 41 +++++++++++++++++++++
 3 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 tests/lean/run/eqnsPrio.lean

diff --git a/src/Init/Tactics.lean b/src/Init/Tactics.lean
index d844392858..3eb891650f 100644
--- a/src/Init/Tactics.lean
+++ b/src/Init/Tactics.lean
@@ -1459,6 +1459,7 @@ have been simplified by using the modifier `↓`. Here is an example
 ```
 
 When multiple simp theorems are applicable, the simplifier uses the one with highest priority.
+The equational theorems of function are applied at very low priority (100 and below).
 If there are several with the same priority, it is uses the "most recent one". Example:
 ```lean
 @[simp high] theorem cond_true (a b : α) : cond true a b = a := rfl
diff --git a/src/Lean/Meta/Tactic/Simp/SimpTheorems.lean b/src/Lean/Meta/Tactic/Simp/SimpTheorems.lean
index 6b7b9e3578..51c55e24c7 100644
--- a/src/Lean/Meta/Tactic/Simp/SimpTheorems.lean
+++ b/src/Lean/Meta/Tactic/Simp/SimpTheorems.lean
@@ -466,8 +466,25 @@ def SimpTheorems.add (s : SimpTheorems) (id : Origin) (levelParams : Array Name)
 def SimpTheorems.addDeclToUnfold (d : SimpTheorems) (declName : Name) : MetaM SimpTheorems := do
   if let some eqns ← getEqnsFor? declName then
     let mut d := d
-    for eqn in eqns do
-      d ← SimpTheorems.addConst d eqn
+    for h : i in [:eqns.size] do
+      let eqn := eqns[i]
+      /-
+      We assign priorities to the equational lemmas so that more specific ones
+      are tried first before a possible catch-all with possible side-conditions.
+
+      We assign very low priorities to match the simplifiers behavior when unfolding
+      a definition, which happens in `simpLoop`’ `visitPreContinue` after applying
+      rewrite rules.
+
+      Definitions with more than 100 equational theorems will use priority 1 for all
+      but the last (a heuristic, not perfect).
+      -/
+      let prio := if eqns.size > 100 then
+        if i + 1 = eqns.size then 0 else 1
+      else
+        100 - i
+      -- We assign very low priority to equational le
+      d ← SimpTheorems.addConst d eqn (prio := prio)
     /-
     Even if a function has equation theorems,
     we also store it in the `toUnfold` set in the following two cases:
diff --git a/tests/lean/run/eqnsPrio.lean b/tests/lean/run/eqnsPrio.lean
new file mode 100644
index 0000000000..97613ad6e8
--- /dev/null
+++ b/tests/lean/run/eqnsPrio.lean
@@ -0,0 +1,41 @@
+import Lean
+
+/-!
+Tests that simp applies the equational lemmas in order. In particular,
+a catch-all at the end is tried afte the others
+-/
+
+def foo : Bool → Nat → Nat
+  | _, 0 => 0
+  | .true, n+1 => foo .true n
+  | _, n+1 => foo .false n
+termination_by _ n => n
+
+/-- info: foo.eq_1 : ∀ (x : Bool), foo x 0 = 0 -/
+#guard_msgs in
+#check foo.eq_1
+/-- info: foo.eq_2 (n : Nat) : foo true n.succ = foo true n -/
+#guard_msgs in
+#check foo.eq_2
+/-- info: foo.eq_3 : ∀ (x : Bool) (n : Nat), (x = true → False) → foo x n.succ = foo false n -/
+#guard_msgs in
+#check foo.eq_3
+
+-- In order to reliably check if simp is not attempting to rewrite with a certain lemma
+-- we can look at te diagnostics. But simply dumping all diangostics is too noisy for a test,
+-- so here we try to get our hands at the `Simp.Stats` and look there.
+open Lean Meta Elab Tactic in
+elab "simp_foo_with_check" : tactic =>
+  withOptions (fun o => diagnostics.set o true) do withMainContext do
+    let stx ← `(tactic|simp [foo])
+    let { ctx, simprocs, dischargeWrapper } ← mkSimpContext stx (eraseLocal := false)
+    let stats ← dischargeWrapper.with fun discharge? => do
+      simpLocation ctx simprocs discharge? (expandOptLocation stx.raw[5])
+    unless stats.diag.triedThmCounter.toList.any (fun (o, _n) => o.key = ``foo.eq_2) do
+        throwError "Simp did not try to use foo.eq_2, is this test still working?"
+    for (origin, n) in stats.diag.triedThmCounter.toList do
+      if origin.key = ``foo.eq_3 then
+        throwError m!"Bad: Simp used foo.eq_3 {n} times"
+
+example : foo true 1 = 0 := by
+  simp_foo_with_check -- essentially simp [foo]