From 40eefb1df670e25c36e328b3de8e457be02ce0f9 Mon Sep 17 00:00:00 2001 From: Kim Morrison Date: Fri, 31 Jan 2025 11:06:51 +1100 Subject: [PATCH] feat: alignment of `List/Array/Vector` lemmas about `range`, `range'`, `zipIdx` (#6878) This PR completes alignments of `List/Array/Vector` lemmas about `range`, `range'`, and `zipIdx`. --- src/Init/Data/Array.lean | 1 + src/Init/Data/Array/Basic.lean | 6 +- src/Init/Data/Array/Lemmas.lean | 29 ++- src/Init/Data/Array/MapIdx.lean | 2 +- src/Init/Data/Array/Range.lean | 297 ++++++++++++++++++++++++++++++ src/Init/Data/List/Nat/Range.lean | 9 +- src/Init/Data/List/Range.lean | 20 +- src/Init/Data/Vector.lean | 1 + src/Init/Data/Vector/Basic.lean | 6 +- src/Init/Data/Vector/MapIdx.lean | 2 +- src/Init/Data/Vector/Range.lean | 271 +++++++++++++++++++++++++++ 11 files changed, 627 insertions(+), 17 deletions(-) create mode 100644 src/Init/Data/Array/Range.lean create mode 100644 src/Init/Data/Vector/Range.lean diff --git a/src/Init/Data/Array.lean b/src/Init/Data/Array.lean index ef277cb58d..2a91c25595 100644 --- a/src/Init/Data/Array.lean +++ b/src/Init/Data/Array.lean @@ -23,5 +23,6 @@ import Init.Data.Array.FinRange import Init.Data.Array.Perm import Init.Data.Array.Find import Init.Data.Array.Lex +import Init.Data.Array.Range import Init.Data.Array.Erase import Init.Data.Array.Zip diff --git a/src/Init/Data/Array/Basic.lean b/src/Init/Data/Array/Basic.lean index 3f7976ae45..473055b35f 100644 --- a/src/Init/Data/Array/Basic.lean +++ b/src/Init/Data/Array/Basic.lean @@ -244,6 +244,10 @@ def ofFn {n} (f : Fin n → α) : Array α := go 0 (mkEmpty n) where def range (n : Nat) : Array Nat := ofFn fun (i : Fin n) => i +/-- The array `#[start, start + step, ..., start + step * (size - 1)]`. -/ +def range' (start size : Nat) (step : Nat := 1) : Array Nat := + ofFn fun (i : Fin size) => start + step * i + @[inline] protected def singleton (v : α) : Array α := #[v] def back! [Inhabited α] (a : Array α) : α := @@ -632,7 +636,7 @@ def mapIdx {α : Type u} {β : Type v} (f : Nat → α → β) (as : Array α) : /-- Turns `#[a, b]` into `#[(a, 0), (b, 1)]`. -/ def zipIdx (arr : Array α) (start := 0) : Array (α × Nat) := - arr.mapIdx fun i a => (a, i + start) + arr.mapIdx fun i a => (a, start + i) @[deprecated zipIdx (since := "2025-01-21")] abbrev zipWithIndex := @zipIdx diff --git a/src/Init/Data/Array/Lemmas.lean b/src/Init/Data/Array/Lemmas.lean index 30ef78c2bc..f55af4793a 100644 --- a/src/Init/Data/Array/Lemmas.lean +++ b/src/Init/Data/Array/Lemmas.lean @@ -3360,15 +3360,35 @@ theorem size_eq_length_toList (as : Array α) : as.size = as.toList.length := rf @[deprecated size_swapIfInBounds (since := "2024-11-24")] abbrev size_swap! := @size_swapIfInBounds @[simp] theorem size_range {n : Nat} : (range n).size = n := by - induction n <;> simp [range] + simp [range] @[simp] theorem toList_range (n : Nat) : (range n).toList = List.range n := by apply List.ext_getElem <;> simp [range] @[simp] -theorem getElem_range {n : Nat} {x : Nat} (h : x < (Array.range n).size) : (Array.range n)[x] = x := by +theorem getElem_range {n : Nat} {i : Nat} (h : i < (Array.range n).size) : (Array.range n)[i] = i := by simp [← getElem_toList] +theorem getElem?_range {n : Nat} {i : Nat} : (Array.range n)[i]? = if i < n then some i else none := by + simp [getElem?_def, getElem_range] + +@[simp] theorem size_range' {start size step} : (range' start size step).size = size := by + simp [range'] + +@[simp] theorem toList_range' {start size step} : + (range' start size step).toList = List.range' start size step := by + apply List.ext_getElem <;> simp [range'] + +@[simp] +theorem getElem_range' {start size step : Nat} {i : Nat} + (h : i < (Array.range' start size step).size) : + (Array.range' start size step)[i] = start + step * i := by + simp [← getElem_toList] + +theorem getElem?_range' {start size step : Nat} {i : Nat} : + (Array.range' start size step)[i]? = if i < size then some (start + step * i) else none := by + simp [getElem?_def, getElem_range'] + /-! ### shrink -/ @[simp] theorem size_shrink_loop (a : Array α) (n : Nat) : (shrink.loop n a).size = a.size - n := by @@ -3649,6 +3669,11 @@ theorem uset_toArray (l : List α) (i : USize) (a : α) (h : i.toNat < l.toArray apply ext' simp +@[simp] theorem toArray_range' (start size step : Nat) : + (range' start size step).toArray = Array.range' start size step := by + apply ext' + simp + @[simp] theorem toArray_ofFn (f : Fin n → α) : (ofFn f).toArray = Array.ofFn f := by ext <;> simp diff --git a/src/Init/Data/Array/MapIdx.lean b/src/Init/Data/Array/MapIdx.lean index 3217801d66..e6dd3e63f8 100644 --- a/src/Init/Data/Array/MapIdx.lean +++ b/src/Init/Data/Array/MapIdx.lean @@ -120,7 +120,7 @@ namespace Array /-! ### zipIdx -/ @[simp] theorem getElem_zipIdx (a : Array α) (k : Nat) (i : Nat) (h : i < (a.zipIdx k).size) : - (a.zipIdx k)[i] = (a[i]'(by simp_all), i + k) := by + (a.zipIdx k)[i] = (a[i]'(by simp_all), k + i) := by simp [zipIdx] @[deprecated getElem_zipIdx (since := "2025-01-21")] diff --git a/src/Init/Data/Array/Range.lean b/src/Init/Data/Array/Range.lean new file mode 100644 index 0000000000..5d154563d0 --- /dev/null +++ b/src/Init/Data/Array/Range.lean @@ -0,0 +1,297 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Kim Morrison +-/ +prelude +import Init.Data.Array.Lemmas +import Init.Data.Array.OfFn +import Init.Data.Array.MapIdx +import Init.Data.Array.Zip +import Init.Data.List.Nat.Range + +/-! +# Lemmas about `Array.range'`, `Array.range`, and `Array.zipIdx` + +-/ + +namespace Array + +open Nat + +/-! ## Ranges and enumeration -/ + +/-! ### range' -/ + +theorem range'_succ (s n step) : range' s (n + 1) step = #[s] ++ range' (s + step) n step := by + rw [← toList_inj] + simp [List.range'_succ] + +@[simp] theorem range'_eq_empty_iff : range' s n step = #[] ↔ n = 0 := by + rw [← size_eq_zero, size_range'] + +theorem range'_ne_empty_iff (s : Nat) {n step : Nat} : range' s n step ≠ #[] ↔ n ≠ 0 := by + cases n <;> simp + +@[simp] theorem range'_zero : range' s 0 step = #[] := by + simp + +@[simp] theorem range'_one {s step : Nat} : range' s 1 step = #[s] := rfl + +@[simp] theorem range'_inj : range' s n = range' s' n' ↔ n = n' ∧ (n = 0 ∨ s = s') := by + rw [← toList_inj] + simp [List.range'_inj] + +theorem mem_range' {n} : m ∈ range' s n step ↔ ∃ i < n, m = s + step * i := by + simp [range'] + constructor + · rintro ⟨⟨i, w⟩, h, h'⟩ + exact ⟨i, w, by simp_all⟩ + · rintro ⟨i, w, h'⟩ + exact ⟨⟨i, w⟩, by simp_all⟩ + +theorem pop_range' : (range' s n step).pop = range' s (n - 1) step := by + ext <;> simp + +theorem map_add_range' (a) (s n step) : map (a + ·) (range' s n step) = range' (a + s) n step := by + ext <;> simp <;> omega + +theorem range'_succ_left : range' (s + 1) n step = (range' s n step).map (· + 1) := by + ext <;> simp <;> omega + +theorem range'_append (s m n step : Nat) : + range' s m step ++ range' (s + step * m) n step = range' s (m + n) step := by + ext i h₁ h₂ + · simp + · simp only [size_append, size_range'] at h₁ h₂ + simp only [getElem_append, size_range', getElem_range', Nat.mul_sub_left_distrib, dite_eq_ite, + ite_eq_left_iff, Nat.not_lt] + intro h + have : step * m ≤ step * i := by exact mul_le_mul_left step h + omega + +@[simp] theorem range'_append_1 (s m n : Nat) : + range' s m ++ range' (s + m) n = range' s (m + n) := by simpa using range'_append s m n 1 + +theorem range'_concat (s n : Nat) : range' s (n + 1) step = range' s n step ++ #[s + step * n] := by + exact (range'_append s n 1 step).symm + +theorem range'_1_concat (s n : Nat) : range' s (n + 1) = range' s n ++ #[s + n] := by + simp [range'_concat] + +@[simp] theorem mem_range'_1 : m ∈ range' s n ↔ s ≤ m ∧ m < s + n := by + simp [mem_range']; exact ⟨ + fun ⟨i, h, e⟩ => e ▸ ⟨Nat.le_add_right .., Nat.add_lt_add_left h _⟩, + fun ⟨h₁, h₂⟩ => ⟨m - s, Nat.sub_lt_left_of_lt_add h₁ h₂, (Nat.add_sub_cancel' h₁).symm⟩⟩ + +theorem map_sub_range' (a s n : Nat) (h : a ≤ s) : + map (· - a) (range' s n step) = range' (s - a) n step := by + conv => lhs; rw [← Nat.add_sub_cancel' h] + rw [← map_add_range', map_map, (?_ : _∘_ = _), map_id] + funext x; apply Nat.add_sub_cancel_left + +@[simp] theorem range'_eq_singleton_iff {s n a : Nat} : range' s n = #[a] ↔ s = a ∧ n = 1 := by + rw [← toList_inj] + simp + +theorem range'_eq_append_iff : range' s n = xs ++ ys ↔ ∃ k, k ≤ n ∧ xs = range' s k ∧ ys = range' (s + k) (n - k) := by + simp [← toList_inj, List.range'_eq_append_iff] + +@[simp] theorem find?_range'_eq_some {s n : Nat} {i : Nat} {p : Nat → Bool} : + (range' s n).find? p = some i ↔ p i ∧ i ∈ range' s n ∧ ∀ j, s ≤ j → j < i → !p j := by + rw [← List.toArray_range'] + simp only [List.find?_toArray, mem_toArray] + simp [List.find?_range'_eq_some] + +@[simp] theorem find?_range'_eq_none {s n : Nat} {p : Nat → Bool} : + (range' s n).find? p = none ↔ ∀ i, s ≤ i → i < s + n → !p i := by + rw [← List.toArray_range'] + simp only [List.find?_toArray] + simp + +theorem erase_range' : + (range' s n).erase i = + range' s (min n (i - s)) ++ range' (max s (i + 1)) (min s (i + 1) + n - (i + 1)) := by + simp only [← List.toArray_range', List.erase_toArray] + simp [List.erase_range'] + +/-! ### range -/ + +theorem range_eq_range' (n : Nat) : range n = range' 0 n := by + simp [range, range'] + +theorem range_succ_eq_map (n : Nat) : range (n + 1) = #[0] ++ map succ (range n) := by + ext i h₁ h₂ + · simp + omega + · simp only [getElem_range, getElem_append, size_toArray, List.length_cons, List.length_nil, + Nat.zero_add, lt_one_iff, List.getElem_toArray, List.getElem_singleton, getElem_map, + succ_eq_add_one, dite_eq_ite] + split <;> omega + +theorem range'_eq_map_range (s n : Nat) : range' s n = map (s + ·) (range n) := by + rw [range_eq_range', map_add_range']; rfl + +@[simp] theorem range_eq_empty_iff {n : Nat} : range n = #[] ↔ n = 0 := by + rw [← size_eq_zero, size_range] + +theorem range_ne_empty_iff {n : Nat} : range n ≠ #[] ↔ n ≠ 0 := by + cases n <;> simp + +theorem range_succ (n : Nat) : range (succ n) = range n ++ #[n] := by + ext i h₁ h₂ + · simp + · simp only [succ_eq_add_one, size_range] at h₁ + simp only [succ_eq_add_one, getElem_range, append_singleton, getElem_push, size_range, + dite_eq_ite] + split <;> omega + +theorem range_add (a b : Nat) : range (a + b) = range a ++ (range b).map (a + ·) := by + rw [← range'_eq_map_range] + simpa [range_eq_range', Nat.add_comm] using (range'_append_1 0 a b).symm + +theorem reverse_range' (s n : Nat) : reverse (range' s n) = map (s + n - 1 - ·) (range n) := by + simp [← toList_inj, List.reverse_range'] + +@[simp] +theorem mem_range {m n : Nat} : m ∈ range n ↔ m < n := by + simp only [range_eq_range', mem_range'_1, Nat.zero_le, true_and, Nat.zero_add] + +theorem not_mem_range_self {n : Nat} : n ∉ range n := by simp + +theorem self_mem_range_succ (n : Nat) : n ∈ range (n + 1) := by simp + +@[simp] theorem take_range (m n : Nat) : take (range n) m = range (min m n) := by + ext <;> simp + +@[simp] theorem find?_range_eq_some {n : Nat} {i : Nat} {p : Nat → Bool} : + (range n).find? p = some i ↔ p i ∧ i ∈ range n ∧ ∀ j, j < i → !p j := by + simp [range_eq_range'] + +@[simp] theorem find?_range_eq_none {n : Nat} {p : Nat → Bool} : + (range n).find? p = none ↔ ∀ i, i < n → !p i := by + simp only [← List.toArray_range, List.find?_toArray, List.find?_range_eq_none] + +theorem erase_range : (range n).erase i = range (min n i) ++ range' (i + 1) (n - (i + 1)) := by + simp [range_eq_range', erase_range'] + + +/-! ### zipIdx -/ + +@[simp] +theorem zipIdx_eq_empty_iff {l : Array α} {n : Nat} : l.zipIdx n = #[] ↔ l = #[] := by + cases l + simp + +@[simp] +theorem getElem?_zipIdx (l : Array α) (n m) : (zipIdx l n)[m]? = l[m]?.map fun a => (a, n + m) := by + simp [getElem?_def] + +theorem map_snd_add_zipIdx_eq_zipIdx (l : Array α) (n k : Nat) : + map (Prod.map id (· + n)) (zipIdx l k) = zipIdx l (n + k) := + ext_getElem? fun i ↦ by simp [(· ∘ ·), Nat.add_comm, Nat.add_left_comm]; rfl + +@[simp] +theorem zipIdx_map_snd (n) (l : Array α) : map Prod.snd (zipIdx l n) = range' n l.size := by + cases l + simp + +@[simp] +theorem zipIdx_map_fst (n) (l : Array α) : map Prod.fst (zipIdx l n) = l := by + cases l + simp + +theorem zipIdx_eq_zip_range' (l : Array α) {n : Nat} : l.zipIdx n = l.zip (range' n l.size) := by + simp [zip_of_prod (zipIdx_map_fst _ _) (zipIdx_map_snd _ _)] + +@[simp] +theorem unzip_zipIdx_eq_prod (l : Array α) {n : Nat} : + (l.zipIdx n).unzip = (l, range' n l.size) := by + simp only [zipIdx_eq_zip_range', unzip_zip, size_range'] + +/-- Replace `zipIdx` with a starting index `n+1` with `zipIdx` starting from `n`, +followed by a `map` increasing the indices by one. -/ +theorem zipIdx_succ (l : Array α) (n : Nat) : + l.zipIdx (n + 1) = (l.zipIdx n).map (fun ⟨a, i⟩ => (a, i + 1)) := by + cases l + simp [List.zipIdx_succ] + +/-- Replace `zipIdx` with a starting index with `zipIdx` starting from 0, +followed by a `map` increasing the indices. -/ +theorem zipIdx_eq_map_add (l : Array α) (n : Nat) : + l.zipIdx n = l.zipIdx.map (fun ⟨a, i⟩ => (a, n + i)) := by + cases l + simp only [zipIdx_toArray, List.map_toArray, mk.injEq] + rw [List.zipIdx_eq_map_add] + +@[simp] +theorem zipIdx_singleton (x : α) (k : Nat) : zipIdx #[x] k = #[(x, k)] := + rfl + +theorem mk_add_mem_zipIdx_iff_getElem? {k i : Nat} {x : α} {l : Array α} : + (x, k + i) ∈ zipIdx l k ↔ l[i]? = some x := by + simp [mem_iff_getElem?, and_left_comm] + +theorem le_snd_of_mem_zipIdx {x : α × Nat} {k : Nat} {l : Array α} (h : x ∈ zipIdx l k) : + k ≤ x.2 := + (mk_mem_zipIdx_iff_le_and_getElem?_sub.1 h).1 + +theorem snd_lt_add_of_mem_zipIdx {x : α × Nat} {l : Array α} {k : Nat} (h : x ∈ zipIdx l k) : + x.2 < k + l.size := by + rcases mem_iff_getElem.1 h with ⟨i, h', rfl⟩ + simpa using h' + +theorem snd_lt_of_mem_zipIdx {x : α × Nat} {l : Array α} {k : Nat} (h : x ∈ l.zipIdx k) : x.2 < l.size + k := by + simpa [Nat.add_comm] using snd_lt_add_of_mem_zipIdx h + +theorem map_zipIdx (f : α → β) (l : Array α) (k : Nat) : + map (Prod.map f id) (zipIdx l k) = zipIdx (l.map f) k := by + cases l + simp [List.map_zipIdx] + +theorem fst_mem_of_mem_zipIdx {x : α × Nat} {l : Array α} {k : Nat} (h : x ∈ zipIdx l k) : x.1 ∈ l := + zipIdx_map_fst k l ▸ mem_map_of_mem _ h + +theorem fst_eq_of_mem_zipIdx {x : α × Nat} {l : Array α} {k : Nat} (h : x ∈ zipIdx l k) : + x.1 = l[x.2 - k]'(by have := le_snd_of_mem_zipIdx h; have := snd_lt_add_of_mem_zipIdx h; omega) := by + cases l + exact List.fst_eq_of_mem_zipIdx (by simpa using h) + +theorem mem_zipIdx {x : α} {i : Nat} {xs : Array α} {k : Nat} (h : (x, i) ∈ xs.zipIdx k) : + k ≤ i ∧ i < k + xs.size ∧ + x = xs[i - k]'(by have := le_snd_of_mem_zipIdx h; have := snd_lt_add_of_mem_zipIdx h; omega) := + ⟨le_snd_of_mem_zipIdx h, snd_lt_add_of_mem_zipIdx h, fst_eq_of_mem_zipIdx h⟩ + +/-- Variant of `mem_zipIdx` specialized at `k = 0`. -/ +theorem mem_zipIdx' {x : α} {i : Nat} {xs : Array α} (h : (x, i) ∈ xs.zipIdx) : + i < xs.size ∧ x = xs[i]'(by have := le_snd_of_mem_zipIdx h; have := snd_lt_add_of_mem_zipIdx h; omega) := + ⟨by simpa using snd_lt_add_of_mem_zipIdx h, fst_eq_of_mem_zipIdx h⟩ + +theorem zipIdx_map (l : Array α) (k : Nat) (f : α → β) : + zipIdx (l.map f) k = (zipIdx l k).map (Prod.map f id) := by + cases l + simp [List.zipIdx_map] + +theorem zipIdx_append (xs ys : Array α) (k : Nat) : + zipIdx (xs ++ ys) k = zipIdx xs k ++ zipIdx ys (k + xs.size) := by + cases xs + cases ys + simp [List.zipIdx_append] + +theorem zipIdx_eq_append_iff {l : Array α} {k : Nat} : + zipIdx l k = l₁ ++ l₂ ↔ + ∃ l₁' l₂', l = l₁' ++ l₂' ∧ l₁ = zipIdx l₁' k ∧ l₂ = zipIdx l₂' (k + l₁'.size) := by + rcases l with ⟨l⟩ + rcases l₁ with ⟨l₁⟩ + rcases l₂ with ⟨l₂⟩ + simp only [zipIdx_toArray, List.append_toArray, mk.injEq, List.zipIdx_eq_append_iff, + toArray_eq_append_iff] + constructor + · rintro ⟨l₁', l₂', rfl, rfl, rfl⟩ + exact ⟨⟨l₁'⟩, ⟨l₂'⟩, by simp⟩ + · rintro ⟨⟨l₁'⟩, ⟨l₂'⟩, rfl, h⟩ + simp only [zipIdx_toArray, mk.injEq, size_toArray] at h + obtain ⟨rfl, rfl⟩ := h + exact ⟨l₁', l₂', by simp⟩ + +end Array diff --git a/src/Init/Data/List/Nat/Range.lean b/src/Init/Data/List/Nat/Range.lean index c1f29d2e99..68eb4040d3 100644 --- a/src/Init/Data/List/Nat/Range.lean +++ b/src/Init/Data/List/Nat/Range.lean @@ -77,12 +77,15 @@ theorem map_sub_range' (a s n : Nat) (h : a ≤ s) : rw [← map_add_range', map_map, (?_ : _∘_ = _), map_id] funext x; apply Nat.add_sub_cancel_left -@[simp] theorem range'_eq_singleton {s n a : Nat} : range' s n = [a] ↔ s = a ∧ n = 1 := by +@[simp] theorem range'_eq_singleton_iff {s n a : Nat} : range' s n = [a] ↔ s = a ∧ n = 1 := by rw [range'_eq_cons_iff] - simp only [nil_eq, range'_eq_nil, and_congr_right_iff] + simp only [nil_eq, range'_eq_nil_iff, and_congr_right_iff] rintro rfl omega +@[deprecated range'_eq_singleton_iff (since := "2025-01-29")] +abbrev range'_eq_singleton := @range'_eq_singleton_iff + theorem range'_eq_append_iff : range' s n = xs ++ ys ↔ ∃ k, k ≤ n ∧ xs = range' s k ∧ ys = range' (s + k) (n - k) := by induction n generalizing s xs ys with | zero => simp @@ -174,7 +177,7 @@ theorem pairwise_lt_range (n : Nat) : Pairwise (· < ·) (range n) := by theorem pairwise_le_range (n : Nat) : Pairwise (· ≤ ·) (range n) := Pairwise.imp Nat.le_of_lt (pairwise_lt_range _) -theorem take_range (m n : Nat) : take m (range n) = range (min m n) := by +@[simp] theorem take_range (m n : Nat) : take m (range n) = range (min m n) := by apply List.ext_getElem · simp · simp +contextual [getElem_take, Nat.lt_min] diff --git a/src/Init/Data/List/Range.lean b/src/Init/Data/List/Range.lean index 094f8e1cdf..2e0b524c49 100644 --- a/src/Init/Data/List/Range.lean +++ b/src/Init/Data/List/Range.lean @@ -8,7 +8,7 @@ import Init.Data.List.Pairwise import Init.Data.List.Zip /-! -# Lemmas about `List.range` and `List.enum` +# Lemmas about `List.range` and `List.zipIdx` Most of the results are deferred to `Data.Init.List.Nat.Range`, where more results about natural arithmetic are available. @@ -29,12 +29,16 @@ theorem range'_succ (s n step) : range' s (n + 1) step = s :: range' (s + step) | 0 => rfl | _ + 1 => congrArg succ (length_range' _ _ _) -@[simp] theorem range'_eq_nil : range' s n step = [] ↔ n = 0 := by +@[simp] theorem range'_eq_nil_iff : range' s n step = [] ↔ n = 0 := by rw [← length_eq_zero, length_range'] -theorem range'_ne_nil (s : Nat) {n : Nat} : range' s n ≠ [] ↔ n ≠ 0 := by +@[deprecated range'_eq_nil_iff (since := "2025-01-29")] abbrev range'_eq_nil := @range'_eq_nil_iff + +theorem range'_ne_nil_iff (s : Nat) {n step : Nat} : range' s n step ≠ [] ↔ n ≠ 0 := by cases n <;> simp +@[deprecated range'_ne_nil_iff (since := "2025-01-29")] abbrev range'_ne_nil := @range'_ne_nil_iff + @[simp] theorem range'_zero : range' s 0 step = [] := by simp @@ -94,18 +98,18 @@ theorem range'_succ_left : range' (s + 1) n step = (range' s n step).map (· + 1 · simp [Nat.add_right_comm] theorem range'_append : ∀ s m n step : Nat, - range' s m step ++ range' (s + step * m) n step = range' s (n + m) step - | _, 0, _, _ => rfl + range' s m step ++ range' (s + step * m) n step = range' s (m + n) step + | _, 0, _, _ => by simp | s, m + 1, n, step => by simpa [range', Nat.mul_succ, Nat.add_assoc, Nat.add_comm] using range'_append (s + step) m n step @[simp] theorem range'_append_1 (s m n : Nat) : - range' s m ++ range' (s + m) n = range' s (n + m) := by simpa using range'_append s m n 1 + range' s m ++ range' (s + m) n = range' s (m + n) := by simpa using range'_append s m n 1 theorem range'_sublist_right {s m n : Nat} : range' s m step <+ range' s n step ↔ m ≤ n := ⟨fun h => by simpa only [length_range'] using h.length_le, - fun h => by rw [← Nat.sub_add_cancel h, ← range'_append]; apply sublist_append_left⟩ + fun h => by rw [← add_sub_of_le h, ← range'_append]; apply sublist_append_left⟩ theorem range'_subset_right {s m n : Nat} (step0 : 0 < step) : range' s m step ⊆ range' s n step ↔ m ≤ n := by @@ -117,7 +121,7 @@ theorem range'_subset_right_1 {s m n : Nat} : range' s m ⊆ range' s n ↔ m range'_subset_right (by decide) theorem range'_concat (s n : Nat) : range' s (n + 1) step = range' s n step ++ [s + step * n] := by - rw [Nat.add_comm n 1]; exact (range'_append s n 1 step).symm + exact (range'_append s n 1 step).symm theorem range'_1_concat (s n : Nat) : range' s (n + 1) = range' s n ++ [s + n] := by simp [range'_concat] diff --git a/src/Init/Data/Vector.lean b/src/Init/Data/Vector.lean index adbb9caf2e..fcbf80477a 100644 --- a/src/Init/Data/Vector.lean +++ b/src/Init/Data/Vector.lean @@ -12,4 +12,5 @@ import Init.Data.Vector.Count import Init.Data.Vector.DecidableEq import Init.Data.Vector.Zip import Init.Data.Vector.OfFn +import Init.Data.Vector.Range import Init.Data.Vector.Erase diff --git a/src/Init/Data/Vector/Basic.lean b/src/Init/Data/Vector/Basic.lean index d4dab1ea60..5c08b02958 100644 --- a/src/Init/Data/Vector/Basic.lean +++ b/src/Init/Data/Vector/Basic.lean @@ -315,9 +315,13 @@ This will perform the update destructively provided that the vector has a refere let a := v.toArray.swapAt! i x ⟨a.fst, a.snd, by simp [a]⟩ -/-- The vector `#v[0,1,2,...,n-1]`. -/ +/-- The vector `#v[0, 1, 2, ..., n-1]`. -/ @[inline] def range (n : Nat) : Vector Nat n := ⟨Array.range n, by simp⟩ +/-- The vector `#v[start, start + step, start + 2 * step, ..., start + (size - 1) * step]`. -/ +@[inline] def range' (start size : Nat) (step : Nat := 1) : Vector Nat size := + ⟨Array.range' start size step, by simp⟩ + /-- Compares two vectors of the same size using a given boolean relation `r`. `isEqv v w r` returns `true` if and only if `r v[i] w[i]` is true for all indices `i`. diff --git a/src/Init/Data/Vector/MapIdx.lean b/src/Init/Data/Vector/MapIdx.lean index 4c3218c172..4c8c878e65 100644 --- a/src/Init/Data/Vector/MapIdx.lean +++ b/src/Init/Data/Vector/MapIdx.lean @@ -59,7 +59,7 @@ namespace Vector simp @[simp] theorem getElem_zipIdx (a : Vector α n) (i : Nat) (h : i < n) : - (a.zipIdx k)[i] = (a[i]'(by simp_all), i + k) := by + (a.zipIdx k)[i] = (a[i]'(by simp_all), k + i) := by rcases a with ⟨a, rfl⟩ simp diff --git a/src/Init/Data/Vector/Range.lean b/src/Init/Data/Vector/Range.lean new file mode 100644 index 0000000000..9c210f963d --- /dev/null +++ b/src/Init/Data/Vector/Range.lean @@ -0,0 +1,271 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Kim Morrison +-/ +prelude +import Init.Data.Vector.Lemmas +import Init.Data.Vector.Zip +import Init.Data.Vector.MapIdx +import Init.Data.Array.Range + +/-! +# Lemmas about `Vector.range'`, `Vector.range`, and `Vector.zipIdx` + +-/ + +namespace Vector + +open Nat + +/-! ## Ranges and enumeration -/ + +/-! ### range' -/ + +@[simp] theorem toArray_range' (start size step) : + (range' start size step).toArray = Array.range' start size step := by + rfl + +theorem range'_eq_mk_range' (start size step) : + range' start size step = Vector.mk (Array.range' start size step) (by simp) := by + rfl + +@[simp] theorem getElem_range' (start size step i) (h : i < size) : + (range' start size step)[i] = start + step * i := by + simp [range', h] + +@[simp] theorem getElem?_range' (start size step i) : + (range' start size step)[i]? = if i < size then some (start + step * i) else none := by + simp [getElem?_def, range'] + +theorem range'_succ (s n step) : + range' s (n + 1) step = (#v[s] ++ range' (s + step) n step).cast (by omega) := by + rw [← toArray_inj] + simp [Array.range'_succ] + +theorem range'_zero : range' s 0 step = #v[] := by + simp + +@[simp] theorem range'_one {s step : Nat} : range' s 1 step = #v[s] := rfl + +@[simp] theorem range'_inj : range' s n = range' s' n ↔ (n = 0 ∨ s = s') := by + rw [← toArray_inj] + simp [List.range'_inj] + +theorem mem_range' {n} : m ∈ range' s n step ↔ ∃ i < n, m = s + step * i := by + simp [range', Array.mem_range'] + +theorem pop_range' : (range' s n step).pop = range' s (n - 1) step := by + ext <;> simp + +theorem map_add_range' (a) (s n step) : map (a + ·) (range' s n step) = range' (a + s) n step := by + ext <;> simp <;> omega + +theorem range'_succ_left : range' (s + 1) n step = (range' s n step).map (· + 1) := by + ext <;> simp <;> omega + +theorem range'_append (s m n step : Nat) : + range' s m step ++ range' (s + step * m) n step = range' s (m + n) step := by + rw [← toArray_inj] + simp [Array.range'_append] + +@[simp] theorem range'_append_1 (s m n : Nat) : + range' s m ++ range' (s + m) n = range' s (m + n) := by simpa using range'_append s m n 1 + +theorem range'_concat (s n : Nat) : range' s (n + 1) step = range' s n step ++ #v[s + step * n] := by + exact (range'_append s n 1 step).symm + +theorem range'_1_concat (s n : Nat) : range' s (n + 1) = range' s n ++ #v[s + n] := by + simp [range'_concat] + +@[simp] theorem mem_range'_1 : m ∈ range' s n ↔ s ≤ m ∧ m < s + n := by + simp [mem_range']; exact ⟨ + fun ⟨i, h, e⟩ => e ▸ ⟨Nat.le_add_right .., Nat.add_lt_add_left h _⟩, + fun ⟨h₁, h₂⟩ => ⟨m - s, Nat.sub_lt_left_of_lt_add h₁ h₂, (Nat.add_sub_cancel' h₁).symm⟩⟩ + +theorem map_sub_range' (a s n : Nat) (h : a ≤ s) : + map (· - a) (range' s n step) = range' (s - a) n step := by + conv => lhs; rw [← Nat.add_sub_cancel' h] + rw [← map_add_range', map_map, (?_ : _∘_ = _), map_id] + funext x; apply Nat.add_sub_cancel_left + +theorem range'_eq_append_iff : range' s (n + m) = xs ++ ys ↔ xs = range' s n ∧ ys = range' (s + n) m := by + simp only [← toArray_inj, toArray_range', toArray_append, Array.range'_eq_append_iff] + constructor + · rintro ⟨k, hk, h₁, h₂⟩ + have w : k = n := by + replace h₁ := congrArg Array.size h₁ + simp_all + subst w + simp_all + omega + · rintro ⟨h₁, h₂⟩ + exact ⟨n, by omega, by simp_all; omega⟩ + +@[simp] theorem find?_range'_eq_some {s n : Nat} {i : Nat} {p : Nat → Bool} : + (range' s n).find? p = some i ↔ p i ∧ i ∈ range' s n ∧ ∀ j, s ≤ j → j < i → !p j := by + simp [range'_eq_mk_range'] + +@[simp] theorem find?_range'_eq_none {s n : Nat} {p : Nat → Bool} : + (range' s n).find? p = none ↔ ∀ i, s ≤ i → i < s + n → !p i := by + simp [range'_eq_mk_range'] + +/-! ### range -/ + +theorem range_eq_range' (n : Nat) : range n = range' 0 n := by + simp [range, range', Array.range_eq_range'] + +theorem range_succ_eq_map (n : Nat) : range (n + 1) = + (#v[0] ++ map succ (range n)).cast (by omega) := by + rw [← toArray_inj] + simp [Array.range_succ_eq_map] + +theorem range'_eq_map_range (s n : Nat) : range' s n = map (s + ·) (range n) := by + rw [range_eq_range', map_add_range']; rfl + +theorem range_succ (n : Nat) : range (succ n) = range n ++ #v[n] := by + rw [← toArray_inj] + simp [Array.range_succ] + +theorem range_add (a b : Nat) : range (a + b) = range a ++ (range b).map (a + ·) := by + rw [← range'_eq_map_range] + simpa [range_eq_range', Nat.add_comm] using (range'_append_1 0 a b).symm + +theorem reverse_range' (s n : Nat) : reverse (range' s n) = map (s + n - 1 - ·) (range n) := by + simp [← toList_inj, List.reverse_range'] + +@[simp] +theorem mem_range {m n : Nat} : m ∈ range n ↔ m < n := by + simp only [range_eq_range', mem_range'_1, Nat.zero_le, true_and, Nat.zero_add] + +theorem not_mem_range_self {n : Nat} : n ∉ range n := by simp + +theorem self_mem_range_succ (n : Nat) : n ∈ range (n + 1) := by simp + +@[simp] theorem take_range (m n : Nat) : take (range n) m = range (min m n) := by + ext <;> simp + erw [getElem_extract] -- Why is an `erw` needed here? This should be by simp! + simp + +@[simp] theorem find?_range_eq_some {n : Nat} {i : Nat} {p : Nat → Bool} : + (range n).find? p = some i ↔ p i ∧ i ∈ range n ∧ ∀ j, j < i → !p j := by + simp [range_eq_range'] + +@[simp] theorem find?_range_eq_none {n : Nat} {p : Nat → Bool} : + (range n).find? p = none ↔ ∀ i, i < n → !p i := by + simp [range_eq_range'] + +/-! ### zipIdx -/ + +@[simp] +theorem getElem?_zipIdx (l : Vector α n) (n m) : (zipIdx l n)[m]? = l[m]?.map fun a => (a, n + m) := by + simp [getElem?_def] + +theorem map_snd_add_zipIdx_eq_zipIdx (l : Vector α n) (m k : Nat) : + map (Prod.map id (· + m)) (zipIdx l k) = zipIdx l (m + k) := by + ext <;> simp <;> omega + +@[simp] +theorem zipIdx_map_snd (m) (l : Vector α n) : map Prod.snd (zipIdx l m) = range' m n := by + rcases l with ⟨l, rfl⟩ + simp [Array.zipIdx_map_snd] + +@[simp] +theorem zipIdx_map_fst (m) (l : Vector α n) : map Prod.fst (zipIdx l m) = l := by + rcases l with ⟨l, rfl⟩ + simp [Array.zipIdx_map_fst] + +theorem zipIdx_eq_zip_range' (l : Vector α n) : l.zipIdx m = l.zip (range' m n) := by + simp [zip_of_prod (zipIdx_map_fst _ _) (zipIdx_map_snd _ _)] + +@[simp] +theorem unzip_zipIdx_eq_prod (l : Vector α n) {m : Nat} : + (l.zipIdx m).unzip = (l, range' m n) := by + simp only [zipIdx_eq_zip_range', unzip_zip] + +/-- Replace `zipIdx` with a starting index `m+1` with `zipIdx` starting from `m`, +followed by a `map` increasing the indices by one. -/ +theorem zipIdx_succ (l : Vector α n) (m : Nat) : + l.zipIdx (m + 1) = (l.zipIdx m).map (fun ⟨a, i⟩ => (a, i + 1)) := by + rcases l with ⟨l, rfl⟩ + simp [Array.zipIdx_succ] + +/-- Replace `zipIdx` with a starting index with `zipIdx` starting from 0, +followed by a `map` increasing the indices. -/ +theorem zipIdx_eq_map_add (l : Vector α n) (m : Nat) : + l.zipIdx m = l.zipIdx.map (fun ⟨a, i⟩ => (a, m + i)) := by + rcases l with ⟨l, rfl⟩ + simp only [zipIdx_mk, map_mk, eq_mk] + rw [Array.zipIdx_eq_map_add] + +@[simp] +theorem zipIdx_singleton (x : α) (k : Nat) : zipIdx #v[x] k = #v[(x, k)] := + rfl + +theorem mk_add_mem_zipIdx_iff_getElem? {k i : Nat} {x : α} {l : Vector α n} : + (x, k + i) ∈ zipIdx l k ↔ l[i]? = some x := by + simp [mem_iff_getElem?, and_left_comm] + +theorem le_snd_of_mem_zipIdx {x : α × Nat} {k : Nat} {l : Vector α n} (h : x ∈ zipIdx l k) : + k ≤ x.2 := + (mk_mem_zipIdx_iff_le_and_getElem?_sub.1 h).1 + +theorem snd_lt_add_of_mem_zipIdx {x : α × Nat} {l : Vector α n} {k : Nat} (h : x ∈ zipIdx l k) : + x.2 < k + n := by + rcases mem_iff_getElem.1 h with ⟨i, h', rfl⟩ + simpa using h' + +theorem snd_lt_of_mem_zipIdx {x : α × Nat} {l : Vector α n} {k : Nat} (h : x ∈ l.zipIdx k) : + x.2 < n + k := by + simpa [Nat.add_comm] using snd_lt_add_of_mem_zipIdx h + +theorem map_zipIdx (f : α → β) (l : Vector α n) (k : Nat) : + map (Prod.map f id) (zipIdx l k) = zipIdx (l.map f) k := by + cases l + simp [Array.map_zipIdx] + +theorem fst_mem_of_mem_zipIdx {x : α × Nat} {l : Vector α n} {k : Nat} (h : x ∈ zipIdx l k) : x.1 ∈ l := + zipIdx_map_fst k l ▸ mem_map_of_mem _ h + +theorem fst_eq_of_mem_zipIdx {x : α × Nat} {l : Vector α n} {k : Nat} (h : x ∈ zipIdx l k) : + x.1 = l[x.2 - k]'(by have := le_snd_of_mem_zipIdx h; have := snd_lt_add_of_mem_zipIdx h; omega) := by + cases l + exact Array.fst_eq_of_mem_zipIdx (by simpa using h) + +theorem mem_zipIdx {x : α} {i : Nat} {xs : Vector α n} {k : Nat} (h : (x, i) ∈ xs.zipIdx k) : + k ≤ i ∧ i < k + n ∧ + x = xs[i - k]'(by have := le_snd_of_mem_zipIdx h; have := snd_lt_add_of_mem_zipIdx h; omega) := + ⟨le_snd_of_mem_zipIdx h, snd_lt_add_of_mem_zipIdx h, fst_eq_of_mem_zipIdx h⟩ + +/-- Variant of `mem_zipIdx` specialized at `k = 0`. -/ +theorem mem_zipIdx' {x : α} {i : Nat} {xs : Vector α n} (h : (x, i) ∈ xs.zipIdx) : + i < n ∧ x = xs[i]'(by have := le_snd_of_mem_zipIdx h; have := snd_lt_add_of_mem_zipIdx h; omega) := + ⟨by simpa using snd_lt_add_of_mem_zipIdx h, fst_eq_of_mem_zipIdx h⟩ + +theorem zipIdx_map (l : Vector α n) (k : Nat) (f : α → β) : + zipIdx (l.map f) k = (zipIdx l k).map (Prod.map f id) := by + cases l + simp [Array.zipIdx_map] + +theorem zipIdx_append (xs : Vector α n) (ys : Vector α m) (k : Nat) : + zipIdx (xs ++ ys) k = zipIdx xs k ++ zipIdx ys (k + n) := by + rcases xs with ⟨xs, rfl⟩ + rcases ys with ⟨ys, rfl⟩ + simp [Array.zipIdx_append] + +theorem zipIdx_eq_append_iff {l : Vector α (n + m)} {k : Nat} : + zipIdx l k = l₁ ++ l₂ ↔ + ∃ (l₁' : Vector α n) (l₂' : Vector α m), + l = l₁' ++ l₂' ∧ l₁ = zipIdx l₁' k ∧ l₂ = zipIdx l₂' (k + n) := by + rcases l with ⟨l, h⟩ + rcases l₁ with ⟨l₁, rfl⟩ + rcases l₂ with ⟨l₂, rfl⟩ + simp only [zipIdx_mk, mk_append_mk, eq_mk, Array.zipIdx_eq_append_iff, mk_eq, toArray_append, + toArray_zipIdx] + constructor + · rintro ⟨l₁', l₂', rfl, rfl, rfl⟩ + exact ⟨⟨l₁', by simp⟩, ⟨l₂', by simp⟩, by simp⟩ + · rintro ⟨⟨l₁', h₁⟩, ⟨l₂', h₂⟩, rfl, w₁, w₂⟩ + exact ⟨l₁', l₂', by simp, w₁, by simp [h₁, w₂]⟩ + +end Vector