From 282e3a13e84b8ba1ea994e517e9e04cf7e7239bf Mon Sep 17 00:00:00 2001
From: Mario Carneiro <di.gama@gmail.com>
Date: Wed, 22 Nov 2023 15:05:08 -0500
Subject: [PATCH] feat: array lemmas for `Array.forIn`

---
 Std/Data/Array/Init/Lemmas.lean |  5 +++++
 Std/Data/Array/Lemmas.lean      | 16 ++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/Std/Data/Array/Init/Lemmas.lean b/Std/Data/Array/Init/Lemmas.lean
index 4ff0f51a2d..e7b4bb3d71 100644
--- a/Std/Data/Array/Init/Lemmas.lean
+++ b/Std/Data/Array/Init/Lemmas.lean
@@ -217,6 +217,11 @@ theorem size_mapM [Monad m] [LawfulMonad m] (f : α → m β) (as : Array α) :
   rw [← appendList_eq_append]; unfold Array.appendList
   induction l generalizing arr <;> simp [*]
 
+@[simp] theorem appendList_nil (arr : Array α) : arr ++ ([] : List α) = arr := Array.ext' (by simp)
+
+@[simp] theorem appendList_cons (arr : Array α) (a : α) (l : List α) :
+    arr ++ (a :: l) = arr.push a ++ l := Array.ext' (by simp)
+
 theorem foldl_data_eq_bind (l : List α) (acc : Array β)
     (F : Array β → α → Array β) (G : α → List β)
     (H : ∀ acc a, (F acc a).data = acc.data ++ G a) :
diff --git a/Std/Data/Array/Lemmas.lean b/Std/Data/Array/Lemmas.lean
index bed23cb9ba..c11190fd6f 100644
--- a/Std/Data/Array/Lemmas.lean
+++ b/Std/Data/Array/Lemmas.lean
@@ -320,3 +320,19 @@ termination_by _ => n - i
 @[simp] theorem getElem_ofFn (f : Fin n → α) (i : Nat) (h) :
     (ofFn f)[i] = f ⟨i, size_ofFn f ▸ h⟩ :=
   getElem_ofFn_go _ _ _ (by simp) (by simp) fun.
+
+theorem forIn_eq_data_forIn [Monad m]
+    (as : Array α) (b : β) (f : α → β → m (ForInStep β)) :
+    forIn as b f = forIn as.data b f := by
+  let rec loop : ∀ {i h b j}, j + i = as.size →
+      Array.forIn.loop as f i h b = forIn (as.data.drop j) b f
+    | 0, _, _, _, rfl => by rw [List.drop_length]; rfl
+    | i+1, _, _, j, ij => by
+      simp [forIn.loop]
+      have j_eq : j = size as - 1 - i := by simp [← ij, ← Nat.add_assoc]
+      have : as.size - 1 - i < as.size := j_eq ▸ ij ▸ Nat.lt_succ_of_le (Nat.le_add_right ..)
+      have : as[size as - 1 - i] :: as.data.drop (j + 1) = as.data.drop j := by
+        rw [j_eq]; exact List.get_cons_drop _ ⟨_, this⟩
+      simp [← this]; congr; funext x; congr; funext b
+      rw [loop (i := i)]; rw [← ij, Nat.succ_add]; rfl
+  simp [forIn, Array.forIn]; rw [loop (Nat.zero_add _)]; rfl