feat: relate Array.forIn and List.forIn (#5799)

src/Init/Data/Array/Lemmas.lean

@@ -102,6 +102,25 @@ We prefer to pull `List.toArray` outwards.
 @[simp] theorem back_toArray [Inhabited α] (l : List α) : l.toArray.back = l.getLast! := by
   simp only [back, size_toArray, Array.get!_eq_getElem!, getElem!_toArray, getLast!_eq_getElem!]
 
+@[simp] theorem forIn_loop_toArray [Monad m] (l : List α) (f : α → β → m (ForInStep β)) (i : Nat)
+    (h : i ≤ l.length) (b : β) :
+    Array.forIn.loop l.toArray f i h b = (l.drop (l.length - i)).forIn b f := by
+  induction i generalizing l b with
+  | zero => simp [Array.forIn.loop]
+  | succ i ih =>
+    simp only [Array.forIn.loop, size_toArray, getElem_toArray, ih, forIn_eq_forIn]
+    rw [Nat.sub_add_eq, List.drop_sub_one (by omega), List.getElem?_eq_getElem (by omega)]
+    simp only [Option.toList_some, singleton_append, forIn_cons]
+    have t : l.length - 1 - i = l.length - i - 1 := by omega
+    simp only [t]
+    congr
+
+@[simp] theorem forIn_toArray [Monad m] (l : List α) (b : β) (f : α → β → m (ForInStep β)) :
+    forIn l.toArray b f = forIn l b f := by
+  change l.toArray.forIn b f = l.forIn b f
+  rw [Array.forIn, forIn_loop_toArray]
+  simp
+
 theorem foldrM_toArray [Monad m] (f : α → β → m β) (init : β) (l : List α) :
     l.toArray.foldrM f init = l.foldrM f init := by
   rw [foldrM_eq_reverse_foldlM_toList]
@@ -699,6 +718,13 @@ theorem getElem_range {n : Nat} {x : Nat} (h : x < (Array.range n).size) : (Arra
 @[simp] theorem toList_take (a : Array α) (n : Nat) : (a.take n).toList = a.toList.take n := by
   apply List.ext_getElem <;> simp
 
+/-! ### forIn -/
+
+@[simp] theorem forIn_toList [Monad m] (as : Array α) (b : β) (f : α → β → m (ForInStep β)) :
+    forIn as.toList b f = forIn as b f := by
+  cases as
+  simp
+
 /-! ### foldl / foldr -/
 
 @[simp] theorem foldlM_loop_empty [Monad m] (f : β → α → m β) (init : β) (i j : Nat) :

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

@@ -187,6 +187,9 @@ theorem take_add (l : List α) (m n : Nat) : l.take (m + n) = l.take m ++ (l.dro
   · apply length_take_le
   · apply Nat.le_add_right
 
+theorem take_one {l : List α} : l.take 1 = l.head?.toList := by
+  induction l <;> simp
+
 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]
@@ -282,14 +285,14 @@ theorem mem_drop_iff_getElem {l : List α} {a : α} :
   · rintro ⟨i, hm, rfl⟩
     refine ⟨i, by simp; omega, by rw [getElem_drop]⟩
 
-theorem head?_drop (l : List α) (n : Nat) :
+@[simp] theorem head?_drop (l : List α) (n : Nat) :
     (l.drop n).head? = l[n]? := by
   rw [head?_eq_getElem?, getElem?_drop, Nat.add_zero]
 
-theorem head_drop {l : List α} {n : Nat} (h : l.drop n ≠ []) :
+@[simp] theorem head_drop {l : List α} {n : Nat} (h : l.drop n ≠ []) :
     (l.drop n).head h = l[n]'(by simp_all) := by
   have w : n < l.length := length_lt_of_drop_ne_nil h
-  simpa [getElem?_eq_getElem, h, w, head_eq_iff_head?_eq_some] using head?_drop l n
+  simp [getElem?_eq_getElem, h, w, head_eq_iff_head?_eq_some]
 
 theorem getLast?_drop {l : List α} : (l.drop n).getLast? = if l.length ≤ n then none else l.getLast? := by
   rw [getLast?_eq_getElem?, getElem?_drop]
@@ -300,7 +303,7 @@ theorem getLast?_drop {l : List α} : (l.drop n).getLast? = if l.length ≤ n th
     congr
     omega
 
-theorem getLast_drop {l : List α} (h : l.drop n ≠ []) :
+@[simp] theorem getLast_drop {l : List α} (h : l.drop n ≠ []) :
     (l.drop n).getLast h = l.getLast (ne_nil_of_length_pos (by simp at h; omega)) := by
   simp only [ne_eq, drop_eq_nil_iff] at h
   apply Option.some_inj.1
@@ -449,6 +452,26 @@ theorem reverse_drop {l : List α} {n : Nat} :
     rw [w, take_zero, drop_of_length_le, reverse_nil]
     omega
 
+theorem take_add_one {l : List α} {n : Nat} :
+    l.take (n + 1) = l.take n ++ l[n]?.toList := by
+  simp [take_add, take_one]
+
+theorem drop_eq_getElem?_toList_append {l : List α} {n : Nat} :
+    l.drop n = l[n]?.toList ++ l.drop (n + 1) := by
+  induction l generalizing n with
+  | nil => simp
+  | cons hd tl ih =>
+    cases n
+    · simp
+    · simp only [drop_succ_cons, getElem?_cons_succ]
+      rw [ih]
+
+theorem drop_sub_one {l : List α} {n : Nat} (h : 0 < n) :
+    l.drop (n - 1) = l[n - 1]?.toList ++ l.drop n := by
+  rw [drop_eq_getElem?_toList_append]
+  congr
+  omega
+
 /-! ### findIdx -/
 
 theorem false_of_mem_take_findIdx {xs : List α} {p : α → Bool} (h : x ∈ xs.take (xs.findIdx p)) :

tests/lean/run/array_simp.lean

@@ -5,3 +5,11 @@
 #check_simp #[1,2,3,4,5][2]! ~> 3
 #check_simp #[1,2,3,4,5][7]! ~> (default : Nat)
 #check_simp (#[] : Array Nat)[0]! ~> (default : Nat)
+
+attribute [local simp] Id.run in
+#check_simp
+  (Id.run do
+    let mut s := 0
+    for i in [1,2,3,4].toArray do
+      s := s + i
+    pure s) ~> 10