chore: preparation for Array.erase lemmas (#6836)

This PR rearranges some material, and adds some missing lemmas, in
preparation for aligning `List/Array/Vector.erase(P)`.

src/Init/Data/Array/Lemmas.lean

@@ -836,9 +836,6 @@ theorem mem_or_eq_of_mem_set
   cases as
   simpa using List.mem_or_eq_of_mem_set (by simpa using h)
 
-@[simp] theorem toList_set (a : Array α) (i x h) :
-    (a.set i x).toList = a.toList.set i x := rfl
-
 /-! ### setIfInBounds -/
 
 @[simp] theorem set!_eq_setIfInBounds : @set! = @setIfInBounds := rfl
@@ -2283,10 +2280,6 @@ theorem flatMap_mkArray {β} (f : α → Array β) : (mkArray n a).flatMap f = (
 
 /-! ### Preliminaries about `swap` needed for `reverse`. -/
 
-theorem swap_def (a : Array α) (i j : Nat) (hi hj) :
-    a.swap i j hi hj = (a.set i a[j]).set j a[i] (by simpa using hj) := by
-  simp [swap]
-
 theorem getElem?_swap (a : Array α) (i j : Nat) (hi hj) (k : Nat) : (a.swap i j hi hj)[k]? =
     if j = k then some a[i] else if i = k then some a[j] else a[k]? := by
   simp [swap_def, getElem?_set]
@@ -3303,9 +3296,6 @@ theorem get_set (a : Array α) (i : Nat) (hi : i < a.size) (j : Nat) (hj : j < a
     (h : i ≠ j) : (a.set i v)[j]'(by simp [*]) = a[j] := by
   simp only [set, ← getElem_toList, List.getElem_set_ne h]
 
-@[simp] theorem toList_swap (a : Array α) (i j : Nat) (hi hj) :
-    (a.swap i j hi hj).toList = (a.toList.set i a[j]).set j a[i] := by simp [swap_def]
-
 @[simp] theorem swapAt_def (a : Array α) (i : Nat) (v : α) (hi) :
     a.swapAt i v hi = (a[i], a.set i v) := rfl
 
@@ -3635,11 +3625,6 @@ theorem toListRev_toArray (l : List α) : l.toArray.toListRev = l.reverse := by
 theorem uset_toArray (l : List α) (i : USize) (a : α) (h : i.toNat < l.toArray.size) :
     l.toArray.uset i a h = (l.set i.toNat a).toArray := by simp
 
-@[simp] theorem swap_toArray (l : List α) (i j : Nat) {hi hj}:
-    l.toArray.swap i j hi hj = ((l.set i l[j]).set j l[i]).toArray := by
-  apply ext'
-  simp
-
 @[simp] theorem modify_toArray (f : α → α) (l : List α) :
     l.toArray.modify i f = (l.modify f i).toArray := by
   apply ext'
@@ -3657,31 +3642,6 @@ theorem uset_toArray (l : List α) (i : USize) (a : α) (h : i.toNat < l.toArray
 @[simp] theorem toArray_ofFn (f : Fin n → α) : (ofFn f).toArray = Array.ofFn f := by
   ext <;> simp
 
-@[simp] theorem eraseIdx_toArray (l : List α) (i : Nat) (h : i < l.toArray.size) :
-    l.toArray.eraseIdx i h = (l.eraseIdx i).toArray := by
-  rw [Array.eraseIdx]
-  split <;> rename_i h'
-  · rw [eraseIdx_toArray]
-    simp only [swap_toArray, Fin.getElem_fin, toList_toArray, mk.injEq]
-    rw [eraseIdx_set_gt (by simp), eraseIdx_set_eq]
-    simp
-  · simp at h h'
-    have t : i = l.length - 1 := by omega
-    simp [t]
-termination_by l.length - i
-decreasing_by
-  rename_i h
-  simp at h
-  simp
-  omega
-
-@[simp] theorem eraseIdxIfInBounds_toArray (l : List α) (i : Nat) :
-    l.toArray.eraseIdxIfInBounds i = (l.eraseIdx i).toArray := by
-  rw [Array.eraseIdxIfInBounds]
-  split
-  · simp
-  · simp_all [eraseIdx_eq_self.2]
-
 end List
 
 namespace Array

src/Init/Data/List/Erase.lean

@@ -271,6 +271,20 @@ theorem head_eraseP_mem (xs : List α) (p : α → Bool) (h) : (xs.eraseP p).hea
 theorem getLast_eraseP_mem (xs : List α) (p : α → Bool) (h) : (xs.eraseP p).getLast h ∈ xs :=
   (eraseP_sublist xs).getLast_mem h
 
+theorem eraseP_eq_eraseIdx {xs : List α} {p : α → Bool} :
+    xs.eraseP p = match xs.findIdx? p with
+    | none => xs
+    | some i => xs.eraseIdx i := by
+  induction xs with
+  | nil => rfl
+  | cons x xs ih =>
+    rw [eraseP_cons, findIdx?_cons]
+    by_cases h : p x
+    · simp [h]
+    · simp only [h]
+      rw [ih]
+      split <;> simp [*]
+
 /-! ### erase -/
 section erase
 variable [BEq α]
@@ -457,6 +471,19 @@ theorem head_erase_mem (xs : List α) (a : α) (h) : (xs.erase a).head h ∈ xs
 theorem getLast_erase_mem (xs : List α) (a : α) (h) : (xs.erase a).getLast h ∈ xs :=
   (erase_sublist a xs).getLast_mem h
 
+theorem erase_eq_eraseIdx [LawfulBEq α] (l : List α) (a : α) :
+    l.erase a = match l.indexOf? a with
+    | none => l
+    | some i => l.eraseIdx i := by
+  induction l with
+  | nil => simp
+  | cons x xs ih =>
+    rw [erase_cons, indexOf?_cons]
+    split
+    · simp
+    · simp [ih]
+      split <;> simp [*]
+
 end erase
 
 /-! ### eraseIdx -/
@@ -573,7 +600,8 @@ protected theorem IsPrefix.eraseIdx {l l' : List α} (h : l <+: l') (k : Nat) :
 -- See also `mem_eraseIdx_iff_getElem` and `mem_eraseIdx_iff_getElem?` in
 -- `Init/Data/List/Nat/Basic.lean`.
 
-theorem erase_eq_eraseIdx [BEq α] [LawfulBEq α] (l : List α) (a : α) (i : Nat) (w : l.indexOf a = i) :
+theorem erase_eq_eraseIdx_of_indexOf [BEq α] [LawfulBEq α]
+    (l : List α) (a : α) (i : Nat) (w : l.indexOf a = i) :
     l.erase a = l.eraseIdx i := by
   subst w
   rw [erase_eq_iff]

src/Init/Data/List/Find.lean

@@ -934,6 +934,12 @@ The verification API for `indexOf?` is still incomplete.
 The lemmas below should be made consistent with those for `findIdx?` (and proved using them).
 -/
 
+@[simp] theorem indexOf?_nil [BEq α] [LawfulBEq α] : ([] : List α).indexOf? a = none := rfl
+
+theorem indexOf?_cons [BEq α] [LawfulBEq α] (a : α) (xs : List α) (b : α) :
+    (a :: xs).indexOf? b = if a == b then some 0 else (xs.indexOf? b).map (· + 1) := by
+  simp [indexOf?]
+
 @[simp] theorem indexOf?_eq_none_iff [BEq α] [LawfulBEq α] {l : List α} {a : α} :
     l.indexOf? a = none ↔ a ∉ l := by
   simp only [indexOf?, findIdx?_eq_none_iff, beq_eq_false_iff_ne, ne_eq]

src/Init/Data/List/Nat/Erase.lean

@@ -65,6 +65,11 @@ theorem getElem_eraseIdx_of_ge (l : List α) (i : Nat) (j : Nat) (h : j < (l.era
   rw [getElem_eraseIdx, dif_neg]
   omega
 
+theorem eraseIdx_eq_dropLast (l : List α) (i : Nat) (h : i + 1 = l.length) :
+    l.eraseIdx i = l.dropLast := by
+  simp [eraseIdx_eq_take_drop_succ, h]
+  rw [take_eq_dropLast h]
+
 theorem eraseIdx_set_eq {l : List α} {i : Nat} {a : α} :
     (l.set i a).eraseIdx i = l.eraseIdx i := by
   apply ext_getElem

src/Init/Data/List/Nat/TakeDrop.lean

@@ -171,6 +171,20 @@ theorem dropLast_take {n : Nat} {l : List α} (h : n < l.length) :
 
 @[deprecated map_eq_append_iff (since := "2024-09-05")] abbrev map_eq_append_split := @map_eq_append_iff
 
+theorem take_eq_dropLast {l : List α} {i : Nat} (h : i + 1 = l.length) :
+    l.take i = l.dropLast := by
+  induction l generalizing i with
+  | nil => simp
+  | cons a as ih =>
+    cases i
+    · simp_all
+    · cases as with
+      | nil => simp_all
+      | cons b bs =>
+        simp only [take_succ_cons, dropLast_cons₂]
+        rw [ih]
+        simpa using h
+
 theorem take_prefix_take_left (l : List α) {m n : Nat} (h : m ≤ n) : take m l <+: take n l := by
   rw [isPrefix_iff]
   intro i w

src/Init/Data/List/ToArray.lean

@@ -13,6 +13,21 @@ import Init.Data.Array.Lex.Basic
 
 We prefer to pull `List.toArray` outwards past `Array` operations.
 -/
+
+namespace Array
+
+@[simp] theorem toList_set (a : Array α) (i x h) :
+    (a.set i x).toList = a.toList.set i x := rfl
+
+theorem swap_def (a : Array α) (i j : Nat) (hi hj) :
+    a.swap i j hi hj = (a.set i a[j]).set j a[i] (by simpa using hj) := by
+  simp [swap]
+
+@[simp] theorem toList_swap (a : Array α) (i j : Nat) (hi hj) :
+    (a.swap i j hi hj).toList = (a.toList.set i a[j]).set j a[i] := by simp [swap_def]
+
+end Array
+
 namespace List
 
 open Array
@@ -417,4 +432,34 @@ theorem flatMap_toArray_cons {β} (f : α → Array β) (a : α) (as : List α)
     apply ext'
     simp [ih, flatMap_toArray_cons]
 
+@[simp] theorem swap_toArray (l : List α) (i j : Nat) {hi hj}:
+    l.toArray.swap i j hi hj = ((l.set i l[j]).set j l[i]).toArray := by
+  apply ext'
+  simp
+
+@[simp] theorem eraseIdx_toArray (l : List α) (i : Nat) (h : i < l.toArray.size) :
+    l.toArray.eraseIdx i h = (l.eraseIdx i).toArray := by
+  rw [Array.eraseIdx]
+  split <;> rename_i h'
+  · rw [eraseIdx_toArray]
+    simp only [swap_toArray, Fin.getElem_fin, toList_toArray, mk.injEq]
+    rw [eraseIdx_set_gt (by simp), eraseIdx_set_eq]
+    simp
+  · simp at h h'
+    have t : i = l.length - 1 := by omega
+    simp [t]
+termination_by l.length - i
+decreasing_by
+  rename_i h
+  simp at h
+  simp
+  omega
+
+@[simp] theorem eraseIdxIfInBounds_toArray (l : List α) (i : Nat) :
+    l.toArray.eraseIdxIfInBounds i = (l.eraseIdx i).toArray := by
+  rw [Array.eraseIdxIfInBounds]
+  split
+  · simp
+  · simp_all [eraseIdx_eq_self.2]
+
 end List