lean4-htt/src/Init/Data/Vector/MapIdx.lean
2025-08-05 02:29:15 +00:00

386 lines
15 KiB
Text
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

/-
Copyright (c) 2025 Lean FRO. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Kim Morrison
-/
module
prelude
public import Init.Data.Array.MapIdx
public import all Init.Data.Array.Basic
public import all Init.Data.Vector.Basic
public import Init.Data.Vector.Attach
public import Init.Data.Vector.Lemmas
public section
set_option linter.listVariables true -- Enforce naming conventions for `List`/`Array`/`Vector` variables.
set_option linter.indexVariables true -- Enforce naming conventions for index variables.
namespace Vector
/-! ### mapFinIdx -/
@[simp, grind =] theorem getElem_mapFinIdx {xs : Vector α n} {f : (i : Nat) → α → (h : i < n) → β} {i : Nat}
(h : i < n) :
(xs.mapFinIdx f)[i] = f i xs[i] h := by
rcases xs with ⟨xs, rfl⟩
simp
@[simp, grind =] theorem getElem?_mapFinIdx {xs : Vector α n} {f : (i : Nat) → α → (h : i < n) → β} {i : Nat} :
(xs.mapFinIdx f)[i]? =
xs[i]?.pbind fun b h => some <| f i b (getElem?_eq_some_iff.1 h).1 := by
simp only [getElem?_def, getElem_mapFinIdx]
split <;> simp_all
/-! ### mapIdx -/
@[simp, grind =] theorem getElem_mapIdx {f : Nat → α → β} {xs : Vector α n} {i : Nat} (h : i < n) :
(xs.mapIdx f)[i] = f i (xs[i]'(by simp_all)) := by
rcases xs with ⟨xs, rfl⟩
simp
@[simp, grind =] theorem getElem?_mapIdx {f : Nat → α → β} {xs : Vector α n} {i : Nat} :
(xs.mapIdx f)[i]? = xs[i]?.map (f i) := by
rcases xs with ⟨xs, rfl⟩
simp
end Vector
namespace Array
@[simp, grind =] theorem mapFinIdx_toVector {xs : Array α} {f : (i : Nat) → α → (h : i < xs.size) → β} :
xs.toVector.mapFinIdx f = (xs.mapFinIdx f).toVector.cast (by simp) := by
ext <;> simp
@[simp, grind =] theorem mapIdx_toVector {f : Nat → α → β} {xs : Array α} :
xs.toVector.mapIdx f = (xs.mapIdx f).toVector.cast (by simp) := by
ext <;> simp
end Array
namespace Vector
/-! ### zipIdx -/
@[simp, grind =] theorem toList_zipIdx {xs : Vector α n} (k : Nat := 0) :
(xs.zipIdx k).toList = xs.toList.zipIdx k := by
rcases xs with ⟨xs, rfl⟩
simp
@[simp, grind =] theorem getElem_zipIdx {xs : Vector α n} {i : Nat} {h : i < n} :
(xs.zipIdx k)[i] = (xs[i]'(by simp_all), k + i) := by
rcases xs with ⟨xs, rfl⟩
simp
theorem mk_mem_zipIdx_iff_le_and_getElem?_sub {x : α} {i : Nat} {xs : Vector α n} {k : Nat} :
(x, i) ∈ xs.zipIdx k ↔ k ≤ i ∧ xs[i - k]? = some x := by
rcases xs with ⟨xs, rfl⟩
simp [Array.mk_mem_zipIdx_iff_le_and_getElem?_sub]
/-- Variant of `mk_mem_zipIdx_iff_le_and_getElem?_sub` specialized at `k = 0`,
to avoid the inequality and the subtraction. -/
theorem mk_mem_zipIdx_iff_getElem? {x : α} {i : Nat} {xs : Vector α n} :
(x, i) ∈ xs.zipIdx ↔ xs[i]? = some x := by
rcases xs with ⟨xs, rfl⟩
simp [Array.mk_mem_zipIdx_iff_le_and_getElem?_sub]
theorem mem_zipIdx_iff_le_and_getElem?_sub {x : α × Nat} {xs : Vector α n} {k : Nat} :
x ∈ xs.zipIdx k ↔ k ≤ x.2 ∧ xs[x.2 - k]? = some x.1 := by
rcases xs with ⟨xs, rfl⟩
simp [Array.mem_zipIdx_iff_le_and_getElem?_sub]
/-- Variant of `mem_zipIdx_iff_le_and_getElem?_sub` specialized at `k = 0`,
to avoid the inequality and the subtraction. -/
theorem mem_zipIdx_iff_getElem? {x : α × Nat} {xs : Vector α n} :
x ∈ xs.zipIdx ↔ xs[x.2]? = some x.1 := by
rcases xs with ⟨xs, rfl⟩
simp [Array.mem_zipIdx_iff_getElem?]
/-! ### mapFinIdx -/
@[congr] theorem mapFinIdx_congr {xs ys : Vector α n} (w : xs = ys)
(f : (i : Nat) → α → (h : i < n) → β) :
mapFinIdx xs f = mapFinIdx ys f := by
subst w
rfl
@[simp, grind =]
theorem mapFinIdx_empty {f : (i : Nat) → α → (h : i < 0) → β} : mapFinIdx #v[] f = #v[] :=
rfl
theorem mapFinIdx_eq_ofFn {as : Vector α n} {f : (i : Nat) → α → (h : i < n) → β} :
as.mapFinIdx f = Vector.ofFn fun i : Fin n => f i as[i] i.2 := by
rcases as with ⟨as, rfl⟩
simp [Array.mapFinIdx_eq_ofFn]
@[grind =]
theorem mapFinIdx_append {xs : Vector α n} {ys : Vector α m} {f : (i : Nat) → α → (h : i < n + m) → β} :
(xs ++ ys).mapFinIdx f =
xs.mapFinIdx (fun i a h => f i a (by omega)) ++
ys.mapFinIdx (fun i a h => f (i + n) a (by omega)) := by
rcases xs with ⟨xs, rfl⟩
rcases ys with ⟨ys, rfl⟩
simp [Array.mapFinIdx_append]
@[simp, grind =]
theorem mapFinIdx_push {xs : Vector α n} {a : α} {f : (i : Nat) → α → (h : i < n + 1) → β} :
mapFinIdx (xs.push a) f =
(mapFinIdx xs (fun i a h => f i a (by omega))).push (f n a (by simp)) := by
simp [← append_singleton, mapFinIdx_append]
theorem mapFinIdx_singleton {a : α} {f : (i : Nat) → α → (h : i < 1) → β} :
#v[a].mapFinIdx f = #v[f 0 a (by simp)] := by
simp
theorem mapFinIdx_eq_zipIdx_map {xs : Vector α n} {f : (i : Nat) → α → (h : i < n) → β} :
xs.mapFinIdx f = xs.zipIdx.attach.map
fun ⟨⟨x, i⟩, m⟩ =>
f i x (by rw [mk_mem_zipIdx_iff_getElem?, getElem?_eq_some_iff] at m; exact m.1) := by
ext <;> simp
theorem exists_of_mem_mapFinIdx {b : β} {xs : Vector α n} {f : (i : Nat) → α → (h : i < n) → β}
(h : b ∈ xs.mapFinIdx f) : ∃ (i : Nat) (h : i < n), f i xs[i] h = b := by
rcases xs with ⟨xs, rfl⟩
exact List.exists_of_mem_mapFinIdx (by simpa using h)
@[simp, grind =] theorem mem_mapFinIdx {b : β} {xs : Vector α n} {f : (i : Nat) → α → (h : i < n) → β} :
b ∈ xs.mapFinIdx f ↔ ∃ (i : Nat) (h : i < n), f i xs[i] h = b := by
rcases xs with ⟨xs, rfl⟩
simp
theorem mapFinIdx_eq_iff {xs : Vector α n} {f : (i : Nat) → α → (h : i < n) → β} :
xs.mapFinIdx f = xs' ↔ ∀ (i : Nat) (h : i < n), xs'[i] = f i xs[i] h := by
rcases xs with ⟨xs, rfl⟩
rcases xs' with ⟨xs', h⟩
simp [mapFinIdx_mk, eq_mk, getElem_mk, Array.mapFinIdx_eq_iff, h]
@[simp] theorem mapFinIdx_eq_singleton_iff {xs : Vector α 1} {f : (i : Nat) → α → (h : i < 1) → β} {b : β} :
xs.mapFinIdx f = #v[b] ↔ ∃ (a : α), xs = #v[a] ∧ f 0 a (by omega) = b := by
rcases xs with ⟨xs, h⟩
simp only [mapFinIdx_mk, eq_mk, Array.mapFinIdx_eq_singleton_iff]
constructor
· rintro ⟨a, rfl, rfl⟩
exact ⟨a, by simp⟩
· rintro ⟨a, rfl, rfl⟩
exact ⟨a, by simp⟩
theorem mapFinIdx_eq_append_iff {xs : Vector α (n + m)} {f : (i : Nat) → α → (h : i < n + m) → β}
{ys : Vector β n} {zs : Vector β m} :
xs.mapFinIdx f = ys ++ zs ↔
∃ (ys' : Vector α n) (zs' : Vector α m), xs = ys' ++ zs' ∧
ys'.mapFinIdx (fun i a h => f i a (by omega)) = ys ∧
zs'.mapFinIdx (fun i a h => f (i + n) a (by omega)) = zs := by
rcases xs with ⟨xs, h⟩
rcases ys with ⟨ys, rfl⟩
rcases zs with ⟨zs, rfl⟩
simp only [mapFinIdx_mk, mk_append_mk, eq_mk, Array.mapFinIdx_eq_append_iff, toArray_mapFinIdx,
mk_eq, toArray_append]
constructor
· rintro ⟨ys', zs', rfl, h₁, h₂⟩
have h₁' := congrArg Array.size h₁
have h₂' := congrArg Array.size h₂
simp only [Array.size_mapFinIdx] at h₁' h₂'
exact ⟨⟨ys', h₁'⟩, ⟨zs', h₂'⟩, by simp_all⟩
· rintro ⟨⟨ys', s₁⟩, ⟨zs', s₂⟩, rfl, h₁, h₂⟩
refine ⟨ys', zs', by simp_all⟩
theorem mapFinIdx_eq_push_iff {xs : Vector α (n + 1)} {b : β} {f : (i : Nat) → α → (h : i < n + 1) → β} {ys : Vector β n} :
xs.mapFinIdx f = ys.push b ↔
∃ (zs : Vector α n) (a : α), xs = zs.push a ∧
zs.mapFinIdx (fun i a h => f i a (by omega)) = ys ∧ b = f n a (by omega) := by
rcases xs with ⟨xs, h⟩
rcases ys with ⟨ys, rfl⟩
simp only [mapFinIdx_mk, push_mk, eq_mk, Array.mapFinIdx_eq_push_iff, mk_eq, toArray_push,
toArray_mapFinIdx]
constructor
· rintro ⟨zs, a, rfl, h₁, rfl⟩
simp only [Array.size_push, Nat.add_right_cancel_iff] at h
exact ⟨⟨zs, h⟩, a, by simp_all⟩
· rintro ⟨⟨zs, h⟩, a, rfl, h₁, rfl⟩
exact ⟨zs, a, by simp_all⟩
theorem mapFinIdx_eq_mapFinIdx_iff {xs : Vector α n} {f g : (i : Nat) → α → (h : i < n) → β} :
xs.mapFinIdx f = xs.mapFinIdx g ↔ ∀ (i : Nat) (h : i < n), f i xs[i] h = g i xs[i] h := by
rw [eq_comm, mapFinIdx_eq_iff]
simp
@[simp, grind =] theorem mapFinIdx_mapFinIdx {xs : Vector α n}
{f : (i : Nat) → α → (h : i < n) → β}
{g : (i : Nat) → β → (h : i < n) → γ} :
(xs.mapFinIdx f).mapFinIdx g = xs.mapFinIdx (fun i a h => g i (f i a h) h) := by
simp [mapFinIdx_eq_iff]
theorem mapFinIdx_eq_replicate_iff {xs : Vector α n} {f : (i : Nat) → α → (h : i < n) → β} {b : β} :
xs.mapFinIdx f = replicate n b ↔ ∀ (i : Nat) (h : i < n), f i xs[i] h = b := by
rcases xs with ⟨xs, rfl⟩
simp [Array.mapFinIdx_eq_replicate_iff]
@[deprecated mapFinIdx_eq_replicate_iff (since := "2025-03-18")]
abbrev mapFinIdx_eq_mkVector_iff := @mapFinIdx_eq_replicate_iff
@[simp, grind =] theorem mapFinIdx_reverse {xs : Vector α n} {f : (i : Nat) → α → (h : i < n) → β} :
xs.reverse.mapFinIdx f = (xs.mapFinIdx (fun i a h => f (n - 1 - i) a (by omega))).reverse := by
rcases xs with ⟨xs, rfl⟩
simp
/-! ### mapIdx -/
@[simp, grind =]
theorem mapIdx_empty {f : Nat → α → β} : mapIdx f #v[] = #v[] :=
rfl
@[simp] theorem mapFinIdx_eq_mapIdx {xs : Vector α n} {f : (i : Nat) → α → (h : i < n) → β} {g : Nat → α → β}
(h : ∀ (i : Nat) (h : i < n), f i xs[i] h = g i xs[i]) :
xs.mapFinIdx f = xs.mapIdx g := by
simp_all [mapFinIdx_eq_iff]
theorem mapIdx_eq_mapFinIdx {xs : Vector α n} {f : Nat → α → β} :
xs.mapIdx f = xs.mapFinIdx (fun i a _ => f i a) := by
simp [mapFinIdx_eq_mapIdx]
theorem mapIdx_eq_zipIdx_map {xs : Vector α n} {f : Nat → α → β} :
xs.mapIdx f = xs.zipIdx.map fun ⟨a, i⟩ => f i a := by
ext <;> simp
@[grind =]
theorem mapIdx_append {xs : Vector α n} {ys : Vector α m} :
(xs ++ ys).mapIdx f = xs.mapIdx f ++ ys.mapIdx fun i => f (i + n) := by
rcases xs with ⟨xs, rfl⟩
rcases ys with ⟨ys, rfl⟩
simp [Array.mapIdx_append]
@[simp, grind =]
theorem mapIdx_push {xs : Vector α n} {a : α} :
mapIdx f (xs.push a) = (mapIdx f xs).push (f n a) := by
simp [← append_singleton, mapIdx_append]
theorem mapIdx_singleton {a : α} : mapIdx f #v[a] = #v[f 0 a] := by
simp
theorem exists_of_mem_mapIdx {b : β} {xs : Vector α n}
(h : b ∈ xs.mapIdx f) : ∃ (i : Nat) (h : i < n), f i xs[i] = b := by
rw [mapIdx_eq_mapFinIdx] at h
simpa [Fin.exists_iff] using exists_of_mem_mapFinIdx h
@[simp, grind =] theorem mem_mapIdx {b : β} {xs : Vector α n} :
b ∈ xs.mapIdx f ↔ ∃ (i : Nat) (h : i < n), f i xs[i] = b := by
constructor
· intro h
exact exists_of_mem_mapIdx h
· rintro ⟨i, h, rfl⟩
rw [mem_iff_getElem]
exact ⟨i, by simpa using h, by simp⟩
theorem mapIdx_eq_push_iff {xs : Vector α (n + 1)} {b : β} :
mapIdx f xs = ys.push b ↔
∃ (a : α) (zs : Vector α n), xs = zs.push a ∧ mapIdx f zs = ys ∧ f n a = b := by
rw [mapIdx_eq_mapFinIdx, mapFinIdx_eq_push_iff]
simp only [mapFinIdx_eq_mapIdx]
constructor
· rintro ⟨zs, a, rfl, rfl, rfl⟩
exact ⟨a, zs, by simp⟩
· rintro ⟨a, zs, rfl, rfl, rfl⟩
exact ⟨zs, a, rfl, by simp⟩
theorem mapIdx_eq_singleton_iff {xs : Vector α 1} {f : Nat → α → β} {b : β} :
mapIdx f xs = #v[b] ↔ ∃ (a : α), xs = #v[a] ∧ f 0 a = b := by
simp
theorem mapIdx_eq_append_iff {xs : Vector α (n + m)} {f : Nat → α → β} {ys : Vector β n} {zs : Vector β m} :
mapIdx f xs = ys ++ zs ↔
∃ (ys' : Vector α n) (zs' : Vector α m), xs = ys' ++ zs' ∧
ys'.mapIdx f = ys ∧
zs'.mapIdx (fun i => f (i + n)) = zs := by
rcases xs with ⟨xs, h⟩
rcases ys with ⟨ys, rfl⟩
rcases zs with ⟨zs, rfl⟩
rw [mapIdx_eq_mapFinIdx, mapFinIdx_eq_append_iff]
simp
theorem mapIdx_eq_iff {xs : Vector α n} {f : Nat → α → β} {ys : Vector β n} :
mapIdx f xs = ys ↔ ∀ (i : Nat) (h : i < n), f i xs[i] = ys[i] := by
rcases xs with ⟨xs, rfl⟩
rcases ys with ⟨ys, h⟩
simp only [mapIdx_mk, eq_mk, Array.mapIdx_eq_iff, getElem_mk]
constructor
· rintro h' i h
specialize h' i
simp_all
· intro h' i
specialize h' i
by_cases w : i < xs.size
· specialize h' w
simp_all
· simp only [Nat.not_lt] at w
simp_all
theorem mapIdx_eq_mapIdx_iff {xs : Vector α n} :
mapIdx f xs = mapIdx g xs ↔ ∀ (i : Nat) (h : i < n), f i xs[i] = g i xs[i] := by
rcases xs with ⟨xs, rfl⟩
simp [Array.mapIdx_eq_mapIdx_iff]
@[simp, grind =] theorem mapIdx_set {xs : Vector α n} {i : Nat} {h : i < n} {a : α} :
(xs.set i a).mapIdx f = (xs.mapIdx f).set i (f i a) (by simpa) := by
rcases xs with ⟨xs, rfl⟩
simp
@[simp] theorem mapIdx_setIfInBounds {xs : Vector α n} {i : Nat} {a : α} :
(xs.setIfInBounds i a).mapIdx f = (xs.mapIdx f).setIfInBounds i (f i a) := by
rcases xs with ⟨xs, rfl⟩
simp
@[simp, grind =] theorem back?_mapIdx {xs : Vector α n} {f : Nat → α → β} :
(mapIdx f xs).back? = (xs.back?).map (f (n - 1)) := by
rcases xs with ⟨xs, rfl⟩
simp
@[simp, grind =] theorem back_mapIdx [NeZero n] {xs : Vector α n} {f : Nat → α → β} :
(mapIdx f xs).back = f (n - 1) (xs.back) := by
rcases xs with ⟨xs, rfl⟩
simp
@[simp, grind =] theorem mapIdx_mapIdx {xs : Vector α n} {f : Nat → α → β} {g : Nat → β → γ} :
(xs.mapIdx f).mapIdx g = xs.mapIdx (fun i => g i ∘ f i) := by
simp [mapIdx_eq_iff]
theorem mapIdx_eq_replicate_iff {xs : Vector α n} {f : Nat → α → β} {b : β} :
mapIdx f xs = replicate n b ↔ ∀ (i : Nat) (h : i < n), f i xs[i] = b := by
rcases xs with ⟨xs, rfl⟩
simp [Array.mapIdx_eq_replicate_iff]
@[deprecated mapIdx_eq_replicate_iff (since := "2025-03-18")]
abbrev mapIdx_eq_mkVector_iff := @mapIdx_eq_replicate_iff
@[simp, grind =] theorem mapIdx_reverse {xs : Vector α n} {f : Nat → α → β} :
xs.reverse.mapIdx f = (mapIdx (fun i => f (n - 1 - i)) xs).reverse := by
rcases xs with ⟨xs, rfl⟩
simp [Array.mapIdx_reverse]
theorem toArray_mapFinIdxM [Monad m] [LawfulMonad m] {xs : Vector α n}
{f : (i : Nat) → α → (h : i < n) → m β} :
toArray <$> xs.mapFinIdxM f = xs.toArray.mapFinIdxM
(fun i x h => f i x (size_toArray xs ▸ h)) := by
let rec go (i j : Nat) (inv : i + j = n) (bs : Vector β (n - i)) :
toArray <$> mapFinIdxM.map xs f i j inv bs
= Array.mapFinIdxM.map xs.toArray (fun i x h => f i x (size_toArray xs ▸ h))
i j (size_toArray _ ▸ inv) bs.toArray := by
match i with
| 0 => simp only [mapFinIdxM.map, map_pure, Array.mapFinIdxM.map, Nat.sub_zero]
| k + 1 =>
simp only [mapFinIdxM.map, map_bind, Array.mapFinIdxM.map, getElem_toArray]
conv => lhs; arg 2; intro; rw [go]
rfl
simp only [mapFinIdxM, Array.mapFinIdxM, size_toArray]
exact go _ _ _ _
theorem toArray_mapIdxM [Monad m] [LawfulMonad m] {xs : Vector α n}
{f : Nat → α → m β} :
toArray <$> xs.mapIdxM f = xs.toArray.mapIdxM f := by
exact toArray_mapFinIdxM
end Vector