feat: lemmas for Array.findSome? and find? (#6111)

This PR fills in the API for `Array.findSome?` and `Array.find?`,
transferring proofs from the corresponding List statements.

src/Init/Data/Array/Find.lean

@@ -0,0 +1,275 @@
+/-
+Copyright (c) 2024 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+prelude
+import Init.Data.List.Find
+import Init.Data.Array.Lemmas
+import Init.Data.Array.Attach
+
+/-!
+# Lemmas about `Array.findSome?`, `Array.find?`.
+-/
+
+namespace Array
+
+open Nat
+
+/-! ### findSome? -/
+
+@[simp] theorem findSomeRev?_push_of_isSome (l : Array α) (h : (f a).isSome) : (l.push a).findSomeRev? f = f a := by
+  cases l; simp_all
+
+@[simp] theorem findSomeRev?_push_of_isNone (l : Array α) (h : (f a).isNone) : (l.push a).findSomeRev? f = l.findSomeRev? f := by
+  cases l; simp_all
+
+theorem exists_of_findSome?_eq_some {f : α → Option β} {l : Array α} (w : l.findSome? f = some b) :
+    ∃ a, a ∈ l ∧ f a = b := by
+  cases l; simp_all [List.exists_of_findSome?_eq_some]
+
+@[simp] theorem findSome?_eq_none_iff : findSome? p l = none ↔ ∀ x ∈ l, p x = none := by
+  cases l; simp
+
+@[simp] theorem findSome?_isSome_iff {f : α → Option β} {l : Array α} :
+    (l.findSome? f).isSome ↔ ∃ x, x ∈ l ∧ (f x).isSome := by
+  cases l; simp
+
+theorem findSome?_eq_some_iff {f : α → Option β} {l : Array α} {b : β} :
+    l.findSome? f = some b ↔ ∃ (l₁ : Array α) (a : α) (l₂ : Array α), l = l₁.push a ++ l₂ ∧ f a = some b ∧ ∀ x ∈ l₁, f x = none := by
+  cases l
+  simp only [List.findSome?_toArray, List.findSome?_eq_some_iff]
+  constructor
+  · rintro ⟨l₁, a, l₂, rfl, h₁, h₂⟩
+    exact ⟨l₁.toArray, a, l₂.toArray, by simp_all⟩
+  · rintro ⟨l₁, a, l₂, h₀, h₁, h₂⟩
+    exact ⟨l₁.toList, a, l₂.toList, by simpa using congrArg toList h₀, h₁, by simpa⟩
+
+@[simp] theorem findSome?_guard (l : Array α) : findSome? (Option.guard fun x => p x) l = find? p l := by
+  cases l; simp
+
+@[simp] theorem getElem?_zero_filterMap (f : α → Option β) (l : Array α) : (l.filterMap f)[0]? = l.findSome? f := by
+  cases l; simp [← List.head?_eq_getElem?]
+
+@[simp] theorem getElem_zero_filterMap (f : α → Option β) (l : Array α) (h) :
+    (l.filterMap f)[0] = (l.findSome? f).get (by cases l; simpa [List.length_filterMap_eq_countP] using h) := by
+  cases l; simp [← List.head_eq_getElem, ← getElem?_zero_filterMap]
+
+@[simp] theorem back?_filterMap (f : α → Option β) (l : Array α) : (l.filterMap f).back? = l.findSomeRev? f := by
+  cases l; simp
+
+@[simp] theorem back!_filterMap [Inhabited β] (f : α → Option β) (l : Array α) :
+    (l.filterMap f).back! = (l.findSomeRev? f).getD default := by
+  cases l; simp
+
+@[simp] theorem map_findSome? (f : α → Option β) (g : β → γ) (l : Array α) :
+    (l.findSome? f).map g = l.findSome? (Option.map g ∘ f) := by
+  cases l; simp
+
+theorem findSome?_map (f : β → γ) (l : Array β) : findSome? p (l.map f) = l.findSome? (p ∘ f) := by
+  cases l; simp [List.findSome?_map]
+
+theorem findSome?_append {l₁ l₂ : Array α} : (l₁ ++ l₂).findSome? f = (l₁.findSome? f).or (l₂.findSome? f) := by
+  cases l₁; cases l₂; simp [List.findSome?_append]
+
+theorem getElem?_zero_flatten (L : Array (Array α)) :
+    (flatten L)[0]? = L.findSome? fun l => l[0]? := by
+  cases L using array_array_induction
+  simp [← List.head?_eq_getElem?, List.head?_flatten, List.findSome?_map, Function.comp_def]
+
+theorem getElem_zero_flatten.proof {L : Array (Array α)} (h : 0 < L.flatten.size) :
+    (L.findSome? fun l => l[0]?).isSome := by
+  cases L using array_array_induction
+  simp only [List.findSome?_toArray, List.findSome?_map, Function.comp_def, List.getElem?_toArray,
+    List.findSome?_isSome_iff, List.isSome_getElem?]
+  simp only [flatten_toArray_map_toArray, size_toArray, List.length_flatten,
+    Nat.sum_pos_iff_exists_pos, List.mem_map] at h
+  obtain ⟨_, ⟨xs, m, rfl⟩, h⟩ := h
+  exact ⟨xs, m, by simpa using h⟩
+
+theorem getElem_zero_flatten {L : Array (Array α)} (h) :
+    (flatten L)[0] = (L.findSome? fun l => l[0]?).get (getElem_zero_flatten.proof h) := by
+  have t := getElem?_zero_flatten L
+  simp [getElem?_eq_getElem, h] at t
+  simp [← t]
+
+theorem back?_flatten {L : Array (Array α)} :
+    (flatten L).back? = (L.findSomeRev? fun l => l.back?) := by
+  cases L using array_array_induction
+  simp [List.getLast?_flatten, ← List.map_reverse, List.findSome?_map, Function.comp_def]
+
+theorem findSome?_mkArray : findSome? f (mkArray n a) = if n = 0 then none else f a := by
+  simp [mkArray_eq_toArray_replicate, List.findSome?_replicate]
+
+@[simp] theorem findSome?_mkArray_of_pos (h : 0 < n) : findSome? f (mkArray n a) = f a := by
+  simp [findSome?_mkArray, Nat.ne_of_gt h]
+
+-- Argument is unused, but used to decide whether `simp` should unfold.
+@[simp] theorem findSome?_mkArray_of_isSome (_ : (f a).isSome) :
+   findSome? f (mkArray n a) = if n = 0 then none else f a := by
+  simp [findSome?_mkArray]
+
+@[simp] theorem findSome?_mkArray_of_isNone (h : (f a).isNone) :
+    findSome? f (mkArray n a) = none := by
+  rw [Option.isNone_iff_eq_none] at h
+  simp [findSome?_mkArray, h]
+
+/-! ### find? -/
+
+@[simp] theorem find?_singleton (a : α) (p : α → Bool) :
+    #[a].find? p = if p a then some a else none := by
+  simp [singleton_eq_toArray_singleton]
+
+@[simp] theorem findRev?_push_of_pos (l : Array α) (h : p a) :
+    findRev? p (l.push a) = some a := by
+  cases l; simp [h]
+
+@[simp] theorem findRev?_cons_of_neg (l : Array α) (h : ¬p a) :
+    findRev? p (l.push a) = findRev? p l := by
+  cases l; simp [h]
+
+@[simp] theorem find?_eq_none : find? p l = none ↔ ∀ x ∈ l, ¬ p x := by
+  cases l; simp
+
+theorem find?_eq_some_iff_append {xs : Array α} :
+    xs.find? p = some b ↔ p b ∧ ∃ (as bs : Array α), xs = as.push b ++ bs ∧ ∀ a ∈ as, !p a := by
+  rcases xs with ⟨xs⟩
+  simp only [List.find?_toArray, List.find?_eq_some_iff_append, Bool.not_eq_eq_eq_not,
+    Bool.not_true, exists_and_right, and_congr_right_iff]
+  intro w
+  constructor
+  · rintro ⟨as, ⟨⟨x, rfl⟩, h⟩⟩
+    exact ⟨as.toArray, ⟨x.toArray, by simp⟩ , by simpa using h⟩
+  · rintro ⟨as, ⟨⟨x, h'⟩, h⟩⟩
+    exact ⟨as.toList, ⟨x.toList, by simpa using congrArg Array.toList h'⟩,
+      by simpa using h⟩
+
+@[simp]
+theorem find?_push_eq_some {xs : Array α} :
+    (xs.push a).find? p = some b ↔ xs.find? p = some b ∨ (xs.find? p = none ∧ (p a ∧ a = b)) := by
+  cases xs; simp
+
+@[simp] theorem find?_isSome {xs : Array α} {p : α → Bool} : (xs.find? p).isSome ↔ ∃ x, x ∈ xs ∧ p x := by
+  cases xs; simp
+
+theorem find?_some {xs : Array α} (h : find? p xs = some a) : p a := by
+  cases xs
+  simp at h
+  exact List.find?_some h
+
+theorem mem_of_find?_eq_some {xs : Array α} (h : find? p xs = some a) : a ∈ xs := by
+  cases xs
+  simp at h
+  simpa using List.mem_of_find?_eq_some h
+
+theorem get_find?_mem {xs : Array α} (h) : (xs.find? p).get h ∈ xs := by
+  cases xs
+  simp [List.get_find?_mem]
+
+@[simp] theorem find?_filter {xs : Array α} (p q : α → Bool) :
+    (xs.filter p).find? q = xs.find? (fun a => p a ∧ q a) := by
+  cases xs; simp
+
+@[simp] theorem getElem?_zero_filter (p : α → Bool) (l : Array α) :
+    (l.filter p)[0]? = l.find? p := by
+  cases l; simp [← List.head?_eq_getElem?]
+
+@[simp] theorem getElem_zero_filter (p : α → Bool) (l : Array α) (h) :
+    (l.filter p)[0] =
+      (l.find? p).get (by cases l; simpa [← List.countP_eq_length_filter] using h) := by
+  cases l
+  simp [List.getElem_zero_eq_head]
+
+@[simp] theorem back?_filter (p : α → Bool) (l : Array α) : (l.filter p).back? = l.findRev? p := by
+  cases l; simp
+
+@[simp] theorem back!_filter [Inhabited α] (p : α → Bool) (l : Array α) :
+    (l.filter p).back! = (l.findRev? p).get! := by
+  cases l; simp [Option.get!_eq_getD]
+
+@[simp] theorem find?_filterMap (xs : Array α) (f : α → Option β) (p : β → Bool) :
+    (xs.filterMap f).find? p = (xs.find? (fun a => (f a).any p)).bind f := by
+  cases xs; simp
+
+@[simp] theorem find?_map (f : β → α) (xs : Array β) :
+    find? p (xs.map f) = (xs.find? (p ∘ f)).map f := by
+  cases xs; simp
+
+@[simp] theorem find?_append {l₁ l₂ : Array α} :
+    (l₁ ++ l₂).find? p = (l₁.find? p).or (l₂.find? p) := by
+  cases l₁
+  cases l₂
+  simp
+
+@[simp] theorem find?_flatten (xs : Array (Array α)) (p : α → Bool) :
+    xs.flatten.find? p = xs.findSome? (·.find? p) := by
+  cases xs using array_array_induction
+  simp [List.findSome?_map, Function.comp_def]
+
+theorem find?_flatten_eq_none {xs : Array (Array α)} {p : α → Bool} :
+    xs.flatten.find? p = none ↔ ∀ ys ∈ xs, ∀ x ∈ ys, !p x := by
+  simp
+
+/--
+If `find? p` returns `some a` from `xs.flatten`, then `p a` holds, and
+some array in `xs` contains `a`, and no earlier element of that array satisfies `p`.
+Moreover, no earlier array in `xs` has an element satisfying `p`.
+-/
+theorem find?_flatten_eq_some {xs : Array (Array α)} {p : α → Bool} {a : α} :
+    xs.flatten.find? p = some a ↔
+      p a ∧ ∃ (as : Array (Array α)) (ys zs : Array α) (bs : Array (Array α)),
+        xs = as.push (ys.push a ++ zs) ++ bs ∧
+        (∀ a ∈ as, ∀ x ∈ a, !p x) ∧ (∀ x ∈ ys, !p x) := by
+  cases xs using array_array_induction
+  simp only [flatten_toArray_map_toArray, List.find?_toArray, List.find?_flatten_eq_some]
+  simp only [Bool.not_eq_eq_eq_not, Bool.not_true, exists_and_right, and_congr_right_iff]
+  intro w
+  constructor
+  · rintro ⟨as, ys, ⟨⟨zs, bs, rfl⟩, h₁, h₂⟩⟩
+    exact ⟨as.toArray.map List.toArray, ys.toArray,
+      ⟨zs.toArray, bs.toArray.map List.toArray, by simp⟩, by simpa using h₁, by simpa using h₂⟩
+  · rintro ⟨as, ys, ⟨⟨zs, bs, h⟩, h₁, h₂⟩⟩
+    replace h := congrArg (·.map Array.toList) (congrArg Array.toList h)
+    simp [Function.comp_def] at h
+    exact ⟨as.toList.map Array.toList, ys.toList,
+      ⟨zs.toList, bs.toList.map Array.toList, by simpa using h⟩,
+        by simpa using h₁, by simpa using h₂⟩
+
+@[simp] theorem find?_flatMap (xs : Array α) (f : α → Array β) (p : β → Bool) :
+    (xs.flatMap f).find? p = xs.findSome? (fun x => (f x).find? p) := by
+  cases xs
+  simp [List.find?_flatMap, Array.flatMap_toArray]
+
+theorem find?_flatMap_eq_none {xs : Array α} {f : α → Array β} {p : β → Bool} :
+    (xs.flatMap f).find? p = none ↔ ∀ x ∈ xs, ∀ y ∈ f x, !p y := by
+  simp
+
+theorem find?_mkArray :
+    find? p (mkArray n a) = if n = 0 then none else if p a then some a else none := by
+  simp [mkArray_eq_toArray_replicate, List.find?_replicate]
+
+@[simp] theorem find?_mkArray_of_length_pos (h : 0 < n) :
+    find? p (mkArray n a) = if p a then some a else none := by
+  simp [find?_mkArray, Nat.ne_of_gt h]
+
+@[simp] theorem find?_mkArray_of_pos (h : p a) :
+    find? p (mkArray n a) = if n = 0 then none else some a := by
+  simp [find?_mkArray, h]
+
+@[simp] theorem find?_mkArray_of_neg (h : ¬ p a) : find? p (mkArray n a) = none := by
+  simp [find?_mkArray, h]
+
+-- This isn't a `@[simp]` lemma since there is already a lemma for `l.find? p = none` for any `l`.
+theorem find?_mkArray_eq_none {n : Nat} {a : α} {p : α → Bool} :
+    (mkArray n a).find? p = none ↔ n = 0 ∨ !p a := by
+  simp [mkArray_eq_toArray_replicate, List.find?_replicate_eq_none, Classical.or_iff_not_imp_left]
+
+@[simp] theorem find?_mkArray_eq_some {n : Nat} {a b : α} {p : α → Bool} :
+    (mkArray n a).find? p = some b ↔ n ≠ 0 ∧ p a ∧ a = b := by
+  simp [mkArray_eq_toArray_replicate]
+
+@[simp] theorem get_find?_mkArray (n : Nat) (a : α) (p : α → Bool) (h) :
+    ((mkArray n a).find? p).get h = a := by
+  simp [mkArray_eq_toArray_replicate]
+
+end Array