feat: relate Array.takeWhile with List.takeWhile (#5950)

src/Init/Data/Array/Bootstrap.lean

@@ -23,7 +23,7 @@ theorem foldlM_eq_foldlM_toList.aux [Monad m]
   · cases Nat.not_le_of_gt ‹_› (Nat.zero_add _ ▸ H)
   · rename_i i; rw [Nat.succ_add] at H
     simp [foldlM_eq_foldlM_toList.aux f arr i (j+1) H]
-    rw (occs := .pos [2]) [← List.get_drop_eq_drop _ _ ‹_›]
+    rw (occs := .pos [2]) [← List.getElem_cons_drop_succ_eq_drop ‹_›]
     rfl
   · rw [List.drop_of_length_le (Nat.ge_of_not_lt ‹_›)]; rfl
 

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 only [List.drop, toListLitAux, getLit_eq, List.get_drop_eq_drop, *]
+    induction i <;> simp only [List.drop, toListLitAux, getLit_eq, List.getElem_cons_drop_succ_eq_drop, *]
 
 end Array

src/Init/Data/Array/Lemmas.lean

@@ -246,7 +246,7 @@ where
   aux (i r) :
       mapM.map f arr i r = (arr.toList.drop i).foldlM (fun bs a => bs.push <$> f a) r := by
     unfold mapM.map; split
-    · rw [← List.get_drop_eq_drop _ i ‹_›]
+    · rw [← List.getElem_cons_drop_succ_eq_drop ‹_›]
       simp only [aux (i + 1), map_eq_pure_bind, length_toList, List.foldlM_cons, bind_assoc,
         pure_bind]
       rfl
@@ -1619,6 +1619,38 @@ theorem filterMap_toArray (f : α → Option β) (l : List α) :
 @[simp] theorem toArray_ofFn (f : Fin n → α) : (ofFn f).toArray = Array.ofFn f := by
   ext <;> simp
 
+theorem takeWhile_go_succ (p : α → Bool) (a : α) (l : List α) (i : Nat) :
+    takeWhile.go p (a :: l).toArray (i+1) r = takeWhile.go p l.toArray i r := by
+  rw [takeWhile.go, takeWhile.go]
+  simp only [size_toArray, length_cons, Nat.add_lt_add_iff_right, Array.get_eq_getElem,
+    getElem_toArray, getElem_cons_succ]
+  split
+  rw [takeWhile_go_succ]
+  rfl
+
+theorem takeWhile_go_toArray (p : α → Bool) (l : List α) (i : Nat) :
+    Array.takeWhile.go p l.toArray i r = r ++ (takeWhile p (l.drop i)).toArray := by
+  induction l generalizing i r with
+  | nil => simp [takeWhile.go]
+  | cons a l ih =>
+    rw [takeWhile.go]
+    cases i with
+    | zero =>
+      simp [takeWhile_go_succ, ih, takeWhile_cons]
+      split <;> simp
+    | succ i =>
+      simp only [size_toArray, length_cons, Nat.add_lt_add_iff_right, Array.get_eq_getElem,
+        getElem_toArray, getElem_cons_succ, drop_succ_cons]
+      split <;> rename_i h₁
+      · rw [takeWhile_go_succ, ih]
+        rw [← getElem_cons_drop_succ_eq_drop h₁, takeWhile_cons]
+        split <;> simp_all
+      · simp_all [drop_eq_nil_of_le]
+
+@[simp] theorem takeWhile_toArray (p : α → Bool) (l : List α) :
+    l.toArray.takeWhile p = (l.takeWhile p).toArray := by
+  simp [Array.takeWhile, takeWhile_go_toArray]
+
 end List
 
 namespace Array
@@ -1629,6 +1661,10 @@ namespace Array
 @[simp] theorem toList_ofFn (f : Fin n → α) : (Array.ofFn f).toList = List.ofFn f := by
   apply List.ext_getElem <;> simp
 
+@[simp] theorem toList_takeWhile (p : α → Bool) (as : Array α) :
+    (as.takeWhile p).toList = as.toList.takeWhile p := by
+  induction as; simp
+
 end Array
 
 /-! ### Deprecations -/

src/Init/Data/List/TakeDrop.lean

@@ -190,7 +190,7 @@ theorem set_drop {l : List α} {n m : Nat} {a : α} :
 theorem take_concat_get (l : List α) (i : Nat) (h : i < l.length) :
     (l.take i).concat l[i] = l.take (i+1) :=
   Eq.symm <| (append_left_inj _).1 <| (take_append_drop (i+1) l).trans <| by
-    rw [concat_eq_append, append_assoc, singleton_append, get_drop_eq_drop, take_append_drop]
+    rw [concat_eq_append, append_assoc, singleton_append, getElem_cons_drop_succ_eq_drop, take_append_drop]
 
 @[deprecated take_succ_cons (since := "2024-07-25")]
 theorem take_cons_succ : (a::as).take (i+1) = a :: as.take i := rfl

src/Init/GetElem.lean

@@ -211,10 +211,13 @@ instance : GetElem (List α) Nat α fun as i => i < as.length where
   | _ :: _, 0, _ => .head ..
   | _ :: l, _+1, _ => .tail _ (getElem_mem (l := l) ..)
 
-theorem get_drop_eq_drop (as : List α) (i : Nat) (h : i < as.length) : as[i] :: as.drop (i+1) = as.drop i :=
+theorem getElem_cons_drop_succ_eq_drop {as : List α} {i : Nat} (h : i < as.length) :
+    as[i] :: as.drop (i+1) = as.drop i :=
   match as, i with
   | _::_, 0   => rfl
-  | _::_, i+1 => get_drop_eq_drop _ i _
+  | _::_, i+1 => getElem_cons_drop_succ_eq_drop (i := i) _
+
+@[deprecated (since := "2024-11-05")] abbrev get_drop_eq_drop := @getElem_cons_drop_succ_eq_drop
 
 end List