feat: Array.zipWithAll (#6191)

This PR adds `Array.zipWithAll`, and the basic lemmas relating it to
`List.zipWithAll`.

src/Init/Data/Array/Basic.lean

@@ -13,6 +13,7 @@ import Init.Data.ToString.Basic
 import Init.GetElem
 import Init.Data.List.ToArray
 import Init.Data.Array.Set
+
 universe u v w
 
 /-! ### Array literal syntax -/
@@ -883,12 +884,12 @@ def isPrefixOf [BEq α] (as bs : Array α) : Bool :=
     false
 
 @[semireducible, specialize] -- This is otherwise irreducible because it uses well-founded recursion.
-def zipWithAux (f : α → β → γ) (as : Array α) (bs : Array β) (i : Nat) (cs : Array γ) : Array γ :=
+def zipWithAux (as : Array α) (bs : Array β) (f : α → β → γ) (i : Nat) (cs : Array γ) : Array γ :=
   if h : i < as.size then
     let a := as[i]
     if h : i < bs.size then
       let b := bs[i]
-      zipWithAux f as bs (i+1) <| cs.push <| f a b
+      zipWithAux as bs f (i+1) <| cs.push <| f a b
     else
       cs
   else
@@ -896,11 +897,23 @@ def zipWithAux (f : α → β → γ) (as : Array α) (bs : Array β) (i : Nat)
 decreasing_by simp_wf; decreasing_trivial_pre_omega
 
 @[inline] def zipWith (as : Array α) (bs : Array β) (f : α → β → γ) : Array γ :=
-  zipWithAux f as bs 0 #[]
+  zipWithAux as bs f 0 #[]
 
 def zip (as : Array α) (bs : Array β) : Array (α × β) :=
   zipWith as bs Prod.mk
 
+def zipWithAll (as : Array α) (bs : Array β) (f : Option α → Option β → γ) : Array γ :=
+  go as bs 0 #[]
+where go (as : Array α) (bs : Array β) (i : Nat) (cs : Array γ) :=
+  if i < max as.size bs.size then
+    let a := as[i]?
+    let b := bs[i]?
+    go as bs (i+1) (cs.push (f a b))
+  else
+    cs
+  termination_by max as.size bs.size - i
+  decreasing_by simp_wf; decreasing_trivial_pre_omega
+
 def unzip (as : Array (α × β)) : Array α × Array β :=
   as.foldl (init := (#[], #[])) fun (as, bs) (a, b) => (as.push a, bs.push b)
 

src/Init/Data/Array/Lemmas.lean

@@ -343,8 +343,8 @@ 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
+theorem zipWithAux_toArray_succ (as : List α) (bs : List β) (f : α → β → γ) (i : Nat) (cs : Array γ) :
+    zipWithAux as.toArray bs.toArray f (i + 1) cs = zipWithAux as.tail.toArray bs.tail.toArray f i cs := by
   rw [zipWithAux]
   conv => rhs; rw [zipWithAux]
   simp only [size_toArray, getElem_toArray, length_tail, getElem_tail]
@@ -355,8 +355,8 @@ theorem zipWithAux_toArray_succ (f : α → β → γ) (as : List α) (bs : List
       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
+theorem zipWithAux_toArray_succ' (as : List α) (bs : List β) (f : α → β → γ) (i : Nat) (cs : Array γ) :
+    zipWithAux as.toArray bs.toArray f (i + 1) cs = zipWithAux (as.drop (i+1)).toArray (bs.drop (i+1)).toArray f 0 cs := by
   induction i generalizing as bs cs with
   | zero => simp [zipWithAux_toArray_succ]
   | succ i ih =>
@@ -364,7 +364,7 @@ theorem zipWithAux_toArray_succ' (f : α → β → γ) (as : List α) (bs : Lis
     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
+    zipWithAux as.toArray bs.toArray f 0 cs = cs ++ (List.zipWith f as bs).toArray := by
   rw [Array.zipWithAux]
   match as, bs with
   | [], _ => simp
@@ -372,7 +372,7 @@ theorem zipWithAux_toArray_zero (f : α → β → γ) (as : List α) (bs : List
   | 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 β) :
+@[simp] theorem zipWith_toArray (as : List α) (bs : List β) (f : α → β → γ) :
     Array.zipWith as.toArray bs.toArray f = (List.zipWith f as bs).toArray := by
   rw [Array.zipWith]
   simp [zipWithAux_toArray_zero]
@@ -381,6 +381,44 @@ theorem zipWithAux_toArray_zero (f : α → β → γ) (as : List α) (bs : List
     Array.zip as.toArray bs.toArray = (List.zip as bs).toArray := by
   simp [Array.zip, zipWith_toArray, zip]
 
+theorem zipWithAll_go_toArray (as : List α) (bs : List β) (f : Option α → Option β → γ) (i : Nat) (cs : Array γ) :
+    zipWithAll.go f as.toArray bs.toArray i cs = cs ++ (List.zipWithAll f (as.drop i) (bs.drop i)).toArray := by
+  unfold zipWithAll.go
+  split <;> rename_i h
+  · rw [zipWithAll_go_toArray]
+    simp at h
+    simp only [getElem?_toArray, push_append_toArray]
+    if ha : i < as.length then
+      if hb : i < bs.length then
+        rw [List.drop_eq_getElem_cons ha, List.drop_eq_getElem_cons hb]
+        simp only [ha, hb, getElem?_eq_getElem, zipWithAll_cons_cons]
+      else
+        simp only [Nat.not_lt] at hb
+        rw [List.drop_eq_getElem_cons ha]
+        rw [(drop_eq_nil_iff (l := bs)).mpr (by omega), (drop_eq_nil_iff (l := bs)).mpr (by omega)]
+        simp only [zipWithAll_nil, map_drop, map_cons]
+        rw [getElem?_eq_getElem ha]
+        rw [getElem?_eq_none hb]
+    else
+      if hb : i < bs.length then
+        simp only [Nat.not_lt] at ha
+        rw [List.drop_eq_getElem_cons hb]
+        rw [(drop_eq_nil_iff (l := as)).mpr (by omega), (drop_eq_nil_iff (l := as)).mpr (by omega)]
+        simp only [nil_zipWithAll, map_drop, map_cons]
+        rw [getElem?_eq_getElem hb]
+        rw [getElem?_eq_none ha]
+      else
+        omega
+  · simp only [size_toArray, Nat.not_lt] at h
+    rw [drop_eq_nil_of_le (by omega), drop_eq_nil_of_le (by omega)]
+    simp
+  termination_by max as.length bs.length - i
+  decreasing_by simp_wf; decreasing_trivial_pre_omega
+
+@[simp] theorem zipWithAll_toArray (f : Option α → Option β → γ) (as : List α) (bs : List β) :
+    Array.zipWithAll as.toArray bs.toArray f = (List.zipWithAll f as bs).toArray := by
+  simp [Array.zipWithAll, zipWithAll_go_toArray]
+
 end List
 
 namespace Array
@@ -1668,6 +1706,12 @@ theorem eraseIdx_eq_eraseIdxIfInBounds {a : Array α} {i : Nat} (h : i < a.size)
     (Array.zip as bs).toList = List.zip as.toList bs.toList := by
   simp [zip, toList_zipWith, List.zip]
 
+@[simp] theorem toList_zipWithAll (f : Option α → Option β → γ) (as : Array α) (bs : Array β) :
+    (Array.zipWithAll as bs f).toList = List.zipWithAll f as.toList bs.toList := by
+  cases as
+  cases bs
+  simp
+
 /-! ### findSomeM?, findM?, findSome?, find? -/
 
 @[simp] theorem findSomeM?_toList [Monad m] [LawfulMonad m] (p : α → m (Option β)) (as : Array α) :

src/Init/Data/List/Basic.lean

@@ -1427,10 +1427,10 @@ def zipWithAll (f : Option α → Option β → γ) : List α → List β → Li
   | a :: as, [] => (a :: as).map fun a => f (some a) none
   | a :: as, b :: bs => f a b :: zipWithAll f as bs
 
-@[simp] theorem zipWithAll_nil_right :
+@[simp] theorem zipWithAll_nil :
     zipWithAll f as [] = as.map fun a => f (some a) none := by
   cases as <;> rfl
-@[simp] theorem zipWithAll_nil_left :
+@[simp] theorem nil_zipWithAll :
     zipWithAll f [] bs = bs.map fun b => f none (some b) := rfl
 @[simp] theorem zipWithAll_cons_cons :
     zipWithAll f (a :: as) (b :: bs) = f (some a) (some b) :: zipWithAll f as bs := rfl