From 126244f1709c977a3dc083631b01f013393ecd20 Mon Sep 17 00:00:00 2001
From: Kim Morrison <kim@tqft.net>
Date: Wed, 11 Sep 2024 12:24:04 +1000
Subject: [PATCH] feat: more List.findIdx theorems

---
 src/Init/Data/List/Find.lean         | 25 +++++++++++++-
 src/Init/Data/List/Nat/TakeDrop.lean | 50 ++++++++++++++++++++++------
 src/Init/Data/Option/Lemmas.lean     | 27 +++++++++++++++
 src/Init/PropLemmas.lean             |  3 ++
 4 files changed, 94 insertions(+), 11 deletions(-)

diff --git a/src/Init/Data/List/Find.lean b/src/Init/Data/List/Find.lean
index 85a337672665..f65c0e0b05b9 100644
--- a/src/Init/Data/List/Find.lean
+++ b/src/Init/Data/List/Find.lean
@@ -627,11 +627,24 @@ theorem IsPrefix.findIdx_eq_of_findIdx_lt_length {l₁ l₂ : List α} {p : α 
 @[simp] theorem findIdx?_cons :
     (x :: xs).findIdx? p i = if p x then some i else findIdx? p xs (i + 1) := rfl
 
-@[simp] theorem findIdx?_succ :
+theorem findIdx?_succ :
     (xs : List α).findIdx? p (i+1) = (xs.findIdx? p i).map fun i => i + 1 := by
   induction xs generalizing i with simp
   | cons _ _ _ => split <;> simp_all
 
+@[simp] theorem findIdx?_start_succ :
+    (xs : List α).findIdx? p (i+1) = (xs.findIdx? p 0).map fun k => k + (i + 1) := by
+  induction xs generalizing i with
+  | nil => simp
+  | cons _ _ _ =>
+    simp only [findIdx?_succ, findIdx?_cons, Nat.zero_add]
+    split
+    · simp_all
+    · simp_all only [findIdx?_succ, Bool.not_eq_true, Option.map_map, Nat.zero_add]
+      congr
+      ext
+      simp only [Nat.add_comm i, Function.comp_apply, Nat.add_assoc]
+
 @[simp]
 theorem findIdx?_eq_none_iff {xs : List α} {p : α → Bool} :
     xs.findIdx? p = none ↔ ∀ x, x ∈ xs → p x = false := by
@@ -683,6 +696,16 @@ theorem findIdx?_eq_none_iff_findIdx_eq {xs : List α} {p : α → Bool} :
     xs.findIdx? p = none ↔ xs.findIdx p = xs.length := by
   simp
 
+theorem findIdx?_eq_guard_findIdx_lt {xs : List α} {p : α → Bool} :
+    xs.findIdx? p = Option.guard (fun i => i < xs.length) (xs.findIdx p) := by
+  match h : xs.findIdx? p with
+  | none =>
+    simp only [findIdx?_eq_none_iff] at h
+    simp [findIdx_eq_length_of_false h, Option.guard]
+  | some i =>
+    simp only [findIdx?_eq_some_iff_findIdx_eq] at h
+    simp [h]
+
 theorem findIdx?_eq_some_iff_getElem {xs : List α} {p : α → Bool} {i : Nat} :
     xs.findIdx? p = some i ↔
       ∃ h : i < xs.length, p xs[i] ∧ ∀ j (hji : j < i), ¬p (xs[j]'(Nat.lt_trans hji h)) := by
diff --git a/src/Init/Data/List/Nat/TakeDrop.lean b/src/Init/Data/List/Nat/TakeDrop.lean
index 428f839f759f..838a24c1b26f 100644
--- a/src/Init/Data/List/Nat/TakeDrop.lean
+++ b/src/Init/Data/List/Nat/TakeDrop.lean
@@ -191,15 +191,7 @@ theorem dropLast_take {n : Nat} {l : List α} (h : n < l.length) :
     (l.take n).dropLast = l.take (n - 1) := by
   simp only [dropLast_eq_take, length_take, Nat.le_of_lt h, Nat.min_eq_left, take_take, sub_le]
 
-theorem map_eq_append_split {f : α → β} {l : List α} {s₁ s₂ : List β}
-    (h : map f l = s₁ ++ s₂) : ∃ l₁ l₂, l = l₁ ++ l₂ ∧ map f l₁ = s₁ ∧ map f l₂ = s₂ := by
-  have := h
-  rw [← take_append_drop (length s₁) l] at this ⊢
-  rw [map_append] at this
-  refine ⟨_, _, rfl, append_inj this ?_⟩
-  rw [length_map, length_take, Nat.min_eq_left]
-  rw [← length_map l f, h, length_append]
-  apply Nat.le_add_right
+@[deprecated map_eq_append_iff (since := "2024-09-05")] abbrev map_eq_append_split := @map_eq_append_iff
 
 theorem take_prefix_take_left (l : List α) {m n : Nat} (h : m ≤ n) : take m l <+: take n l := by
   rw [isPrefix_iff]
@@ -464,7 +456,7 @@ theorem false_of_mem_take_findIdx {xs : List α} {p : α → Bool} (h : x ∈ xs
   obtain ⟨i, h, rfl⟩ := h
   exact not_of_lt_findIdx (by omega)
 
-theorem findIdx_take {xs : List α} {n : Nat} {p : α → Bool} :
+@[simp] theorem findIdx_take {xs : List α} {n : Nat} {p : α → Bool} :
     (xs.take n).findIdx p = min n (xs.findIdx p) := by
   induction xs generalizing n with
   | nil => simp
@@ -476,6 +468,44 @@ theorem findIdx_take {xs : List α} {n : Nat} {p : α → Bool} :
       · simp
       · rw [Nat.add_min_add_right]
 
+@[simp] theorem findIdx?_take {xs : List α} {n : Nat} {p : α → Bool} :
+    (xs.take n).findIdx? p = (xs.findIdx? p).bind (Option.guard (fun i => i < n)) := by
+  induction xs generalizing n with
+  | nil => simp
+  | cons x xs ih =>
+    cases n
+    · simp
+    · simp only [take_succ_cons, findIdx?_cons]
+      split
+      · simp
+      · simp [ih, Option.guard_comp]
+
+@[simp] theorem min_findIdx_findIdx {xs : List α} {p q : α → Bool} :
+    min (xs.findIdx p) (xs.findIdx q) = xs.findIdx (fun a => p a || q a) := by
+  induction xs with
+  | nil => simp
+  | cons x xs ih =>
+    simp [findIdx_cons, cond_eq_if, Bool.not_eq_eq_eq_not, Bool.not_true]
+    split <;> split <;> simp_all [Nat.add_min_add_right]
+
+/-! ### takeWhile -/
+
+theorem takeWhile_eq_take_findIdx_not {xs : List α} {p : α → Bool} :
+    takeWhile p xs = take (xs.findIdx (fun a => !p a)) xs := by
+  induction xs with
+  | nil => simp
+  | cons x xs ih =>
+    simp only [takeWhile_cons, ih, findIdx_cons, cond_eq_if, Bool.not_eq_eq_eq_not, Bool.not_true]
+    split <;> simp_all
+
+theorem dropWhile_eq_drop_findIdx_not {xs : List α} {p : α → Bool} :
+    dropWhile p xs = drop (xs.findIdx (fun a => !p a)) xs := by
+  induction xs with
+  | nil => simp
+  | cons x xs ih =>
+    simp only [dropWhile_cons, ih, findIdx_cons, cond_eq_if, Bool.not_eq_eq_eq_not, Bool.not_true]
+    split <;> simp_all
+
 /-! ### rotateLeft -/
 
 @[simp] theorem rotateLeft_replicate (n) (a : α) : rotateLeft (replicate m a) n = replicate m a := by
diff --git a/src/Init/Data/Option/Lemmas.lean b/src/Init/Data/Option/Lemmas.lean
index 45b024761179..dd196be56af3 100644
--- a/src/Init/Data/Option/Lemmas.lean
+++ b/src/Init/Data/Option/Lemmas.lean
@@ -247,6 +247,12 @@ theorem isSome_filter_of_isSome (p : α → Bool) (o : Option α) (h : (o.filter
 theorem bind_map_comm {α β} {x : Option (Option α)} {f : α → β} :
     x.bind (Option.map f) = (x.map (Option.map f)).bind id := by cases x <;> simp
 
+@[simp] theorem bind_map {f : α → β} {g : β → Option γ} {x : Option α} :
+    (x.map f).bind g = x.bind (g ∘ f) := by cases x <;> simp
+
+@[simp] theorem map_bind {f : α → Option β} {g : β → γ} {x : Option α} :
+    (x.bind f).map g = x.bind (Option.map g ∘ f) := by cases x <;> simp
+
 theorem join_map_eq_map_join {f : α → β} {x : Option (Option α)} :
     (x.map (Option.map f)).join = x.join.map f := by cases x <;> simp
 
@@ -277,6 +283,27 @@ theorem map_orElse {x y : Option α} : (x <|> y).map f = (x.map f <|> y.map f) :
 @[simp] theorem guard_pos [DecidablePred p] (h : p a) : Option.guard p a = some a := by
   simp [Option.guard, h]
 
+@[congr] theorem guard_congr {f g : α → Prop} [DecidablePred f] [DecidablePred g]
+    (h : ∀ a, f a ↔ g a):
+    guard f = guard g := by
+  funext a
+  simp [guard, h]
+
+@[simp] theorem guard_false {α} :
+    guard (fun (_ : α) => False) = fun _ => none := by
+  funext a
+  simp [guard]
+
+@[simp] theorem guard_true {α} :
+    guard (fun (_ : α) => True) = some := by
+  funext a
+  simp [guard]
+
+theorem guard_comp {p : α → Prop} [DecidablePred p] {f : β → α} :
+    guard p ∘ f = Option.map f ∘ guard (p ∘ f) := by
+  ext1 b
+  simp [guard]
+
 theorem liftOrGet_eq_or_eq {f : α → α → α} (h : ∀ a b, f a b = a ∨ f a b = b) :
     ∀ o₁ o₂, liftOrGet f o₁ o₂ = o₁ ∨ liftOrGet f o₁ o₂ = o₂
   | none, none => .inl rfl
diff --git a/src/Init/PropLemmas.lean b/src/Init/PropLemmas.lean
index d754768fd893..aa8946f6a4e5 100644
--- a/src/Init/PropLemmas.lean
+++ b/src/Init/PropLemmas.lean
@@ -556,6 +556,9 @@ This is the same as `decidable_of_iff` but the iff is flipped. -/
 instance Decidable.predToBool (p : α → Prop) [DecidablePred p] :
     CoeDep (α → Prop) p (α → Bool) := ⟨fun b => decide <| p b⟩
 
+instance [DecidablePred p] : DecidablePred (p ∘ f) :=
+  fun x => inferInstanceAs (Decidable (p (f x)))
+
 /-- Prove that `a` is decidable by constructing a boolean `b` and a proof that `b ↔ a`.
 (This is sometimes taken as an alternate definition of decidability.) -/
 def decidable_of_bool : ∀ (b : Bool), (b ↔ a) → Decidable a