From 51377afd6c8c2c4e4826993a406be38daaf259bd Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Mon, 21 Oct 2024 18:37:40 +1100
Subject: [PATCH] feat: simp lemmas for Array.isEqv and beq (#5786)

- [ ] depends on: #5785
---
 src/Init/Data/Array/Basic.lean       | 20 +++++++++++
 src/Init/Data/Array/Bootstrap.lean   |  2 +-
 src/Init/Data/Array/DecidableEq.lean | 51 ++++++++++++++++++++++++++++
 src/Init/Data/Array/GetLit.lean      |  2 +-
 src/Init/Data/Array/Lemmas.lean      | 15 ++------
 src/Init/Data/List/Basic.lean        |  5 +++
 src/Init/Data/List/Nat.lean          |  1 +
 src/Init/Data/List/Nat/BEq.lean      | 47 +++++++++++++++++++++++++
 tests/lean/run/array_isEqvAux.lean   |  2 ++
 9 files changed, 131 insertions(+), 14 deletions(-)
 create mode 100644 src/Init/Data/List/Nat/BEq.lean

diff --git a/src/Init/Data/Array/Basic.lean b/src/Init/Data/Array/Basic.lean
index 1db7a0dbd78d..d057e5e87df6 100644
--- a/src/Init/Data/Array/Basic.lean
+++ b/src/Init/Data/Array/Basic.lean
@@ -80,6 +80,26 @@ theorem ext' {as bs : Array α} (h : as.toList = bs.toList) : as = bs := by
 
 @[simp] theorem size_toArray (as : List α) : as.toArray.size = as.length := by simp [size]
 
+@[simp] theorem getElem_toList {a : Array α} {i : Nat} (h : i < a.size) : a.toList[i] = a[i] := rfl
+
+end Array
+
+namespace List
+
+@[simp] theorem toArray_toList (a : Array α) : a.toList.toArray = a := rfl
+
+@[simp] theorem getElem_toArray {a : List α} {i : Nat} (h : i < a.toArray.size) :
+    a.toArray[i] = a[i]'(by simpa using h) := rfl
+
+@[simp] theorem getElem?_toArray {a : List α} {i : Nat} : a.toArray[i]? = a[i]? := rfl
+
+@[simp] theorem getElem!_toArray [Inhabited α] {a : List α} {i : Nat} :
+    a.toArray[i]! = a[i]! := rfl
+
+end List
+
+namespace Array
+
 @[deprecated toList_toArray (since := "2024-09-09")] abbrev data_toArray := @toList_toArray
 
 @[deprecated Array.toList (since := "2024-09-10")] abbrev Array.data := @Array.toList
diff --git a/src/Init/Data/Array/Bootstrap.lean b/src/Init/Data/Array/Bootstrap.lean
index 998e90c377bd..bb322c4c62b1 100644
--- a/src/Init/Data/Array/Bootstrap.lean
+++ b/src/Init/Data/Array/Bootstrap.lean
@@ -42,7 +42,7 @@ theorem foldrM_eq_reverse_foldlM_toList.aux [Monad m]
   unfold foldrM.fold
   match i with
   | 0 => simp [List.foldlM, List.take]
-  | i+1 => rw [← List.take_concat_get _ _ h]; simp [← (aux f arr · i)]; rfl
+  | i+1 => rw [← List.take_concat_get _ _ h]; simp [← (aux f arr · i)]
 
 theorem foldrM_eq_reverse_foldlM_toList [Monad m] (f : α → β → m β) (init : β) (arr : Array α) :
     arr.foldrM f init = arr.toList.reverse.foldlM (fun x y => f y x) init := by
diff --git a/src/Init/Data/Array/DecidableEq.lean b/src/Init/Data/Array/DecidableEq.lean
index 567a7c25a1de..bda267db790f 100644
--- a/src/Init/Data/Array/DecidableEq.lean
+++ b/src/Init/Data/Array/DecidableEq.lean
@@ -6,6 +6,8 @@ Authors: Leonardo de Moura
 prelude
 import Init.Data.Array.Basic
 import Init.Data.BEq
+import Init.Data.Nat.Lemmas
+import Init.Data.List.Nat.BEq
 import Init.ByCases
 
 namespace Array
@@ -26,6 +28,14 @@ theorem rel_of_isEqvAux
       subst hj'
       exact heqv.left
 
+theorem isEqvAux_of_rel (r : α → α → Bool) (a b : Array α) (hsz : a.size = b.size) (i : Nat) (hi : i ≤ a.size)
+    (w : ∀ j, (hj : j < i) → r (a[j]'(Nat.lt_of_lt_of_le hj hi)) (b[j]'(Nat.lt_of_lt_of_le hj (hsz ▸ hi)))) : Array.isEqvAux a b hsz r i hi := by
+  induction i with
+  | zero => simp [Array.isEqvAux]
+  | succ i ih =>
+    simp only [isEqvAux, Bool.and_eq_true]
+    exact ⟨w i (Nat.lt_add_one i), ih _ fun j hj => w j (Nat.lt_add_right 1 hj)⟩
+
 theorem rel_of_isEqv (r : α → α → Bool) (a b : Array α) :
     Array.isEqv a b r → ∃ h : a.size = b.size, ∀ (i : Nat) (h' : i < a.size), r (a[i]) (b[i]'(h ▸ h')) := by
   simp only [isEqv]
@@ -33,6 +43,29 @@ theorem rel_of_isEqv (r : α → α → Bool) (a b : Array α) :
   · exact fun h' => ⟨h, rel_of_isEqvAux r a b h a.size (Nat.le_refl ..) h'⟩
   · intro; contradiction
 
+theorem isEqv_iff_rel (a b : Array α) (r) :
+    Array.isEqv a b r ↔ ∃ h : a.size = b.size, ∀ (i : Nat) (h' : i < a.size), r (a[i]) (b[i]'(h ▸ h')) :=
+  ⟨rel_of_isEqv r a b, fun ⟨h, w⟩ => by
+    simp only [isEqv, ← h, ↓reduceDIte]
+    exact isEqvAux_of_rel r a b h a.size (by simp [h]) w⟩
+
+theorem isEqv_eq_decide (a b : Array α) (r) :
+    Array.isEqv a b r =
+      if h : a.size = b.size then decide (∀ (i : Nat) (h' : i < a.size), r (a[i]) (b[i]'(h ▸ h'))) else false := by
+  by_cases h : Array.isEqv a b r
+  · simp only [h, Bool.true_eq]
+    simp only [isEqv_iff_rel] at h
+    obtain ⟨h, w⟩ := h
+    simp [h, w]
+  · let h' := h
+    simp only [Bool.not_eq_true] at h
+    simp only [h, Bool.false_eq, dite_eq_right_iff, decide_eq_false_iff_not, Classical.not_forall,
+      Bool.not_eq_true]
+    simpa [isEqv_iff_rel] using h'
+
+@[simp] theorem isEqv_toList [BEq α] (a b : Array α) : (a.toList.isEqv b.toList r) = (a.isEqv b r) := by
+  simp [isEqv_eq_decide, List.isEqv_eq_decide]
+
 theorem eq_of_isEqv [DecidableEq α] (a b : Array α) (h : Array.isEqv a b (fun x y => x = y)) : a = b := by
   have ⟨h, h'⟩ := rel_of_isEqv (fun x y => x = y) a b h
   exact ext _ _ h (fun i lt _ => by simpa using h' i lt)
@@ -56,4 +89,22 @@ instance [DecidableEq α] : DecidableEq (Array α) :=
     | true  => isTrue (eq_of_isEqv a b h)
     | false => isFalse fun h' => by subst h'; rw [isEqv_self] at h; contradiction
 
+theorem beq_eq_decide [BEq α] (a b : Array α) :
+    (a == b) = if h : a.size = b.size then
+      decide (∀ (i : Nat) (h' : i < a.size), a[i] == b[i]'(h ▸ h')) else false := by
+  simp [BEq.beq, isEqv_eq_decide]
+
+@[simp] theorem beq_toList [BEq α] (a b : Array α) : (a.toList == b.toList) = (a == b) := by
+  simp [beq_eq_decide, List.beq_eq_decide]
+
 end Array
+
+namespace List
+
+@[simp] theorem isEqv_toArray [BEq α] (a b : List α) : (a.toArray.isEqv b.toArray r) = (a.isEqv b r) := by
+  simp [isEqv_eq_decide, Array.isEqv_eq_decide]
+
+@[simp] theorem beq_toArray [BEq α] (a b : List α) : (a.toArray == b.toArray) = (a == b) := by
+  simp [beq_eq_decide, Array.beq_eq_decide]
+
+end List
diff --git a/src/Init/Data/Array/GetLit.lean b/src/Init/Data/Array/GetLit.lean
index c294c0b7e14b..73594b793788 100644
--- a/src/Init/Data/Array/GetLit.lean
+++ b/src/Init/Data/Array/GetLit.lean
@@ -41,6 +41,6 @@ where
   getLit_eq (as : Array α) (i : Nat) (h₁ : as.size = n) (h₂ : i < n) : as.getLit i h₁ h₂ = getElem as.toList i ((id (α := as.toList.length = n) h₁) ▸ h₂) :=
     rfl
   go (i : Nat) (hi : i ≤ as.size) : toListLitAux as n hsz i hi (as.toList.drop i) = as.toList := by
-    induction i <;> simp [getLit_eq, List.get_drop_eq_drop, toListLitAux, List.drop, *]
+    induction i <;> simp only [List.drop, toListLitAux, getLit_eq, List.get_drop_eq_drop, *]
 
 end Array
diff --git a/src/Init/Data/Array/Lemmas.lean b/src/Init/Data/Array/Lemmas.lean
index 37f5509f2d6c..0a60920f7c97 100644
--- a/src/Init/Data/Array/Lemmas.lean
+++ b/src/Init/Data/Array/Lemmas.lean
@@ -18,8 +18,6 @@ import Init.TacticsExtra
 
 namespace Array
 
-@[simp] theorem getElem_toList {a : Array α} {i : Nat} (h : i < a.size) : a.toList[i] = a[i] := rfl
-
 @[simp] theorem getElem_mk {xs : List α} {i : Nat} (h : i < xs.length) : (Array.mk xs)[i] = xs[i] := rfl
 
 theorem getElem_eq_getElem_toList {a : Array α} (h : i < a.size) : a[i] = a.toList[i] := by
@@ -86,16 +84,6 @@ We prefer to pull `List.toArray` outwards.
     (a.toArrayAux b).size = b.size + a.length := by
   simp [size]
 
-@[simp] theorem toArray_toList (a : Array α) : a.toList.toArray = a := rfl
-
-@[simp] theorem getElem_toArray {a : List α} {i : Nat} (h : i < a.toArray.size) :
-    a.toArray[i] = a[i]'(by simpa using h) := rfl
-
-@[simp] theorem getElem?_toArray {a : List α} {i : Nat} : a.toArray[i]? = a[i]? := rfl
-
-@[simp] theorem getElem!_toArray [Inhabited α] {a : List α} {i : Nat} :
-    a.toArray[i]! = a[i]! := rfl
-
 @[simp] theorem push_toArray (l : List α) (a : α) : l.toArray.push a = (l ++ [a]).toArray := by
   apply ext'
   simp
@@ -170,6 +158,9 @@ namespace Array
 
 @[simp] theorem singleton_def (v : α) : singleton v = #[v] := rfl
 
+-- This is a duplicate of `List.toArray_toList`.
+-- It's confusing to guess which namespace this theorem should live in,
+-- so we provide both.
 @[simp] theorem toArray_toList (a : Array α) : a.toList.toArray = a := rfl
 
 @[simp] theorem length_toList {l : Array α} : l.toList.length = l.size := rfl
diff --git a/src/Init/Data/List/Basic.lean b/src/Init/Data/List/Basic.lean
index e66e18962007..8e31bf9d2cd9 100644
--- a/src/Init/Data/List/Basic.lean
+++ b/src/Init/Data/List/Basic.lean
@@ -122,6 +122,11 @@ protected def beq [BEq α] : List α → List α → Bool
   | a::as, b::bs => a == b && List.beq as bs
   | _,     _     => false
 
+@[simp] theorem beq_nil_nil [BEq α] : List.beq ([] : List α) ([] : List α) = true := rfl
+@[simp] theorem beq_cons_nil [BEq α] (a : α) (as : List α) : List.beq (a::as) [] = false := rfl
+@[simp] theorem beq_nil_cons [BEq α] (a : α) (as : List α) : List.beq [] (a::as) = false := rfl
+theorem beq_cons₂ [BEq α] (a b : α) (as bs : List α) : List.beq (a::as) (b::bs) = (a == b && List.beq as bs) := rfl
+
 instance [BEq α] : BEq (List α) := ⟨List.beq⟩
 
 instance [BEq α] [LawfulBEq α] : LawfulBEq (List α) where
diff --git a/src/Init/Data/List/Nat.lean b/src/Init/Data/List/Nat.lean
index 60e33a1027b6..ae479a9fead3 100644
--- a/src/Init/Data/List/Nat.lean
+++ b/src/Init/Data/List/Nat.lean
@@ -12,3 +12,4 @@ import Init.Data.List.Nat.TakeDrop
 import Init.Data.List.Nat.Count
 import Init.Data.List.Nat.Erase
 import Init.Data.List.Nat.Find
+import Init.Data.List.Nat.BEq
diff --git a/src/Init/Data/List/Nat/BEq.lean b/src/Init/Data/List/Nat/BEq.lean
new file mode 100644
index 000000000000..a15b3bf935c2
--- /dev/null
+++ b/src/Init/Data/List/Nat/BEq.lean
@@ -0,0 +1,47 @@
+/-
+Copyright (c) 2024 Lean FRO All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+prelude
+import Init.Data.Nat.Lemmas
+import Init.Data.List.Basic
+
+namespace List
+
+/-! ### isEqv-/
+
+theorem isEqv_eq_decide (a b : List α) (r) :
+    isEqv a b r = if h : a.length = b.length then
+      decide (∀ (i : Nat) (h' : i < a.length), r (a[i]'(h ▸ h')) (b[i]'(h ▸ h'))) else false := by
+  induction a generalizing b with
+  | nil =>
+    cases b <;> simp
+  | cons a as ih =>
+    cases b with
+    | nil => simp
+    | cons b bs =>
+      simp only [isEqv, ih, length_cons, Nat.add_right_cancel_iff]
+      split <;> simp [Nat.forall_lt_succ_left']
+
+/-! ### beq -/
+
+theorem beq_eq_isEqv [BEq α] (a b : List α) : a.beq b = isEqv a b (· == ·) := by
+  induction a generalizing b with
+  | nil =>
+    cases b <;> simp
+  | cons a as ih =>
+    cases b with
+    | nil => simp
+    | cons b bs =>
+      simp only [beq_cons₂, ih, isEqv_eq_decide, length_cons, Nat.add_right_cancel_iff,
+        Nat.forall_lt_succ_left', getElem_cons_zero, getElem_cons_succ, Bool.decide_and,
+        Bool.decide_eq_true]
+      split <;> simp
+
+theorem beq_eq_decide [BEq α] (a b : List α) :
+    (a == b) = if h : a.length = b.length then
+      decide (∀ (i : Nat) (h' : i < a.length), a[i] == b[i]'(h ▸ h')) else false := by
+  simp [BEq.beq, beq_eq_isEqv, isEqv_eq_decide]
+
+end List
diff --git a/tests/lean/run/array_isEqvAux.lean b/tests/lean/run/array_isEqvAux.lean
index 029b6a4edadd..04c19852bcd7 100644
--- a/tests/lean/run/array_isEqvAux.lean
+++ b/tests/lean/run/array_isEqvAux.lean
@@ -17,6 +17,8 @@ After unfolding the instances 'instDecidableEqNat', 'Array.instDecidableEq' and
 
 example : #[0, 1] = #[0, 1] := by decide
 
+example : let a := Array.range (10^6); a == a := by native_decide
+
 /-!
 There are other `Array` functions that use well-founded recursion,
 which we've marked as `@[semireducible]`. We test that `decide` can unfold them here.