From 22c81548116e0829b6c335e91c247bed2bb06ac2 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Tue, 26 Dec 2023 17:02:09 -0800
Subject: [PATCH] feat: add pre simp lemmas for if-then-else terms

See new test for example that takes exponential time without new simp
theorems.
TODO: replace auxiliary theorems with simprocs as soon as we implement them.
---
 src/Init/SimpLemmas.lean                   |  4 ++++
 tests/lean/1079.lean.expected.out          | 15 +++++++++-----
 tests/lean/run/simpIfPre.lean              | 23 ++++++++++++++++++++++
 tests/lean/simpZetaFalse.lean.expected.out |  3 ++-
 4 files changed, 39 insertions(+), 6 deletions(-)
 create mode 100644 tests/lean/run/simpIfPre.lean

diff --git a/src/Init/SimpLemmas.lean b/src/Init/SimpLemmas.lean
index 4e1cbfa809..cc7de05f71 100644
--- a/src/Init/SimpLemmas.lean
+++ b/src/Init/SimpLemmas.lean
@@ -84,6 +84,10 @@ theorem dite_congr {_ : Decidable b} [Decidable c]
 @[simp] theorem ite_false (a b : α) : (if False then a else b) = b := rfl
 @[simp] theorem dite_true {α : Sort u} {t : True → α} {e : ¬ True → α} : (dite True t e) = t True.intro := rfl
 @[simp] theorem dite_false {α : Sort u} {t : False → α} {e : ¬ False → α} : (dite False t e) = e not_false := rfl
+@[simp ↓] theorem ite_cond_true {_ : Decidable c} (a b : α) (h : c) : (if c then a else b) = a := by simp [h]
+@[simp ↓] theorem ite_cond_false {_ : Decidable c} (a b : α) (h : ¬ c) : (if c then a else b) = b := by simp [h]
+@[simp ↓] theorem dite_cond_true {α : Sort u} {_ : Decidable c} {t : c → α} {e : ¬ c → α} (h : c) : (dite c t e) = t h := by simp [h]
+@[simp ↓] theorem dite_cond_false {α : Sort u} {_ : Decidable c} {t : c → α} {e : ¬ c → α} (h : ¬ c) : (dite c t e) = e h := by simp [h]
 @[simp] theorem ite_self {α : Sort u} {c : Prop} {d : Decidable c} (a : α) : ite c a a = a := by cases d <;> rfl
 @[simp] theorem and_self (p : Prop) : (p ∧ p) = p := propext ⟨(·.1), fun h => ⟨h, h⟩⟩
 @[simp] theorem and_true (p : Prop) : (p ∧ True) = p := propext ⟨(·.1), (⟨·, trivial⟩)⟩
diff --git a/tests/lean/1079.lean.expected.out b/tests/lean/1079.lean.expected.out
index e01a614637..06e23eede1 100644
--- a/tests/lean/1079.lean.expected.out
+++ b/tests/lean/1079.lean.expected.out
@@ -1,11 +1,16 @@
 1079.lean:3:2-6:12: error: alternative 'isFalse' has not been provided
+[Meta.Tactic.simp.discharge] >> discharge?: m = n
+[Meta.Tactic.simp.unify] @eq_self:1000, failed to unify
+      ?a = ?a
+    with
+      m = n
+[Meta.Tactic.simp.unify] Nat.succ.injEq:1000, failed to unify
+      Nat.succ ?n = Nat.succ ?n
+    with
+      m = n
 [Meta.Tactic.simp.rewrite] h:1000, m ==> n
 [Meta.Tactic.simp.rewrite] @eq_self:1000, n = n ==> True
-[Meta.Tactic.simp.unify] @ite_self:1000, failed to unify
-      if ?c then ?a else ?a
-    with
-      if True then Ordering.eq else Ordering.gt
-[Meta.Tactic.simp.rewrite] @ite_true:1000, if True then Ordering.eq else Ordering.gt ==> Ordering.eq
+[Meta.Tactic.simp.rewrite] @ite_cond_true:1000, if m = n then Ordering.eq else Ordering.gt ==> Ordering.eq
 [Meta.Tactic.simp.unify] @eq_self:1000, failed to unify
       ?a = ?a
     with
diff --git a/tests/lean/run/simpIfPre.lean b/tests/lean/run/simpIfPre.lean
new file mode 100644
index 0000000000..3b9c5245a9
--- /dev/null
+++ b/tests/lean/run/simpIfPre.lean
@@ -0,0 +1,23 @@
+/-!
+Test support for `if-then-else` terms in the simplifier.
+The condition should be simplified before trying to apply congruence.
+We are currently accomplished that using pre-simp theorems.
+TODO: replace them with simprocs.
+In the following example, the term `g (a + <num>)` takes an
+exponential amount of time to be simplified without the pre-simp theorems.
+-/
+
+def myid (x : Nat) := 0 + x
+
+@[simp] theorem myid_eq : myid x = x := by simp [myid]
+
+def f (x : Nat) (y z : Nat) : Nat :=
+  if myid x = 0 then y else z
+
+def g (x : Nat) : Nat :=
+  match x with
+  | 0 => 1
+  | a+1 => f x (g a + 1) (g a)
+
+theorem ex (h : a = 1) : g (a+32) = a := by
+  simp [g, f, h]
diff --git a/tests/lean/simpZetaFalse.lean.expected.out b/tests/lean/simpZetaFalse.lean.expected.out
index b532343f64..0af829e1a7 100644
--- a/tests/lean/simpZetaFalse.lean.expected.out
+++ b/tests/lean/simpZetaFalse.lean.expected.out
@@ -17,7 +17,8 @@ fun x h =>
             ite_congr (Eq.trans (congrFun (congrArg Eq h) x) (eq_self x)) (fun a => Eq.refl 1) fun a =>
               Eq.refl (y + 1)))
         1))
-    (of_eq_true (eq_self 1))
+    (of_eq_true
+      (Eq.trans (congrFun (congrArg Eq (ite_cond_true 1 (x * x + 1) (of_eq_true (Eq.refl True)))) 1) (eq_self 1)))
 x z : Nat
 h : f (f x) = x
 h' : z = x