original behavour as Array.mapFinIdx

src/Init/Data/Array/Basic.lean

diff --git a/src/Init/Data/Array/Basic.lean b/src/Init/Data/Array/Basic.lean
index 62a8997..1db7a0d 100644
--- a/src/Init/Data/Array/Basic.lean
+++ b/src/Init/Data/Array/Basic.lean
@@ -398,8 +398,10 @@
   decreasing_by simp_wf; decreasing_trivial_pre_omega
   map 0 (mkEmpty as.size)
 
+/-- Variant of `mapIdxM` which receives the index as a `Fin as.size`. -/
 @[inline]
-def mapIdxM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (f : Nat → α → m β) : m (Array β) :=
+def mapFinIdxM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m]
+    (as : Array α) (f : Fin as.size → α → m β) : m (Array β) :=
   let rec @[specialize] map (i : Nat) (j : Nat) (inv : i + j = as.size) (bs : Array β) : m (Array β) := do
     match i, inv with
     | 0,    _  => pure bs
@@ -408,10 +410,14 @@
         rw [← inv, Nat.add_assoc, Nat.add_comm 1 j, Nat.add_comm]
         apply Nat.le_add_right
       have : i + (j + 1) = as.size := by rw [← inv, Nat.add_comm j 1, Nat.add_assoc]
-      map i (j+1) this (bs.push (← f j (as.get ⟨j, j_lt⟩)))
+      map i (j+1) this (bs.push (← f ⟨j, j_lt⟩ (as.get ⟨j, j_lt⟩)))
   map as.size 0 rfl (mkEmpty as.size)
 
 @[inline]
+def mapIdxM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (f : Nat → α → m β) : m (Array β) :=
+  as.mapFinIdxM fun i a => f i a
+
+@[inline]
 def findSomeM? {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (f : α → m (Option β)) : m (Option β) := do
   for a in as do
     match (← f a) with
@@ -516,6 +522,11 @@
 def map {α : Type u} {β : Type v} (f : α → β) (as : Array α) : Array β :=
   Id.run <| as.mapM f
 
+/-- Variant of `mapIdx` which receives the index as a `Fin as.size`. -/
+@[inline]
+def mapFinIdx {α : Type u} {β : Type v} (as : Array α) (f : Fin as.size → α → β) : Array β :=
+  Id.run <| as.mapFinIdxM f
+
 @[inline]
 def mapIdx {α : Type u} {β : Type v} (as : Array α) (f : Nat → α → β) : Array β :=
   Id.run <| as.mapIdxM f

src/Init/Data/Array/MapIdx.lean

diff --git a/src/Init/Data/Array/MapIdx.lean b/src/Init/Data/Array/MapIdx.lean
index 4b910b3..8c90ae8 100644
--- a/src/Init/Data/Array/MapIdx.lean
+++ b/src/Init/Data/Array/MapIdx.lean
@@ -9,25 +9,24 @@
 
 namespace Array
 
-
-/-! ### mapIdx -/
+/-! ### mapFinIdx -/
 
 -- This could also be proved from `SatisfiesM_mapIdxM` in Batteries.
-theorem mapIdx_induction (as : Array α) (f : Nat → α → β)
+theorem mapFinIdx_induction (as : Array α) (f : Fin as.size → α → β)
     (motive : Nat → Prop) (h0 : motive 0)
     (p : Fin as.size → β → Prop)
     (hs : ∀ i, motive i.1 → p i (f i as[i]) ∧ motive (i + 1)) :
-    motive as.size ∧ ∃ eq : (Array.mapIdx as f).size = as.size,
-      ∀ i h, p ⟨i, h⟩ ((Array.mapIdx as f)[i]) := by
+    motive as.size ∧ ∃ eq : (Array.mapFinIdx as f).size = as.size,
+      ∀ i h, p ⟨i, h⟩ ((Array.mapFinIdx as f)[i]) := by
   let rec go {bs i j h} (h₁ : j = bs.size) (h₂ : ∀ i h h', p ⟨i, h⟩ bs[i]) (hm : motive j) :
-    let arr : Array β := Array.mapIdxM.map (m := Id) as f i j h bs
+    let arr : Array β := Array.mapFinIdxM.map (m := Id) as f i j h bs
     motive as.size ∧ ∃ eq : arr.size = as.size, ∀ i h, p ⟨i, h⟩ arr[i] := by
-    induction i generalizing j bs with simp [mapIdxM.map]
+    induction i generalizing j bs with simp [mapFinIdxM.map]
     | zero =>
       have := (Nat.zero_add _).symm.trans h
       exact ⟨this ▸ hm, h₁ ▸ this, fun _ _ => h₂ ..⟩
     | succ i ih =>
-      apply @ih (bs.push (f j as[j])) (j + 1) (by omega) (by simp; omega)
+      apply @ih (bs.push (f ⟨j, by omega⟩ as[j])) (j + 1) (by omega) (by simp; omega)
       · intro i i_lt h'
         rw [get_push]
         split
@@ -36,7 +35,40 @@
           obtain rfl : i = j := by omega
           apply (hs ⟨i, by omega⟩ hm).1
       · exact (hs ⟨j, by omega⟩ hm).2
-  simp [mapIdx, mapIdxM]; exact go rfl nofun h0
+  simp [mapFinIdx, mapFinIdxM]; exact go rfl nofun h0
+
+theorem mapFinIdx_spec (as : Array α) (f : Fin as.size → α → β)
+    (p : Fin as.size → β → Prop) (hs : ∀ i, p i (f i as[i])) :
+    ∃ eq : (Array.mapFinIdx as f).size = as.size,
+      ∀ i h, p ⟨i, h⟩ ((Array.mapFinIdx as f)[i]) :=
+  (mapFinIdx_induction _ _ (fun _ => True) trivial p fun _ _ => ⟨hs .., trivial⟩).2
+
+@[simp] theorem size_mapFinIdx (a : Array α) (f : Fin a.size → α → β) : (a.mapFinIdx f).size = a.size :=
+  (mapFinIdx_spec (p := fun _ _ => True) (hs := fun _ => trivial)).1
+
+@[simp] theorem size_zipWithIndex (as : Array α) : as.zipWithIndex.size = as.size :=
+  Array.size_mapFinIdx _ _
+
+@[simp] theorem getElem_mapFinIdx (a : Array α) (f : Fin a.size → α → β) (i : Nat)
+    (h : i < (mapFinIdx a f).size) :
+    (a.mapFinIdx f)[i] = f ⟨i, by simp_all⟩ (a[i]'(by simp_all)) :=
+  (mapFinIdx_spec _ _ (fun i b => b = f i a[i]) fun _ => rfl).2 i _
+
+@[simp] theorem getElem?_mapFinIdx (a : Array α) (f : Fin a.size → α → β) (i : Nat) :
+    (a.mapFinIdx f)[i]? =
+      a[i]?.pbind fun b h => f ⟨i, (getElem?_eq_some_iff.1 h).1⟩ b := by
+  simp only [getElem?_def, size_mapFinIdx, getElem_mapFinIdx]
+  split <;> simp_all
+
+/-! ### mapIdx -/
+
+theorem mapIdx_induction (as : Array α) (f : Nat → α → β)
+    (motive : Nat → Prop) (h0 : motive 0)
+    (p : Fin as.size → β → Prop)
+    (hs : ∀ i, motive i.1 → p i (f i as[i]) ∧ motive (i + 1)) :
+    motive as.size ∧ ∃ eq : (Array.mapIdx as f).size = as.size,
+      ∀ i h, p ⟨i, h⟩ ((Array.mapIdx as f)[i]) :=
+  mapFinIdx_induction as (fun i a => f i a) motive h0 p hs
 
 theorem mapIdx_spec (as : Array α) (f : Nat → α → β)
     (p : Fin as.size → β → Prop) (hs : ∀ i, p i (f i as[i])) :
@@ -47,17 +79,14 @@
 @[simp] theorem size_mapIdx (a : Array α) (f : Nat → α → β) : (a.mapIdx f).size = a.size :=
   (mapIdx_spec (p := fun _ _ => True) (hs := fun _ => trivial)).1
 
-@[simp] theorem size_zipWithIndex (as : Array α) : as.zipWithIndex.size = as.size :=
-  Array.size_mapIdx _ _
-
 @[simp] theorem getElem_mapIdx (a : Array α) (f : Nat → α → β) (i : Nat)
     (h : i < (mapIdx a f).size) :
     (a.mapIdx f)[i] = f i (a[i]'(by simp_all)) :=
-  (mapIdx_spec _ _ (fun i b => b = f i a[i]) fun _ => rfl).2 i (by simpa using h)
+  (mapIdx_spec _ _ (fun i b => b = f i a[i]) fun _ => rfl).2 i (by simp_all)
 
 @[simp] theorem getElem?_mapIdx (a : Array α) (f : Nat → α → β) (i : Nat) :
-    (a.mapIdx f)[i]? = a[i]?.map (f i) := by
-  simp only [getElem?_def, size_mapIdx, getElem_mapIdx]
-  split <;> simp_all
+    (a.mapIdx f)[i]? =
+      a[i]?.map (f i) := by
+  simp [getElem?_def, size_mapIdx, getElem_mapIdx]
 
 end Array