From 15139b6ef65d51d480b455af21c44d46feb004f0 Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Wed, 6 Nov 2024 20:22:08 +1100
Subject: [PATCH] feat: relate Array.zipWith/zip/unzip with List versions
 (#5972)

---
 src/Init/Data/Array/Bootstrap.lean |  11 +++
 src/Init/Data/Array/Lemmas.lean    | 111 ++++++++++++++++++++++++-----
 src/Init/Data/List/Lemmas.lean     |  10 ++-
 3 files changed, 114 insertions(+), 18 deletions(-)

diff --git a/src/Init/Data/Array/Bootstrap.lean b/src/Init/Data/Array/Bootstrap.lean
index bba5c37d74..09686e2b09 100644
--- a/src/Init/Data/Array/Bootstrap.lean
+++ b/src/Init/Data/Array/Bootstrap.lean
@@ -79,6 +79,17 @@ theorem foldr_eq_foldr_toList (f : α → β → β) (init : β) (arr : Array α
   rw [foldl_eq_foldl_toList]
   induction arr'.toList generalizing arr <;> simp [*]
 
+@[simp] theorem toList_empty : (#[] : Array α).toList = [] := rfl
+
+@[simp] theorem append_nil (as : Array α) : as ++ #[] = as := by
+  apply ext'; simp only [toList_append, toList_empty, List.append_nil]
+
+@[simp] theorem nil_append (as : Array α) : #[] ++ as = as := by
+  apply ext'; simp only [toList_append, toList_empty, List.nil_append]
+
+@[simp] theorem append_assoc (as bs cs : Array α) : as ++ bs ++ cs = as ++ (bs ++ cs) := by
+  apply ext'; simp only [toList_append, List.append_assoc]
+
 @[simp] theorem appendList_eq_append
     (arr : Array α) (l : List α) : arr.appendList l = arr ++ l := rfl
 
diff --git a/src/Init/Data/Array/Lemmas.lean b/src/Init/Data/Array/Lemmas.lean
index eb843a835a..e4ec396ffa 100644
--- a/src/Init/Data/Array/Lemmas.lean
+++ b/src/Init/Data/Array/Lemmas.lean
@@ -190,6 +190,13 @@ theorem foldl_toArray (f : β → α → β) (init : β) (l : List α) :
   apply ext'
   simp
 
+@[simp] theorem push_append_toArray {as : Array α} {a : α} {bs : List α} : as.push a ++ bs.toArray = as ++ (a ::bs).toArray := by
+  cases as
+  simp
+
+@[simp] theorem foldl_push {l : List α} {as : Array α} : l.foldl Array.push as = as ++ l.toArray := by
+  induction l generalizing as <;> simp [*]
+
 theorem isPrefixOfAux_toArray_succ [BEq α] (l₁ l₂ : List α) (hle : l₁.length ≤ l₂.length) (i : Nat) :
     Array.isPrefixOfAux l₁.toArray l₂.toArray hle (i + 1) =
       Array.isPrefixOfAux l₁.tail.toArray l₂.tail.toArray (by simp; omega) i := by
@@ -238,6 +245,44 @@ theorem isPrefixOfAux_toArray_zero [BEq α] (l₁ l₂ : List α) (hle : l₁.le
         rw [ih]
         simp_all
 
+theorem zipWithAux_toArray_succ (f : α → β → γ) (as : List α) (bs : List β) (i : Nat) (cs : Array γ) :
+    zipWithAux f as.toArray bs.toArray (i + 1) cs = zipWithAux f as.tail.toArray bs.tail.toArray i cs := by
+  rw [zipWithAux]
+  conv => rhs; rw [zipWithAux]
+  simp only [size_toArray, getElem_toArray, length_tail, getElem_tail]
+  split <;> rename_i h₁
+  · split <;> rename_i h₂
+    · rw [dif_pos (by omega), dif_pos (by omega), zipWithAux_toArray_succ]
+    · rw [dif_pos (by omega)]
+      rw [dif_neg (by omega)]
+  · rw [dif_neg (by omega)]
+
+theorem zipWithAux_toArray_succ' (f : α → β → γ) (as : List α) (bs : List β) (i : Nat) (cs : Array γ) :
+    zipWithAux f as.toArray bs.toArray (i + 1) cs = zipWithAux f (as.drop (i+1)).toArray (bs.drop (i+1)).toArray 0 cs := by
+  induction i generalizing as bs cs with
+  | zero => simp [zipWithAux_toArray_succ]
+  | succ i ih =>
+    rw [zipWithAux_toArray_succ, ih]
+    simp
+
+theorem zipWithAux_toArray_zero (f : α → β → γ) (as : List α) (bs : List β) (cs : Array γ) :
+    zipWithAux f as.toArray bs.toArray 0 cs = cs ++ (List.zipWith f as bs).toArray := by
+  rw [Array.zipWithAux]
+  match as, bs with
+  | [], _ => simp
+  | _, [] => simp
+  | a :: as, b :: bs =>
+    simp [zipWith_cons_cons, zipWithAux_toArray_succ', zipWithAux_toArray_zero, push_append_toArray]
+
+@[simp] theorem zipWith_toArray (f : α → β → γ) (as : List α) (bs : List β) :
+    Array.zipWith as.toArray bs.toArray f = (List.zipWith f as bs).toArray := by
+  rw [Array.zipWith]
+  simp [zipWithAux_toArray_zero]
+
+@[simp] theorem zip_toArray (as : List α) (bs : List β) :
+    Array.zip as.toArray bs.toArray = (List.zip as bs).toArray := by
+  simp [Array.zip, zipWith_toArray, zip]
+
 end List
 
 namespace Array
@@ -1099,8 +1144,6 @@ theorem filterMap_congr {as bs : Array α} (h : as = bs)
 
 theorem size_empty : (#[] : Array α).size = 0 := rfl
 
-@[simp] theorem toList_empty : (#[] : Array α).toList = [] := rfl
-
 /-! ### append -/
 
 theorem push_eq_append_singleton (as : Array α) (x) : as.push x = as ++ #[x] := rfl
@@ -1150,15 +1193,6 @@ theorem getElem?_append {as bs : Array α} {n : Nat} :
   · exact getElem?_append_left h
   · exact getElem?_append_right (by simpa using h)
 
-@[simp] theorem append_nil (as : Array α) : as ++ #[] = as := by
-  apply ext'; simp only [toList_append, toList_empty, List.append_nil]
-
-@[simp] theorem nil_append (as : Array α) : #[] ++ as = as := by
-  apply ext'; simp only [toList_append, toList_empty, List.nil_append]
-
-@[simp] theorem append_assoc (as bs cs : Array α) : as ++ bs ++ cs = as ++ (bs ++ cs) := by
-  apply ext'; simp only [toList_append, List.append_assoc]
-
 /-! ### flatten -/
 
 @[simp] theorem toList_flatten {l : Array (Array α)} :
@@ -1523,6 +1557,18 @@ theorem feraseIdx_eq_eraseIdx {a : Array α} {i : Fin a.size} :
   cases bs
   simp
 
+/-! ### zipWith -/
+
+@[simp] theorem toList_zipWith (f : α → β → γ) (as : Array α) (bs : Array β) :
+    (Array.zipWith as bs f).toList = List.zipWith f as.toList bs.toList := by
+  cases as
+  cases bs
+  simp
+
+@[simp] theorem toList_zip (as : Array α) (bs : Array β) :
+    (Array.zip as bs).toList = List.zip as.toList bs.toList := by
+  simp [zip, toList_zipWith, List.zip]
+
 end Array
 
 open Array
@@ -1538,11 +1584,6 @@ Our goal is to have `simp` "pull `List.toArray` outwards" as much as possible.
 @[simp] theorem toListRev_toArray (l : List α) : l.toArray.toListRev = l.reverse := by
   simp
 
-@[simp] theorem push_append_toArray (as : Array α) (a : α) (l : List α) :
-    as.push a ++ l.toArray = as ++ (a :: l).toArray := by
-  apply ext'
-  simp
-
 @[simp] theorem take_toArray (l : List α) (n : Nat) : l.toArray.take n = (l.take n).toArray := by
   apply ext'
   simp
@@ -1764,6 +1805,44 @@ namespace Array
   induction as
   simp
 
+/-! ### unzip -/
+
+@[simp] theorem fst_unzip (as : Array (α × β)) : (Array.unzip as).fst = as.map Prod.fst := by
+  simp only [unzip]
+  rcases as with ⟨as⟩
+  simp only [List.foldl_toArray']
+  rw [← List.foldl_hom (f := Prod.fst) (g₂ := fun bs x => bs.push x.1) (H := by simp), ← List.foldl_map]
+  simp
+
+@[simp] theorem snd_unzip (as : Array (α × β)) : (Array.unzip as).snd = as.map Prod.snd := by
+  simp only [unzip]
+  rcases as with ⟨as⟩
+  simp only [List.foldl_toArray']
+  rw [← List.foldl_hom (f := Prod.snd) (g₂ := fun bs x => bs.push x.2) (H := by simp), ← List.foldl_map]
+  simp
+
+end Array
+
+namespace List
+
+@[simp] theorem unzip_toArray (as : List (α × β)) :
+    as.toArray.unzip = Prod.map List.toArray List.toArray as.unzip := by
+  ext1 <;> simp
+
+end List
+
+namespace Array
+
+@[simp] theorem toList_fst_unzip (as : Array (α × β)) :
+    as.unzip.1.toList = as.toList.unzip.1 := by
+  cases as
+  simp
+
+@[simp] theorem toList_snd_unzip (as : Array (α × β)) :
+    as.unzip.2.toList = as.toList.unzip.2 := by
+  cases as
+  simp
+
 end Array
 
 /-! ### Deprecations -/
diff --git a/src/Init/Data/List/Lemmas.lean b/src/Init/Data/List/Lemmas.lean
index cca16ba548..388ad7b8b9 100644
--- a/src/Init/Data/List/Lemmas.lean
+++ b/src/Init/Data/List/Lemmas.lean
@@ -863,14 +863,14 @@ theorem foldr_map (f : α₁ → α₂) (g : α₂ → β → β) (l : List α
     (l.map f).foldr g init = l.foldr (fun x y => g (f x) y) init := by
   induction l generalizing init <;> simp [*]
 
-theorem foldl_map' {α β : Type u} (g : α → β) (f : α → α → α) (f' : β → β → β) (a : α) (l : List α)
+theorem foldl_map' (g : α → β) (f : α → α → α) (f' : β → β → β) (a : α) (l : List α)
     (h : ∀ x y, f' (g x) (g y) = g (f x y)) :
     (l.map g).foldl f' (g a) = g (l.foldl f a) := by
   induction l generalizing a
   · simp
   · simp [*, h]
 
-theorem foldr_map' {α β : Type u} (g : α → β) (f : α → α → α) (f' : β → β → β) (a : α) (l : List α)
+theorem foldr_map' (g : α → β) (f : α → α → α) (f' : β → β → β) (a : α) (l : List α)
     (h : ∀ x y, f' (g x) (g y) = g (f x y)) :
     (l.map g).foldr f' (g a) = g (l.foldr f a) := by
   induction l generalizing a
@@ -2700,6 +2700,12 @@ theorem flatMap_reverse {β} (l : List α) (f : α → List β) : (l.reverse.fla
     l.reverse.foldr f b = l.foldl (fun x y => f y x) b :=
   (foldl_reverse ..).symm.trans <| by simp
 
+theorem foldl_eq_foldr_reverse (l : List α) (f : β → α → β) (b) :
+    l.foldl f b = l.reverse.foldr (fun x y => f y x) b := by simp
+
+theorem foldr_eq_foldl_reverse (l : List α) (f : α → β → β) (b) :
+    l.foldr f b = l.reverse.foldl (fun x y => f y x) b := by simp
+
 @[simp] theorem reverse_replicate (n) (a : α) : reverse (replicate n a) = replicate n a :=
   eq_replicate_iff.2
     ⟨by rw [length_reverse, length_replicate],