2141 lines
72 KiB
Text
2141 lines
72 KiB
Text
/-
|
||
Copyright (c) 2018 Microsoft Corporation. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Leonardo de Moura
|
||
-/
|
||
module
|
||
|
||
prelude
|
||
public import Init.GetElem
|
||
public import Init.Data.List.ToArrayImpl
|
||
import all Init.Data.List.ToArrayImpl
|
||
public import Init.Data.Array.Set
|
||
import all Init.Data.Array.Set
|
||
public import Init.WF
|
||
meta import Init.MetaTypes
|
||
import Init.WFTactics
|
||
|
||
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.
|
||
|
||
universe u v w
|
||
|
||
/-! ### Array literal syntax -/
|
||
|
||
/-- Syntax for `Array α`. -/
|
||
syntax (name := «term#[_,]») "#[" withoutPosition(term,*,?) "]" : term
|
||
|
||
macro_rules
|
||
| `(#[ $elems,* ]) => `(List.toArray [ $elems,* ])
|
||
|
||
recommended_spelling "empty" for "#[]" in [«term#[_,]»]
|
||
recommended_spelling "singleton" for "#[x]" in [«term#[_,]»]
|
||
|
||
variable {α : Type u}
|
||
|
||
namespace Array
|
||
|
||
/-! ### Preliminary theorems -/
|
||
|
||
@[simp, grind =] theorem size_set {xs : Array α} {i : Nat} {v : α} (h : i < xs.size) :
|
||
(set xs i v h).size = xs.size :=
|
||
List.length_set ..
|
||
|
||
@[simp, grind =] theorem size_push {xs : Array α} (v : α) : (push xs v).size = xs.size + 1 :=
|
||
List.length_concat ..
|
||
|
||
theorem ext {xs ys : Array α}
|
||
(h₁ : xs.size = ys.size)
|
||
(h₂ : (i : Nat) → (hi₁ : i < xs.size) → (hi₂ : i < ys.size) → xs[i] = ys[i])
|
||
: xs = ys := by
|
||
let rec extAux (as bs : List α)
|
||
(h₁ : as.length = bs.length)
|
||
(h₂ : (i : Nat) → (hi₁ : i < as.length) → (hi₂ : i < bs.length) → as[i] = bs[i])
|
||
: as = bs := by
|
||
induction as generalizing bs with
|
||
| nil =>
|
||
cases bs with
|
||
| nil => rfl
|
||
| cons b bs => rw [List.length_cons] at h₁; injection h₁
|
||
| cons a as ih =>
|
||
cases bs with
|
||
| nil => rw [List.length_cons] at h₁; injection h₁
|
||
| cons b bs =>
|
||
have hz₁ : 0 < (a::as).length := by rw [List.length_cons]; apply Nat.zero_lt_succ
|
||
have hz₂ : 0 < (b::bs).length := by rw [List.length_cons]; apply Nat.zero_lt_succ
|
||
have headEq : a = b := h₂ 0 hz₁ hz₂
|
||
have h₁' : as.length = bs.length := by rw [List.length_cons, List.length_cons] at h₁; injection h₁
|
||
have h₂' : (i : Nat) → (hi₁ : i < as.length) → (hi₂ : i < bs.length) → as[i] = bs[i] := by
|
||
intro i hi₁ hi₂
|
||
have hi₁' : i+1 < (a::as).length := by rw [List.length_cons]; apply Nat.succ_lt_succ; assumption
|
||
have hi₂' : i+1 < (b::bs).length := by rw [List.length_cons]; apply Nat.succ_lt_succ; assumption
|
||
have : (a::as)[i+1] = (b::bs)[i+1] := h₂ (i+1) hi₁' hi₂'
|
||
apply this
|
||
have tailEq : as = bs := ih bs h₁' h₂'
|
||
rw [headEq, tailEq]
|
||
cases xs; cases ys
|
||
apply congrArg
|
||
apply extAux
|
||
assumption
|
||
assumption
|
||
|
||
theorem ext' {xs ys : Array α} (h : xs.toList = ys.toList) : xs = ys := by
|
||
cases xs; cases ys; simp at h; rw [h]
|
||
|
||
@[simp] theorem toArrayAux_eq {as : List α} {acc : Array α} : (as.toArrayAux acc).toList = acc.toList ++ as := by
|
||
induction as generalizing acc <;> simp [*, List.toArrayAux, Array.push, List.append_assoc, List.concat_eq_append]
|
||
|
||
@[simp, grind =] theorem toArray_toList {xs : Array α} : xs.toList.toArray = xs := rfl
|
||
|
||
@[simp, grind =] theorem getElem_toList {xs : Array α} {i : Nat} (h : i < xs.size) : xs.toList[i] = xs[i] := rfl
|
||
|
||
@[simp, grind =] theorem getElem?_toList {xs : Array α} {i : Nat} : xs.toList[i]? = xs[i]? := by
|
||
simp only [getElem?_def, getElem_toList]
|
||
simp only [Array.size]
|
||
|
||
/-- `a ∈ as` is a predicate which asserts that `a` is in the array `as`. -/
|
||
-- NB: This is defined as a structure rather than a plain def so that a lemma
|
||
-- like `sizeOf_lt_of_mem` will not apply with no actual arrays around.
|
||
structure Mem (as : Array α) (a : α) : Prop where
|
||
val : a ∈ as.toList
|
||
|
||
instance : Membership α (Array α) where
|
||
mem := Mem
|
||
|
||
theorem mem_def {a : α} {as : Array α} : a ∈ as ↔ a ∈ as.toList :=
|
||
⟨fun | .mk h => h, Array.Mem.mk⟩
|
||
|
||
@[simp, grind =] theorem _root_.List.mem_toArray {a : α} {l : List α} : a ∈ l.toArray ↔ a ∈ l := by
|
||
simp [mem_def]
|
||
|
||
@[deprecated List.mem_toArray (since := "2025-09-04")]
|
||
theorem mem_toArray {a : α} {l : List α} : a ∈ l.toArray ↔ a ∈ l :=
|
||
List.mem_toArray
|
||
|
||
@[simp] theorem getElem_mem {xs : Array α} {i : Nat} (h : i < xs.size) : xs[i] ∈ xs := by
|
||
rw [Array.mem_def, ← getElem_toList]
|
||
apply List.getElem_mem
|
||
|
||
grind_pattern getElem_mem => xs[i] ∈ xs
|
||
|
||
@[simp, grind =] theorem emptyWithCapacity_eq {α n} : @emptyWithCapacity α n = #[] := rfl
|
||
|
||
@[simp] theorem mkEmpty_eq {α n} : @mkEmpty α n = #[] := rfl
|
||
|
||
end Array
|
||
|
||
namespace List
|
||
|
||
-- This does not need to be a simp lemma, as already after the `whnfR` the right hand side is `as`.
|
||
theorem toList_toArray {as : List α} : as.toArray.toList = as := rfl
|
||
|
||
@[simp, grind =] theorem size_toArray {as : List α} : as.toArray.size = as.length := by simp [Array.size]
|
||
|
||
@[simp, grind =] theorem getElem_toArray {xs : List α} {i : Nat} (h : i < xs.toArray.size) :
|
||
xs.toArray[i] = xs[i]'(by simpa using h) := rfl
|
||
|
||
@[simp, grind =] theorem getElem?_toArray {xs : List α} {i : Nat} : xs.toArray[i]? = xs[i]? := by
|
||
simp [getElem?_def]
|
||
|
||
@[simp, grind =] theorem getElem!_toArray [Inhabited α] {xs : List α} {i : Nat} :
|
||
xs.toArray[i]! = xs[i]! := by
|
||
simp [getElem!_def]
|
||
|
||
end List
|
||
|
||
namespace Array
|
||
|
||
theorem size_eq_length_toList {xs : Array α} : xs.size = xs.toList.length := rfl
|
||
|
||
/-! ### Externs -/
|
||
|
||
/--
|
||
Returns the size of the array as a platform-native unsigned integer.
|
||
|
||
This is a low-level version of `Array.size` that directly queries the runtime system's
|
||
representation of arrays. While this is not provable, `Array.usize` always returns the exact size of
|
||
the array since the implementation only supports arrays of size less than `USize.size`.
|
||
-/
|
||
@[extern "lean_array_size", simp, expose]
|
||
def usize (xs : @& Array α) : USize := xs.size.toUSize
|
||
|
||
/--
|
||
Low-level indexing operator which is as fast as a C array read.
|
||
|
||
This avoids overhead due to unboxing a `Nat` used as an index.
|
||
-/
|
||
@[extern "lean_array_uget", simp, expose]
|
||
def uget (xs : @& Array α) (i : USize) (h : i.toNat < xs.size) : α :=
|
||
xs[i.toNat]
|
||
|
||
/--
|
||
Low-level modification operator which is as fast as a C array write. The modification is performed
|
||
in-place when the reference to the array is unique.
|
||
|
||
This avoids overhead due to unboxing a `Nat` used as an index.
|
||
-/
|
||
@[extern "lean_array_uset", expose]
|
||
def uset (xs : Array α) (i : USize) (v : α) (h : i.toNat < xs.size) : Array α :=
|
||
xs.set i.toNat v h
|
||
|
||
/--
|
||
Removes the last element of an array. If the array is empty, then it is returned unmodified. The
|
||
modification is performed in-place when the reference to the array is unique.
|
||
|
||
Examples:
|
||
* `#[1, 2, 3].pop = #[1, 2]`
|
||
* `#["orange", "yellow"].pop = #["orange"]`
|
||
* `(#[] : Array String).pop = #[]`
|
||
-/
|
||
@[extern "lean_array_pop", expose]
|
||
def pop (xs : Array α) : Array α where
|
||
toList := xs.toList.dropLast
|
||
|
||
@[simp, grind =] theorem size_pop {xs : Array α} : xs.pop.size = xs.size - 1 := by
|
||
match xs with
|
||
| ⟨[]⟩ => rfl
|
||
| ⟨a::as⟩ => simp [pop, Nat.succ_sub_succ_eq_sub, size]
|
||
|
||
/--
|
||
Creates an array that contains `n` repetitions of `v`.
|
||
|
||
The corresponding `List` function is `List.replicate`.
|
||
|
||
Examples:
|
||
* `Array.replicate 2 true = #[true, true]`
|
||
* `Array.replicate 3 () = #[(), (), ()]`
|
||
* `Array.replicate 0 "anything" = #[]`
|
||
-/
|
||
@[extern "lean_mk_array", expose]
|
||
def replicate {α : Type u} (n : Nat) (v : α) : Array α where
|
||
toList := List.replicate n v
|
||
|
||
/--
|
||
Swaps two elements of an array. The modification is performed in-place when the reference to the
|
||
array is unique.
|
||
|
||
Examples:
|
||
* `#["red", "green", "blue", "brown"].swap 0 3 = #["brown", "green", "blue", "red"]`
|
||
* `#["red", "green", "blue", "brown"].swap 0 2 = #["blue", "green", "red", "brown"]`
|
||
* `#["red", "green", "blue", "brown"].swap 1 2 = #["red", "blue", "green", "brown"]`
|
||
* `#["red", "green", "blue", "brown"].swap 3 0 = #["brown", "green", "blue", "red"]`
|
||
-/
|
||
@[extern "lean_array_fswap", expose]
|
||
def swap (xs : Array α) (i j : @& Nat) (hi : i < xs.size := by get_elem_tactic) (hj : j < xs.size := by get_elem_tactic) : Array α :=
|
||
let v₁ := xs[i]
|
||
let v₂ := xs[j]
|
||
let xs' := xs.set i v₂
|
||
xs'.set j v₁ (Nat.lt_of_lt_of_eq hj (size_set _).symm)
|
||
|
||
@[simp, grind =] theorem size_swap {xs : Array α} {i j : Nat} {hi hj} : (xs.swap i j hi hj).size = xs.size := by
|
||
change ((xs.set i xs[j]).set j xs[i]
|
||
(Nat.lt_of_lt_of_eq hj (size_set _).symm)).size = xs.size
|
||
rw [size_set, size_set]
|
||
|
||
/--
|
||
Swaps two elements of an array, returning the array unchanged if either index is out of bounds. The
|
||
modification is performed in-place when the reference to the array is unique.
|
||
|
||
Examples:
|
||
* `#["red", "green", "blue", "brown"].swapIfInBounds 0 3 = #["brown", "green", "blue", "red"]`
|
||
* `#["red", "green", "blue", "brown"].swapIfInBounds 0 2 = #["blue", "green", "red", "brown"]`
|
||
* `#["red", "green", "blue", "brown"].swapIfInBounds 1 2 = #["red", "blue", "green", "brown"]`
|
||
* `#["red", "green", "blue", "brown"].swapIfInBounds 0 4 = #["red", "green", "blue", "brown"]`
|
||
* `#["red", "green", "blue", "brown"].swapIfInBounds 9 2 = #["red", "green", "blue", "brown"]`
|
||
-/
|
||
@[extern "lean_array_swap", expose]
|
||
def swapIfInBounds (xs : Array α) (i j : @& Nat) : Array α :=
|
||
if h₁ : i < xs.size then
|
||
if h₂ : j < xs.size then swap xs i j
|
||
else xs
|
||
else xs
|
||
|
||
/-! ### GetElem instance for `USize`, backed by `uget` -/
|
||
|
||
instance : GetElem (Array α) USize α fun xs i => i.toNat < xs.size where
|
||
getElem xs i h := xs.uget i h
|
||
|
||
/-! ### Definitions -/
|
||
|
||
instance : EmptyCollection (Array α) := ⟨Array.empty⟩
|
||
instance : Inhabited (Array α) where
|
||
default := Array.empty
|
||
|
||
/--
|
||
Checks whether an array is empty.
|
||
|
||
An array is empty if its size is `0`.
|
||
|
||
Examples:
|
||
* `(#[] : Array String).isEmpty = true`
|
||
* `#[1, 2].isEmpty = false`
|
||
* `#[()].isEmpty = false`
|
||
-/
|
||
@[expose]
|
||
def isEmpty (xs : Array α) : Bool :=
|
||
xs.size = 0
|
||
|
||
@[specialize]
|
||
def isEqvAux (xs ys : Array α) (hsz : xs.size = ys.size) (p : α → α → Bool) :
|
||
∀ (i : Nat) (_ : i ≤ xs.size), Bool
|
||
| 0, _ => true
|
||
| i+1, h =>
|
||
p xs[i] (ys[i]'(hsz ▸ h)) && isEqvAux xs ys hsz p i (Nat.le_trans (Nat.le_add_right i 1) h)
|
||
|
||
/--
|
||
Returns `true` if `as` and `bs` have the same length and they are pairwise related by `eqv`.
|
||
|
||
Short-circuits at the first non-related pair of elements.
|
||
|
||
Examples:
|
||
* `#[1, 2, 3].isEqv #[2, 3, 4] (· < ·) = true`
|
||
* `#[1, 2, 3].isEqv #[2, 2, 4] (· < ·) = false`
|
||
* `#[1, 2, 3].isEqv #[2, 3] (· < ·) = false`
|
||
-/
|
||
@[inline] def isEqv (xs ys : Array α) (p : α → α → Bool) : Bool :=
|
||
if h : xs.size = ys.size then
|
||
isEqvAux xs ys h p xs.size (Nat.le_refl xs.size)
|
||
else
|
||
false
|
||
|
||
instance [BEq α] : BEq (Array α) :=
|
||
⟨fun xs ys => isEqv xs ys BEq.beq⟩
|
||
|
||
/-
|
||
`ofFn f` with `f : Fin n → α` returns the list whose ith element is `f i`.
|
||
```
|
||
ofFn f = #[f 0, f 1, ... , f(n - 1)]
|
||
``` -/
|
||
/--
|
||
Creates an array by applying `f` to each potential index in order, starting at `0`.
|
||
|
||
Examples:
|
||
* `Array.ofFn (n := 3) toString = #["0", "1", "2"]`
|
||
* `Array.ofFn (fun i => #["red", "green", "blue"].get i.val i.isLt) = #["red", "green", "blue"]`
|
||
-/
|
||
def ofFn {n} (f : Fin n → α) : Array α := go (emptyWithCapacity n) n (Nat.le_refl n) where
|
||
/-- Auxiliary for `ofFn`. `ofFn.go f acc i h = acc ++ #[f (n - i), ..., f(n - 1)]` -/
|
||
go (acc : Array α) : (i : Nat) → i ≤ n → Array α
|
||
| i + 1, h =>
|
||
have w : n - i - 1 < n :=
|
||
Nat.lt_of_lt_of_le (Nat.sub_one_lt (Nat.sub_ne_zero_iff_lt.mpr h)) (Nat.sub_le n i)
|
||
go (acc.push (f ⟨n - i - 1, w⟩)) i (Nat.le_of_succ_le h)
|
||
| 0, _ => acc
|
||
|
||
-- See also `Array.ofFnM` defined in `Init.Data.Array.OfFn`.
|
||
|
||
/--
|
||
Constructs an array that contains all the numbers from `0` to `n`, exclusive.
|
||
|
||
Examples:
|
||
* `Array.range 5 := #[0, 1, 2, 3, 4]`
|
||
* `Array.range 0 := #[]`
|
||
* `Array.range 1 := #[0]`
|
||
-/
|
||
def range (n : Nat) : Array Nat :=
|
||
ofFn fun (i : Fin n) => i
|
||
|
||
/--
|
||
Constructs an array of numbers of size `size`, starting at `start` and increasing by
|
||
`step` at each element.
|
||
|
||
In other words, `Array.range' start size step` is `#[start, start+step, ..., start+(len-1)*step]`.
|
||
|
||
Examples:
|
||
* `Array.range' 0 3 (step := 1) = #[0, 1, 2]`
|
||
* `Array.range' 0 3 (step := 2) = #[0, 2, 4]`
|
||
* `Array.range' 0 4 (step := 2) = #[0, 2, 4, 6]`
|
||
* `Array.range' 3 4 (step := 2) = #[3, 5, 7, 9]`
|
||
-/
|
||
def range' (start size : Nat) (step : Nat := 1) : Array Nat :=
|
||
ofFn fun (i : Fin size) => start + step * i
|
||
|
||
/--
|
||
Constructs a single-element array that contains `v`.
|
||
|
||
Examples:
|
||
* `Array.singleton 5 = #[5]`
|
||
* `Array.singleton "one" = #["one"]`
|
||
-/
|
||
@[inline, expose] protected def singleton (v : α) : Array α := #[v]
|
||
|
||
/--
|
||
Returns the last element of an array, or panics if the array is empty.
|
||
|
||
Safer alternatives include `Array.back`, which requires a proof the array is non-empty, and
|
||
`Array.back?`, which returns an `Option`.
|
||
-/
|
||
def back! [Inhabited α] (xs : Array α) : α :=
|
||
xs[xs.size - 1]!
|
||
|
||
/--
|
||
Returns the last element of an array, given a proof that the array is not empty.
|
||
|
||
See `Array.back!` for the version that panics if the array is empty, or `Array.back?` for the
|
||
version that returns an option.
|
||
-/
|
||
def back (xs : Array α) (h : 0 < xs.size := by get_elem_tactic) : α :=
|
||
xs[xs.size - 1]'(Nat.sub_one_lt_of_lt h)
|
||
|
||
/--
|
||
Returns the last element of an array, or `none` if the array is empty.
|
||
|
||
See `Array.back!` for the version that panics if the array is empty, or `Array.back` for the version
|
||
that requires a proof the array is non-empty.
|
||
-/
|
||
def back? (xs : Array α) : Option α :=
|
||
xs[xs.size - 1]?
|
||
|
||
/--
|
||
Swaps a new element with the element at the given index.
|
||
|
||
Returns the value formerly found at `i`, paired with an array in which the value at `i` has been
|
||
replaced with `v`.
|
||
|
||
Examples:
|
||
* `#["spinach", "broccoli", "carrot"].swapAt 1 "pepper" = ("broccoli", #["spinach", "pepper", "carrot"])`
|
||
* `#["spinach", "broccoli", "carrot"].swapAt 2 "pepper" = ("carrot", #["spinach", "broccoli", "pepper"])`
|
||
-/
|
||
@[inline, expose] def swapAt (xs : Array α) (i : Nat) (v : α) (hi : i < xs.size := by get_elem_tactic) : α × Array α :=
|
||
let e := xs[i]
|
||
let xs' := xs.set i v
|
||
(e, xs')
|
||
|
||
/--
|
||
Swaps a new element with the element at the given index. Panics if the index is out of bounds.
|
||
|
||
Returns the value formerly found at `i`, paired with an array in which the value at `i` has been
|
||
replaced with `v`.
|
||
|
||
Examples:
|
||
* `#["spinach", "broccoli", "carrot"].swapAt! 1 "pepper" = (#["spinach", "pepper", "carrot"], "broccoli")`
|
||
* `#["spinach", "broccoli", "carrot"].swapAt! 2 "pepper" = (#["spinach", "broccoli", "pepper"], "carrot")`
|
||
-/
|
||
@[inline, expose]
|
||
def swapAt! (xs : Array α) (i : Nat) (v : α) : α × Array α :=
|
||
if h : i < xs.size then
|
||
swapAt xs i v
|
||
else
|
||
have : Inhabited (α × Array α) := ⟨(v, xs)⟩
|
||
panic! String.Internal.append (String.Internal.append "index " (toString i)) " out of bounds"
|
||
|
||
/--
|
||
Returns the first `n` elements of an array. The resulting array is produced by repeatedly calling
|
||
`Array.pop`. If `n` is greater than the size of the array, it is returned unmodified.
|
||
|
||
If the reference to the array is unique, then this function uses in-place modification.
|
||
|
||
Examples:
|
||
* `#[0, 1, 2, 3, 4].shrink 2 = #[0, 1]`
|
||
* `#[0, 1, 2, 3, 4].shrink 0 = #[]`
|
||
* `#[0, 1, 2, 3, 4].shrink 10 = #[0, 1, 2, 3, 4]`
|
||
-/
|
||
def shrink (xs : Array α) (n : Nat) : Array α :=
|
||
let rec loop
|
||
| 0, xs => xs
|
||
| n+1, xs => loop n xs.pop
|
||
loop (xs.size - n) xs
|
||
|
||
/--
|
||
Returns a new array that contains the first `i` elements of `xs`. If `xs` has fewer than `i`
|
||
elements, the new array contains all the elements of `xs`.
|
||
|
||
The returned array is always a new array, even if it contains the same elements as the input array.
|
||
|
||
Examples:
|
||
* `#["red", "green", "blue"].take 1 = #["red"]`
|
||
* `#["red", "green", "blue"].take 2 = #["red", "green"]`
|
||
* `#["red", "green", "blue"].take 5 = #["red", "green", "blue"]`
|
||
-/
|
||
abbrev take (xs : Array α) (i : Nat) : Array α := extract xs 0 i
|
||
|
||
@[simp, grind =] theorem take_eq_extract {xs : Array α} {i : Nat} : xs.take i = xs.extract 0 i := rfl
|
||
|
||
/--
|
||
Removes the first `i` elements of `xs`. If `xs` has fewer than `i` elements, the new array is empty.
|
||
|
||
The returned array is always a new array, even if it contains the same elements as the input array.
|
||
|
||
Examples:
|
||
* `#["red", "green", "blue"].drop 1 = #["green", "blue"]`
|
||
* `#["red", "green", "blue"].drop 2 = #["blue"]`
|
||
* `#["red", "green", "blue"].drop 5 = #[]`
|
||
-/
|
||
abbrev drop (xs : Array α) (i : Nat) : Array α := extract xs i xs.size
|
||
|
||
@[simp, grind =] theorem drop_eq_extract {xs : Array α} {i : Nat} : xs.drop i = xs.extract i xs.size := rfl
|
||
|
||
@[inline]
|
||
unsafe def modifyMUnsafe [Monad m] (xs : Array α) (i : Nat) (f : α → m α) : m (Array α) := do
|
||
if h : i < xs.size then
|
||
let v := xs[i]
|
||
-- Replace a[i] by `box(0)`. This ensures that `v` remains unshared if possible.
|
||
-- Note: we assume that arrays have a uniform representation irrespective
|
||
-- of the element type, and that it is valid to store `box(0)` in any array.
|
||
let xs' := xs.set i (unsafeCast ())
|
||
let v ← f v
|
||
pure <| xs'.set i v (Nat.lt_of_lt_of_eq h (size_set ..).symm)
|
||
else
|
||
pure xs
|
||
|
||
/--
|
||
Replaces the element at the given index, if it exists, with the result of applying the monadic
|
||
function `f` to it. If the index is invalid, the array is returned unmodified and `f` is not called.
|
||
|
||
Examples:
|
||
```lean example
|
||
#eval #[1, 2, 3, 4].modifyM 2 fun x => do
|
||
IO.println s!"It was {x}"
|
||
return x * 10
|
||
```
|
||
```output
|
||
It was 3
|
||
```
|
||
```output
|
||
#[1, 2, 30, 4]
|
||
```
|
||
|
||
```lean example
|
||
#eval #[1, 2, 3, 4].modifyM 6 fun x => do
|
||
IO.println s!"It was {x}"
|
||
return x * 10
|
||
```
|
||
```output
|
||
#[1, 2, 3, 4]
|
||
```
|
||
-/
|
||
@[implemented_by modifyMUnsafe]
|
||
def modifyM [Monad m] (xs : Array α) (i : Nat) (f : α → m α) : m (Array α) := do
|
||
if h : i < xs.size then
|
||
let v := xs[i]
|
||
let v ← f v
|
||
pure <| xs.set i v
|
||
else
|
||
pure xs
|
||
|
||
/--
|
||
Replaces the element at the given index, if it exists, with the result of applying `f` to it. If the
|
||
index is invalid, the array is returned unmodified.
|
||
|
||
Examples:
|
||
* `#[1, 2, 3].modify 0 (· * 10) = #[10, 2, 3]`
|
||
* `#[1, 2, 3].modify 2 (· * 10) = #[1, 2, 30]`
|
||
* `#[1, 2, 3].modify 3 (· * 10) = #[1, 2, 3]`
|
||
-/
|
||
@[inline]
|
||
def modify (xs : Array α) (i : Nat) (f : α → α) : Array α :=
|
||
Id.run <| modifyM xs i (pure <| f ·)
|
||
|
||
set_option linter.indexVariables false in -- Changing `idx` causes bootstrapping issues, haven't investigated.
|
||
/--
|
||
Replaces the element at the given index, if it exists, with the result of applying `f` to it. If the
|
||
index is invalid, the array is returned unmodified.
|
||
|
||
Examples:
|
||
* `#[1, 2, 3].modifyOp 0 (· * 10) = #[10, 2, 3]`
|
||
* `#[1, 2, 3].modifyOp 2 (· * 10) = #[1, 2, 30]`
|
||
* `#[1, 2, 3].modifyOp 3 (· * 10) = #[1, 2, 3]`
|
||
-/
|
||
@[inline]
|
||
def modifyOp (xs : Array α) (idx : Nat) (f : α → α) : Array α :=
|
||
xs.modify idx f
|
||
|
||
/--
|
||
We claim this unsafe implementation is correct because an array cannot have more than `usizeSz` elements in our runtime.
|
||
|
||
This kind of low level trick can be removed with a little bit of compiler support. For example, if the compiler simplifies `as.size < usizeSz` to true. -/
|
||
@[inline] unsafe def forIn'Unsafe {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (b : β) (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β :=
|
||
let sz := as.usize
|
||
let rec @[specialize] loop (i : USize) (b : β) : m β := do
|
||
if i < sz then
|
||
let a := as.uget i lcProof
|
||
match (← f a lcProof b) with
|
||
| ForInStep.done b => pure b
|
||
| ForInStep.yield b => loop (i+1) b
|
||
else
|
||
pure b
|
||
loop 0 b
|
||
|
||
/-- Reference implementation for `forIn'` -/
|
||
@[implemented_by Array.forIn'Unsafe, expose]
|
||
protected def forIn' {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (b : β) (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β :=
|
||
let rec loop (i : Nat) (h : i ≤ as.size) (b : β) : m β := do
|
||
match i, h with
|
||
| 0, _ => pure b
|
||
| i+1, h =>
|
||
have h' : i < as.size := Nat.lt_of_lt_of_le (Nat.lt_succ_self i) h
|
||
have : as.size - 1 < as.size := Nat.sub_lt (Nat.zero_lt_of_lt h') (by decide)
|
||
have : as.size - 1 - i < as.size := Nat.lt_of_le_of_lt (Nat.sub_le (as.size - 1) i) this
|
||
match (← f as[as.size - 1 - i] (getElem_mem this) b) with
|
||
| ForInStep.done b => pure b
|
||
| ForInStep.yield b => loop i (Nat.le_of_lt h') b
|
||
loop as.size (Nat.le_refl _) b
|
||
|
||
instance [Monad m] : ForIn' m (Array α) α inferInstance where
|
||
forIn' := Array.forIn'
|
||
|
||
-- No separate `ForIn` instance is required because it can be derived from `ForIn'`.
|
||
|
||
-- We simplify `Array.forIn'` to `forIn'`.
|
||
@[simp] theorem forIn'_eq_forIn' [Monad m] : @Array.forIn' α β m _ = forIn' := rfl
|
||
|
||
/-- See comment at `forIn'Unsafe` -/
|
||
@[inline]
|
||
unsafe def foldlMUnsafe {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : β → α → m β) (init : β) (as : Array α) (start := 0) (stop := as.size) : m β :=
|
||
let rec @[specialize] fold (i : USize) (stop : USize) (b : β) : m β := do
|
||
if i == stop then
|
||
pure b
|
||
else
|
||
fold (i+1) stop (← f b (as.uget i lcProof))
|
||
if start < stop then
|
||
if stop ≤ as.size then
|
||
fold (USize.ofNat start) (USize.ofNat stop) init
|
||
else if start < as.size then
|
||
fold (USize.ofNat start) (USize.ofNat as.size) init
|
||
else
|
||
pure init
|
||
else
|
||
pure init
|
||
|
||
/--
|
||
Folds a monadic function over a list from the left, accumulating a value starting with `init`. The
|
||
accumulated value is combined with the each element of the list in order, using `f`.
|
||
|
||
The optional parameters `start` and `stop` control the region of the array to be folded. Folding
|
||
proceeds from `start` (inclusive) to `stop` (exclusive), so no folding occurs unless `start < stop`.
|
||
By default, the entire array is folded.
|
||
|
||
Examples:
|
||
```lean example
|
||
example [Monad m] (f : α → β → m α) :
|
||
Array.foldlM (m := m) f x₀ #[a, b, c] = (do
|
||
let x₁ ← f x₀ a
|
||
let x₂ ← f x₁ b
|
||
let x₃ ← f x₂ c
|
||
pure x₃)
|
||
:= by rfl
|
||
```
|
||
|
||
```lean example
|
||
example [Monad m] (f : α → β → m α) :
|
||
Array.foldlM (m := m) f x₀ #[a, b, c] (start := 1) = (do
|
||
let x₁ ← f x₀ b
|
||
let x₂ ← f x₁ c
|
||
pure x₂)
|
||
:= by rfl
|
||
```
|
||
-/
|
||
-- Reference implementation for `foldlM`
|
||
@[implemented_by foldlMUnsafe, expose]
|
||
def foldlM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : β → α → m β) (init : β) (as : Array α) (start := 0) (stop := as.size) : m β :=
|
||
let fold (stop : Nat) (h : stop ≤ as.size) :=
|
||
let rec loop (i : Nat) (j : Nat) (b : β) : m β := do
|
||
if hlt : j < stop then
|
||
match i with
|
||
| 0 => pure b
|
||
| i'+1 =>
|
||
have : j < as.size := Nat.lt_of_lt_of_le hlt h
|
||
loop i' (j+1) (← f b as[j])
|
||
else
|
||
pure b
|
||
loop (stop - start) start init
|
||
if h : stop ≤ as.size then
|
||
fold stop h
|
||
else
|
||
fold as.size (Nat.le_refl _)
|
||
|
||
/-- See comment at `forIn'Unsafe` -/
|
||
@[inline]
|
||
unsafe def foldrMUnsafe {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → β → m β) (init : β) (as : Array α) (start := as.size) (stop := 0) : m β :=
|
||
let rec @[specialize] fold (i : USize) (stop : USize) (b : β) : m β := do
|
||
if i == stop then
|
||
pure b
|
||
else
|
||
fold (i-1) stop (← f (as.uget (i-1) lcProof) b)
|
||
if start ≤ as.size then
|
||
if stop < start then
|
||
fold (USize.ofNat start) (USize.ofNat stop) init
|
||
else
|
||
pure init
|
||
else if stop < as.size then
|
||
fold (USize.ofNat as.size) (USize.ofNat stop) init
|
||
else
|
||
pure init
|
||
|
||
/--
|
||
Folds a monadic function over an array from the right, accumulating a value starting with `init`.
|
||
The accumulated value is combined with the each element of the list in reverse order, using `f`.
|
||
|
||
The optional parameters `start` and `stop` control the region of the array to be folded. Folding
|
||
proceeds from `start` (exclusive) to `stop` (inclusive), so no folding occurs unless `start > stop`.
|
||
By default, the entire array is folded.
|
||
|
||
Examples:
|
||
```lean example
|
||
example [Monad m] (f : α → β → m β) :
|
||
Array.foldrM (m := m) f x₀ #[a, b, c] = (do
|
||
let x₁ ← f c x₀
|
||
let x₂ ← f b x₁
|
||
let x₃ ← f a x₂
|
||
pure x₃)
|
||
:= by rfl
|
||
```
|
||
|
||
```lean example
|
||
example [Monad m] (f : α → β → m β) :
|
||
Array.foldrM (m := m) f x₀ #[a, b, c] (start := 2) = (do
|
||
let x₁ ← f b x₀
|
||
let x₂ ← f a x₁
|
||
pure x₂)
|
||
:= by rfl
|
||
```
|
||
-/
|
||
-- Reference implementation for `foldrM`
|
||
@[implemented_by foldrMUnsafe, expose]
|
||
def foldrM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → β → m β) (init : β) (as : Array α) (start := as.size) (stop := 0) : m β :=
|
||
let rec fold (i : Nat) (h : i ≤ as.size) (b : β) : m β := do
|
||
if i == stop then
|
||
pure b
|
||
else match i, h with
|
||
| 0, _ => pure b
|
||
| i+1, h =>
|
||
have : i < as.size := Nat.lt_of_lt_of_le (Nat.lt_succ_self _) h
|
||
fold i (Nat.le_of_lt this) (← f as[i] b)
|
||
if h : start ≤ as.size then
|
||
if stop < start then
|
||
fold start h init
|
||
else
|
||
pure init
|
||
else if stop < as.size then
|
||
fold as.size (Nat.le_refl _) init
|
||
else
|
||
pure init
|
||
|
||
/-- See comment at `forIn'Unsafe` -/
|
||
@[inline]
|
||
unsafe def mapMUnsafe {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → m β) (as : Array α) : m (Array β) :=
|
||
let sz := as.usize
|
||
let rec @[specialize] map (i : USize) (bs : Array NonScalar) : m (Array PNonScalar.{v}) := do
|
||
if i < sz then
|
||
let v := bs.uget i lcProof
|
||
-- Replace bs[i] by `box(0)`. This ensures that `v` remains unshared if possible.
|
||
-- Note: we assume that arrays have a uniform representation irrespective
|
||
-- of the element type, and that it is valid to store `box(0)` in any array.
|
||
let bs' := bs.uset i default lcProof
|
||
let vNew ← f (unsafeCast v)
|
||
map (i+1) (bs'.uset i (unsafeCast vNew) lcProof)
|
||
else
|
||
pure (unsafeCast bs)
|
||
unsafeCast <| map 0 (unsafeCast as)
|
||
|
||
/--
|
||
Applies the monadic action `f` to every element in the array, left-to-right, and returns the array
|
||
of results.
|
||
-/
|
||
-- Reference implementation for `mapM`
|
||
@[implemented_by mapMUnsafe]
|
||
def mapM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → m β) (as : Array α) : m (Array β) :=
|
||
-- Note: we cannot use `foldlM` here for the reference implementation because this calls
|
||
-- `bind` and `pure` too many times. (We are not assuming `m` is a `LawfulMonad`)
|
||
let rec map (i : Nat) (bs : Array β) : m (Array β) := do
|
||
if hlt : i < as.size then
|
||
map (i+1) (bs.push (← f as[i]))
|
||
else
|
||
pure bs
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
map 0 (emptyWithCapacity as.size)
|
||
|
||
/--
|
||
Applies the monadic action `f` to every element in the array, along with the element's index and a
|
||
proof that the index is in bounds, from left to right. Returns the array of results.
|
||
-/
|
||
@[inline, expose]
|
||
def mapFinIdxM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m]
|
||
(as : Array α) (f : (i : Nat) → α → (h : i < 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
|
||
| i+1, inv =>
|
||
have j_lt : j < as.size := by
|
||
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[j] j_lt))
|
||
map as.size 0 rfl (emptyWithCapacity as.size)
|
||
|
||
/--
|
||
Applies the monadic action `f` to every element in the array, along with the element's index, from
|
||
left to right. Returns the array of results.
|
||
-/
|
||
@[inline, expose]
|
||
def mapIdxM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : Nat → α → m β) (as : Array α) : m (Array β) :=
|
||
as.mapFinIdxM fun i a _ => f i a
|
||
|
||
/--
|
||
Maps `f` over the array and collects the results with `<|>`. The result for the end of the array is
|
||
`failure`.
|
||
|
||
Examples:
|
||
* `#[[], [1, 2], [], [2]].firstM List.head? = some 1`
|
||
* `#[[], [], []].firstM List.head? = none`
|
||
* `#[].firstM List.head? = none`
|
||
-/
|
||
@[inline]
|
||
def firstM {α : Type u} {m : Type v → Type w} [Alternative m] (f : α → m β) (as : Array α) : m β :=
|
||
go 0
|
||
where
|
||
go (i : Nat) : m β :=
|
||
if hlt : i < as.size then
|
||
f as[i] <|> go (i+1)
|
||
else
|
||
failure
|
||
termination_by as.size - i
|
||
decreasing_by exact Nat.sub_succ_lt_self as.size i hlt
|
||
|
||
/--
|
||
Returns the first non-`none` result of applying the monadic function `f` to each element of the
|
||
array, in order. Returns `none` if `f` returns `none` for all elements.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #[7, 6, 5, 8, 1, 2, 6].findSomeM? fun i => do
|
||
if i < 5 then
|
||
return some (i * 10)
|
||
if i ≤ 6 then
|
||
IO.println s!"Almost! {i}"
|
||
return none
|
||
```
|
||
```output
|
||
Almost! 6
|
||
Almost! 5
|
||
```
|
||
```output
|
||
some 10
|
||
```
|
||
-/
|
||
@[inline, expose]
|
||
def findSomeM? {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → m (Option β)) (as : Array α) : m (Option β) := do
|
||
for a in as do
|
||
match (← f a) with
|
||
| some b => return some b
|
||
| _ => pure ⟨⟩
|
||
return none
|
||
|
||
/--
|
||
Returns the first element of the array for which the monadic predicate `p` returns `true`, or `none`
|
||
if no such element is found. Elements of the array are checked in order.
|
||
|
||
The monad `m` is restricted to `Type → Type` to avoid needing to use `ULift Bool` in `p`'s type.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #[7, 6, 5, 8, 1, 2, 6].findM? fun i => do
|
||
if i < 5 then
|
||
return true
|
||
if i ≤ 6 then
|
||
IO.println s!"Almost! {i}"
|
||
return false
|
||
```
|
||
```output
|
||
Almost! 6
|
||
Almost! 5
|
||
```
|
||
```output
|
||
some 1
|
||
```
|
||
-/
|
||
@[inline]
|
||
def findM? {α : Type} [Monad m] (p : α → m Bool) (as : Array α) : m (Option α) := do
|
||
for a in as do
|
||
if (← p a) then
|
||
return some a
|
||
return none
|
||
|
||
/--
|
||
Finds the index of the first element of an array for which the monadic predicate `p` returns `true`.
|
||
Elements are examined in order from left to right, and the search is terminated when an element that
|
||
satisfies `p` is found. If no such element exists in the array, then `none` is returned.
|
||
-/
|
||
@[inline]
|
||
def findIdxM? [Monad m] (p : α → m Bool) (as : Array α) : m (Option Nat) := do
|
||
let mut i := 0
|
||
for a in as do
|
||
if (← p a) then
|
||
return some i
|
||
i := i + 1
|
||
return none
|
||
|
||
@[inline]
|
||
unsafe def anyMUnsafe {α : Type u} {m : Type → Type w} [Monad m] (p : α → m Bool) (as : Array α) (start := 0) (stop := as.size) : m Bool :=
|
||
let rec @[specialize] any (i : USize) (stop : USize) : m Bool := do
|
||
if i == stop then
|
||
pure false
|
||
else
|
||
if (← p (as.uget i lcProof)) then
|
||
pure true
|
||
else
|
||
any (i+1) stop
|
||
if start < stop then
|
||
let stop' := min stop as.size
|
||
if start < stop' then
|
||
any (USize.ofNat start) (USize.ofNat stop')
|
||
else
|
||
pure false
|
||
else
|
||
pure false
|
||
|
||
/--
|
||
Returns `true` if the monadic predicate `p` returns `true` for any element of `as`.
|
||
|
||
Short-circuits upon encountering the first `true`. The elements in `as` are examined in order from
|
||
left to right.
|
||
|
||
The optional parameters `start` and `stop` control the region of the array to be checked. Only the
|
||
elements with indices from `start` (inclusive) to `stop` (exclusive) are checked. By default, the
|
||
entire array is checked.
|
||
-/
|
||
@[implemented_by anyMUnsafe, expose]
|
||
def anyM {α : Type u} {m : Type → Type w} [Monad m] (p : α → m Bool) (as : Array α) (start := 0) (stop := as.size) : m Bool :=
|
||
let any (stop : Nat) (h : stop ≤ as.size) :=
|
||
let rec loop (j : Nat) : m Bool := do
|
||
if hlt : j < stop then
|
||
have : j < as.size := Nat.lt_of_lt_of_le hlt h
|
||
if (← p as[j]) then
|
||
pure true
|
||
else
|
||
loop (j+1)
|
||
else
|
||
pure false
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
loop start
|
||
if h : stop ≤ as.size then
|
||
any stop h
|
||
else
|
||
any as.size (Nat.le_refl _)
|
||
|
||
/--
|
||
Returns `true` if the monadic predicate `p` returns `true` for every element of `as`.
|
||
|
||
Short-circuits upon encountering the first `false`. The elements in `as` are examined in order from
|
||
left to right.
|
||
|
||
The optional parameters `start` and `stop` control the region of the array to be checked. Only the
|
||
elements with indices from `start` (inclusive) to `stop` (exclusive) are checked. By default, the
|
||
entire array is checked.
|
||
-/
|
||
@[inline]
|
||
def allM {α : Type u} {m : Type → Type w} [Monad m] (p : α → m Bool) (as : Array α) (start := 0) (stop := as.size) : m Bool :=
|
||
return !(← as.anyM (start := start) (stop := stop) fun v => return !(← p v))
|
||
|
||
/--
|
||
Returns the first non-`none` result of applying the monadic function `f` to each element of the
|
||
array in reverse order, from right to left. Once a non-`none` result is found, no further elements
|
||
are checked. Returns `none` if `f` returns `none` for all elements of the array.
|
||
|
||
Examples:
|
||
```lean example
|
||
#eval #[1, 2, 0, -4, 1].findSomeRevM? (m := Except String) fun x => do
|
||
if x = 0 then throw "Zero!"
|
||
else if x < 0 then return (some x)
|
||
else return none
|
||
```
|
||
```output
|
||
Except.ok (some (-4))
|
||
```
|
||
```lean example
|
||
#eval #[1, 2, 0, 4, 1].findSomeRevM? (m := Except String) fun x => do
|
||
if x = 0 then throw "Zero!"
|
||
else if x < 0 then return (some x)
|
||
else return none
|
||
```
|
||
```output
|
||
Except.error "Zero!"
|
||
```
|
||
-/
|
||
@[inline]
|
||
def findSomeRevM? {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → m (Option β)) (as : Array α) : m (Option β) :=
|
||
let rec @[specialize] find : (i : Nat) → i ≤ as.size → m (Option β)
|
||
| 0, _ => pure none
|
||
| i+1, h => do
|
||
have : i < as.size := Nat.lt_of_lt_of_le (Nat.lt_succ_self _) h
|
||
let r ← f as[i]
|
||
match r with
|
||
| some _ => pure r
|
||
| none =>
|
||
have : i ≤ as.size := Nat.le_of_lt this
|
||
find i this
|
||
find as.size (Nat.le_refl _)
|
||
|
||
/--
|
||
Returns the last element of the array for which the monadic predicate `p` returns `true`, or `none`
|
||
if no such element is found. Elements of the array are checked in reverse, from right to left..
|
||
|
||
The monad `m` is restricted to `Type → Type` to avoid needing to use `ULift Bool` in `p`'s type.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #[7, 5, 8, 1, 2, 6, 5, 8].findRevM? fun i => do
|
||
if i < 5 then
|
||
return true
|
||
if i ≤ 6 then
|
||
IO.println s!"Almost! {i}"
|
||
return false
|
||
```
|
||
```output
|
||
Almost! 5
|
||
Almost! 6
|
||
```
|
||
```output
|
||
some 2
|
||
```
|
||
-/
|
||
@[inline]
|
||
def findRevM? {α : Type} {m : Type → Type w} [Monad m] (p : α → m Bool) (as : Array α) : m (Option α) :=
|
||
as.findSomeRevM? fun a => return if (← p a) then some a else none
|
||
|
||
/--
|
||
Applies the monadic action `f` to each element of an array, in order.
|
||
|
||
The optional parameters `start` and `stop` control the region of the array to which `f` should be
|
||
applied. Iteration proceeds from `start` (inclusive) to `stop` (exclusive), so `f` is not invoked
|
||
unless `start < stop`. By default, the entire array is used.
|
||
-/
|
||
@[inline, expose]
|
||
protected def forM {α : Type u} {m : Type v → Type w} [Monad m] (f : α → m PUnit) (as : Array α) (start := 0) (stop := as.size) : m PUnit :=
|
||
as.foldlM (fun _ => f) ⟨⟩ start stop
|
||
|
||
instance [Monad m] : ForM m (Array α) α where
|
||
forM xs f := Array.forM f xs
|
||
|
||
-- We simplify `Array.forM` to `forM`.
|
||
@[simp] theorem forM_eq_forM [Monad m] {f : α → m PUnit} :
|
||
Array.forM f as 0 as.size = forM as f := rfl
|
||
|
||
/--
|
||
Applies the monadic action `f` to each element of an array from right to left, in reverse order.
|
||
|
||
The optional parameters `start` and `stop` control the region of the array to which `f` should be
|
||
applied. Iteration proceeds from `start` (exclusive) to `stop` (inclusive), so no `f` is not invoked
|
||
unless `start > stop`. By default, the entire array is used.
|
||
-/
|
||
@[inline]
|
||
def forRevM {α : Type u} {m : Type v → Type w} [Monad m] (f : α → m PUnit) (as : Array α) (start := as.size) (stop := 0) : m PUnit :=
|
||
as.foldrM (fun a _ => f a) ⟨⟩ start stop
|
||
|
||
/--
|
||
Folds a function over an array from the left, accumulating a value starting with `init`. The
|
||
accumulated value is combined with the each element of the array in order, using `f`.
|
||
|
||
The optional parameters `start` and `stop` control the region of the array to be folded. Folding
|
||
proceeds from `start` (inclusive) to `stop` (exclusive), so no folding occurs unless `start < stop`.
|
||
By default, the entire array is used.
|
||
|
||
Examples:
|
||
* `#[a, b, c].foldl f z = f (f (f z a) b) c`
|
||
* `#[1, 2, 3].foldl (· ++ toString ·) "" = "123"`
|
||
* `#[1, 2, 3].foldl (s!"({·} {·})") "" = "((( 1) 2) 3)"`
|
||
-/
|
||
@[inline, expose]
|
||
def foldl {α : Type u} {β : Type v} (f : β → α → β) (init : β) (as : Array α) (start := 0) (stop := as.size) : β :=
|
||
Id.run <| as.foldlM (pure <| f · ·) init start stop
|
||
|
||
/--
|
||
Folds a function over an array from the right, accumulating a value starting with `init`. The
|
||
accumulated value is combined with the each element of the array in reverse order, using `f`.
|
||
|
||
The optional parameters `start` and `stop` control the region of the array to be folded. Folding
|
||
proceeds from `start` (exclusive) to `stop` (inclusive), so no folding occurs unless `start > stop`.
|
||
By default, the entire array is used.
|
||
|
||
Examples:
|
||
* `#[a, b, c].foldr f init = f a (f b (f c init))`
|
||
* `#[1, 2, 3].foldr (toString · ++ ·) "" = "123"`
|
||
* `#[1, 2, 3].foldr (s!"({·} {·})") "!" = "(1 (2 (3 !)))"`
|
||
-/
|
||
@[inline, expose]
|
||
def foldr {α : Type u} {β : Type v} (f : α → β → β) (init : β) (as : Array α) (start := as.size) (stop := 0) : β :=
|
||
Id.run <| as.foldrM (pure <| f · ·) init start stop
|
||
|
||
/--
|
||
Computes the sum of the elements of an array.
|
||
|
||
Examples:
|
||
* `#[a, b, c].sum = a + (b + (c + 0))`
|
||
* `#[1, 2, 5].sum = 8`
|
||
-/
|
||
@[inline, expose]
|
||
def sum {α} [Add α] [Zero α] : Array α → α :=
|
||
foldr (· + ·) 0
|
||
|
||
/--
|
||
Counts the number of elements in the array `as` that satisfy the Boolean predicate `p`.
|
||
|
||
Examples:
|
||
* `#[1, 2, 3, 4, 5].countP (· % 2 == 0) = 2`
|
||
* `#[1, 2, 3, 4, 5].countP (· < 5) = 4`
|
||
* `#[1, 2, 3, 4, 5].countP (· > 5) = 0`
|
||
-/
|
||
@[inline, expose]
|
||
def countP {α : Type u} (p : α → Bool) (as : Array α) : Nat :=
|
||
as.foldr (init := 0) fun a acc => bif p a then acc + 1 else acc
|
||
|
||
/--
|
||
Counts the number of times an element occurs in an array.
|
||
|
||
Examples:
|
||
* `#[1, 1, 2, 3, 5].count 1 = 2`
|
||
* `#[1, 1, 2, 3, 5].count 5 = 1`
|
||
* `#[1, 1, 2, 3, 5].count 4 = 0`
|
||
-/
|
||
@[inline, expose]
|
||
def count {α : Type u} [BEq α] (a : α) (as : Array α) : Nat :=
|
||
countP (· == a) as
|
||
|
||
/--
|
||
Applies a function to each element of the array, returning the resulting array of values.
|
||
|
||
Examples:
|
||
* `#[a, b, c].map f = #[f a, f b, f c]`
|
||
* `#[].map Nat.succ = #[]`
|
||
* `#["one", "two", "three"].map (·.length) = #[3, 3, 5]`
|
||
* `#["one", "two", "three"].map (·.reverse) = #["eno", "owt", "eerht"]`
|
||
-/
|
||
@[inline, expose]
|
||
def map {α : Type u} {β : Type v} (f : α → β) (as : Array α) : Array β :=
|
||
Id.run <| as.mapM (pure <| f ·)
|
||
|
||
instance : Functor Array where
|
||
map := map
|
||
|
||
/--
|
||
Applies a function to each element of the array along with the index at which that element is found,
|
||
returning the array of results. In addition to the index, the function is also provided with a proof
|
||
that the index is valid.
|
||
|
||
`Array.mapIdx` is a variant that does not provide the function with evidence that the index is
|
||
valid.
|
||
-/
|
||
@[inline, expose]
|
||
def mapFinIdx {α : Type u} {β : Type v} (as : Array α) (f : (i : Nat) → α → (h : i < as.size) → β) : Array β :=
|
||
Id.run <| as.mapFinIdxM (pure <| f · · ·)
|
||
|
||
/--
|
||
Applies a function to each element of the array along with the index at which that element is found,
|
||
returning the array of results.
|
||
|
||
`Array.mapFinIdx` is a variant that additionally provides the function with a proof that the index
|
||
is valid.
|
||
-/
|
||
@[inline, expose]
|
||
def mapIdx {α : Type u} {β : Type v} (f : Nat → α → β) (as : Array α) : Array β :=
|
||
Id.run <| as.mapIdxM (pure <| f · ·)
|
||
|
||
/--
|
||
Pairs each element of an array with its index, optionally starting from an index other than `0`.
|
||
|
||
Examples:
|
||
* `#[a, b, c].zipIdx = #[(a, 0), (b, 1), (c, 2)]`
|
||
* `#[a, b, c].zipIdx 5 = #[(a, 5), (b, 6), (c, 7)]`
|
||
-/
|
||
@[expose]
|
||
def zipIdx (xs : Array α) (start := 0) : Array (α × Nat) :=
|
||
xs.mapIdx fun i a => (a, start + i)
|
||
|
||
|
||
|
||
/--
|
||
Returns the first element of the array for which the predicate `p` returns `true`, or `none` if no
|
||
such element is found.
|
||
|
||
Examples:
|
||
* `#[7, 6, 5, 8, 1, 2, 6].find? (· < 5) = some 1`
|
||
* `#[7, 6, 5, 8, 1, 2, 6].find? (· < 1) = none`
|
||
-/
|
||
@[inline, expose]
|
||
def find? {α : Type u} (p : α → Bool) (as : Array α) : Option α :=
|
||
Id.run do
|
||
for a in as do
|
||
if p a then
|
||
return some a
|
||
return none
|
||
|
||
/--
|
||
Returns the first non-`none` result of applying the function `f` to each element of the
|
||
array, in order. Returns `none` if `f` returns `none` for all elements.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #[7, 6, 5, 8, 1, 2, 6].findSome? fun i =>
|
||
if i < 5 then
|
||
some (i * 10)
|
||
else
|
||
none
|
||
```
|
||
```output
|
||
some 10
|
||
```
|
||
-/
|
||
@[inline, expose]
|
||
def findSome? {α : Type u} {β : Type v} (f : α → Option β) (as : Array α) : Option β :=
|
||
Id.run <| as.findSomeM? (pure <| f ·)
|
||
|
||
/--
|
||
Returns the first non-`none` result of applying the function `f` to each element of the
|
||
array, in order. Panics if `f` returns `none` for all elements.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #[7, 6, 5, 8, 1, 2, 6].findSome? fun i =>
|
||
if i < 5 then
|
||
some (i * 10)
|
||
else
|
||
none
|
||
```
|
||
```output
|
||
some 10
|
||
```
|
||
-/
|
||
@[inline]
|
||
def findSome! {α : Type u} {β : Type v} [Inhabited β] (f : α → Option β) (xs : Array α) : β :=
|
||
match xs.findSome? f with
|
||
| some b => b
|
||
| none => panic! "failed to find element"
|
||
|
||
/--
|
||
Returns the first non-`none` result of applying `f` to each element of the array in reverse order,
|
||
from right to left. Returns `none` if `f` returns `none` for all elements of the array.
|
||
|
||
Examples:
|
||
* `#[7, 6, 5, 8, 1, 2, 6].findSome? (fun x => if x < 5 then some (10 * x) else none) = some 10`
|
||
* `#[7, 6, 5, 8, 1, 2, 6].findSome? (fun x => if x < 1 then some (10 * x) else none) = none`
|
||
-/
|
||
@[inline]
|
||
def findSomeRev? {α : Type u} {β : Type v} (f : α → Option β) (as : Array α) : Option β :=
|
||
Id.run <| as.findSomeRevM? (pure <| f ·)
|
||
|
||
/--
|
||
Returns the last element of the array for which the predicate `p` returns `true`, or `none` if no
|
||
such element is found.
|
||
|
||
Examples:
|
||
* `#[7, 6, 5, 8, 1, 2, 6].findRev? (· < 5) = some 2`
|
||
* `#[7, 6, 5, 8, 1, 2, 6].findRev? (· < 1) = none`
|
||
-/
|
||
@[inline]
|
||
def findRev? {α : Type} (p : α → Bool) (as : Array α) : Option α :=
|
||
Id.run <| as.findRevM? (pure <| p ·)
|
||
|
||
/--
|
||
Returns the index of the first element for which `p` returns `true`, or `none` if there is no such
|
||
element.
|
||
|
||
Examples:
|
||
* `#[7, 6, 5, 8, 1, 2, 6].findIdx (· < 5) = some 4`
|
||
* `#[7, 6, 5, 8, 1, 2, 6].findIdx (· < 1) = none`
|
||
-/
|
||
@[inline, expose]
|
||
def findIdx? {α : Type u} (p : α → Bool) (as : Array α) : Option Nat :=
|
||
let rec loop (j : Nat) :=
|
||
if h : j < as.size then
|
||
if p as[j] then some j else loop (j + 1)
|
||
else none
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
loop 0
|
||
|
||
/--
|
||
Returns the index of the first element for which `p` returns `true`, or `none` if there is no such
|
||
element. The index is returned as a `Fin`, which guarantees that it is in bounds.
|
||
|
||
Examples:
|
||
* `#[7, 6, 5, 8, 1, 2, 6].findFinIdx? (· < 5) = some (4 : Fin 7)`
|
||
* `#[7, 6, 5, 8, 1, 2, 6].findFinIdx? (· < 1) = none`
|
||
-/
|
||
@[inline]
|
||
def findFinIdx? {α : Type u} (p : α → Bool) (as : Array α) : Option (Fin as.size) :=
|
||
let rec loop (j : Nat) :=
|
||
if h : j < as.size then
|
||
if p as[j] then some ⟨j, h⟩ else loop (j + 1)
|
||
else none
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
loop 0
|
||
|
||
private theorem findIdx?_loop_eq_map_findFinIdx?_loop_val {xs : Array α} {p : α → Bool} {j} :
|
||
findIdx?.loop p xs j = (findFinIdx?.loop p xs j).map (·.val) := by
|
||
unfold findIdx?.loop
|
||
unfold findFinIdx?.loop
|
||
split <;> rename_i h
|
||
case isTrue =>
|
||
split
|
||
case isTrue => simp
|
||
case isFalse =>
|
||
have : xs.size - (j + 1) < xs.size - j := Nat.sub_succ_lt_self xs.size j h
|
||
rw [findIdx?_loop_eq_map_findFinIdx?_loop_val (j := j + 1)]
|
||
case isFalse => simp
|
||
termination_by xs.size - j
|
||
|
||
theorem findIdx?_eq_map_findFinIdx?_val {xs : Array α} {p : α → Bool} :
|
||
xs.findIdx? p = (xs.findFinIdx? p).map (·.val) := by
|
||
simp [findIdx?, findFinIdx?, findIdx?_loop_eq_map_findFinIdx?_loop_val]
|
||
|
||
/--
|
||
Returns the index of the first element for which `p` returns `true`, or the size of the array if
|
||
there is no such element.
|
||
|
||
Examples:
|
||
* `#[7, 6, 5, 8, 1, 2, 6].findIdx (· < 5) = 4`
|
||
* `#[7, 6, 5, 8, 1, 2, 6].findIdx (· < 1) = 7`
|
||
-/
|
||
@[inline, expose]
|
||
def findIdx (p : α → Bool) (as : Array α) : Nat := (as.findIdx? p).getD as.size
|
||
|
||
def idxOfAux [BEq α] (xs : Array α) (v : α) (i : Nat) : Option (Fin xs.size) :=
|
||
if h : i < xs.size then
|
||
if xs[i] == v then some ⟨i, h⟩
|
||
else idxOfAux xs v (i+1)
|
||
else none
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
|
||
|
||
|
||
/--
|
||
Returns the index of the first element equal to `a`, or `none` if no element is equal
|
||
to `a`. The index is returned as a `Fin`, which guarantees that it is in bounds.
|
||
|
||
Examples:
|
||
* `#["carrot", "potato", "broccoli"].finIdxOf? "carrot" = some 0`
|
||
* `#["carrot", "potato", "broccoli"].finIdxOf? "broccoli" = some 2`
|
||
* `#["carrot", "potato", "broccoli"].finIdxOf? "tomato" = none`
|
||
* `#["carrot", "potato", "broccoli"].finIdxOf? "anything else" = none`
|
||
-/
|
||
def finIdxOf? [BEq α] (xs : Array α) (v : α) : Option (Fin xs.size) :=
|
||
idxOfAux xs v 0
|
||
|
||
|
||
|
||
/--
|
||
Returns the index of the first element equal to `a`, or the size of the array if no element is equal
|
||
to `a`.
|
||
|
||
Examples:
|
||
* `#["carrot", "potato", "broccoli"].idxOf "carrot" = 0`
|
||
* `#["carrot", "potato", "broccoli"].idxOf "broccoli" = 2`
|
||
* `#["carrot", "potato", "broccoli"].idxOf "tomato" = 3`
|
||
* `#["carrot", "potato", "broccoli"].idxOf "anything else" = 3`
|
||
-/
|
||
def idxOf [BEq α] (a : α) : Array α → Nat := findIdx (· == a)
|
||
|
||
/--
|
||
Returns the index of the first element equal to `a`, or `none` if no element is equal to `a`.
|
||
|
||
Examples:
|
||
* `#["carrot", "potato", "broccoli"].idxOf? "carrot" = some 0`
|
||
* `#["carrot", "potato", "broccoli"].idxOf? "broccoli" = some 2`
|
||
* `#["carrot", "potato", "broccoli"].idxOf? "tomato" = none`
|
||
* `#["carrot", "potato", "broccoli"].idxOf? "anything else" = none`
|
||
-/
|
||
def idxOf? [BEq α] (xs : Array α) (v : α) : Option Nat :=
|
||
(xs.finIdxOf? v).map (·.val)
|
||
|
||
/--
|
||
Returns `true` if `p` returns `true` for any element of `as`.
|
||
|
||
Short-circuits upon encountering the first `true`.
|
||
|
||
The optional parameters `start` and `stop` control the region of the array to be checked. Only the
|
||
elements with indices from `start` (inclusive) to `stop` (exclusive) are checked. By default, the
|
||
entire array is checked.
|
||
|
||
Examples:
|
||
* `#[2, 4, 6].any (· % 2 = 0) = true`
|
||
* `#[2, 4, 6].any (· % 2 = 1) = false`
|
||
* `#[2, 4, 5, 6].any (· % 2 = 0) = true`
|
||
* `#[2, 4, 5, 6].any (· % 2 = 1) = true`
|
||
-/
|
||
@[inline, expose, suggest_for Array.some]
|
||
def any (as : Array α) (p : α → Bool) (start := 0) (stop := as.size) : Bool :=
|
||
Id.run <| as.anyM (pure <| p ·) start stop
|
||
|
||
/--
|
||
Returns `true` if `p` returns `true` for every element of `as`.
|
||
|
||
Short-circuits upon encountering the first `false`.
|
||
|
||
The optional parameters `start` and `stop` control the region of the array to be checked. Only the
|
||
elements with indices from `start` (inclusive) to `stop` (exclusive) are checked. By default, the
|
||
entire array is checked.
|
||
|
||
Examples:
|
||
* `#[a, b, c].all p = (p a && (p b && p c))`
|
||
* `#[2, 4, 6].all (· % 2 = 0) = true`
|
||
* `#[2, 4, 5, 6].all (· % 2 = 0) = false`
|
||
-/
|
||
@[inline, suggest_for Array.every]
|
||
def all (as : Array α) (p : α → Bool) (start := 0) (stop := as.size) : Bool :=
|
||
Id.run <| as.allM (pure <| p ·) start stop
|
||
|
||
/--
|
||
Checks whether `a` is an element of `as`, using `==` to compare elements.
|
||
|
||
`Array.elem` is a synonym that takes the element before the array.
|
||
|
||
Examples:
|
||
* `#[1, 4, 2, 3, 3, 7].contains 3 = true`
|
||
* `Array.contains #[1, 4, 2, 3, 3, 7] 5 = false`
|
||
-/
|
||
@[expose]
|
||
def contains [BEq α] (as : Array α) (a : α) : Bool :=
|
||
as.any (a == ·)
|
||
|
||
/--
|
||
Checks whether `a` is an element of `as`, using `==` to compare elements.
|
||
|
||
`Array.contains` is a synonym that takes the array before the element.
|
||
|
||
For verification purposes, `Array.elem` is simplified to `Array.contains`.
|
||
|
||
Example:
|
||
* `Array.elem 3 #[1, 4, 2, 3, 3, 7] = true`
|
||
* `Array.elem 5 #[1, 4, 2, 3, 3, 7] = false`
|
||
-/
|
||
def elem [BEq α] (a : α) (as : Array α) : Bool :=
|
||
as.contains a
|
||
|
||
/-- Convert a `Array α` into an `List α`. This is O(n) in the size of the array. -/
|
||
-- This function is exported to C, where it is called by `Array.toList`
|
||
-- (the projection) to implement this functionality.
|
||
@[export lean_array_to_list_impl]
|
||
def toListImpl (as : Array α) : List α :=
|
||
as.foldr List.cons []
|
||
|
||
/--
|
||
Prepends an array to a list. The elements of the array are at the beginning of the resulting list.
|
||
|
||
Equivalent to `as.toList ++ l`.
|
||
|
||
Examples:
|
||
* `#[1, 2].toListAppend [3, 4] = [1, 2, 3, 4]`
|
||
* `#[1, 2].toListAppend [] = [1, 2]`
|
||
* `#[].toListAppend [3, 4, 5] = [3, 4, 5]`
|
||
-/
|
||
@[inline]
|
||
def toListAppend (as : Array α) (l : List α) : List α :=
|
||
as.foldr List.cons l
|
||
|
||
/--
|
||
Appends two arrays. Normally used via the `++` operator.
|
||
|
||
Appending arrays takes time proportional to the length of the second array.
|
||
|
||
Examples:
|
||
* `#[1, 2, 3] ++ #[4, 5] = #[1, 2, 3, 4, 5]`.
|
||
* `#[] ++ #[4, 5] = #[4, 5]`.
|
||
* `#[1, 2, 3] ++ #[] = #[1, 2, 3]`.
|
||
-/
|
||
@[expose]
|
||
protected def append (as : Array α) (bs : Array α) : Array α :=
|
||
bs.foldl (init := as) fun xs v => xs.push v
|
||
|
||
instance : Append (Array α) := ⟨Array.append⟩
|
||
|
||
/--
|
||
Appends an array and a list.
|
||
|
||
Takes time proportional to the length of the list..
|
||
|
||
Examples:
|
||
* `#[1, 2, 3].appendList [4, 5] = #[1, 2, 3, 4, 5]`.
|
||
* `#[].appendList [4, 5] = #[4, 5]`.
|
||
* `#[1, 2, 3].appendList [] = #[1, 2, 3]`.
|
||
-/
|
||
protected def appendList (as : Array α) (bs : List α) : Array α :=
|
||
bs.foldl (init := as) fun xs v => xs.push v
|
||
|
||
instance : HAppend (Array α) (List α) (Array α) := ⟨Array.appendList⟩
|
||
|
||
|
||
/--
|
||
Applies a monadic function that returns an array to each element of an array, from left to right.
|
||
The resulting arrays are appended.
|
||
-/
|
||
@[inline]
|
||
def flatMapM [Monad m] (f : α → m (Array β)) (as : Array α) : m (Array β) :=
|
||
as.foldlM (init := empty) fun bs a => do return bs ++ (← f a)
|
||
|
||
/--
|
||
Applies a function that returns an array to each element of an array. The resulting arrays are
|
||
appended.
|
||
|
||
Examples:
|
||
* `#[2, 3, 2].flatMap Array.range = #[0, 1, 0, 1, 2, 0, 1]`
|
||
* `#[['a', 'b'], ['c', 'd', 'e']].flatMap List.toArray = #['a', 'b', 'c', 'd', 'e']`
|
||
-/
|
||
@[inline, expose]
|
||
def flatMap (f : α → Array β) (as : Array α) : Array β :=
|
||
as.foldl (init := empty) fun bs a => bs ++ f a
|
||
|
||
/--
|
||
Appends the contents of array of arrays into a single array. The resulting array contains the same
|
||
elements as the nested arrays in the same order.
|
||
|
||
Examples:
|
||
* `#[#[5], #[4], #[3, 2]].flatten = #[5, 4, 3, 2]`
|
||
* `#[#[0, 1], #[], #[2], #[1, 0, 1]].flatten = #[0, 1, 2, 1, 0, 1]`
|
||
* `(#[] : Array Nat).flatten = #[]`
|
||
-/
|
||
@[inline, expose] def flatten (xss : Array (Array α)) : Array α :=
|
||
xss.foldl (init := empty) fun acc xs => acc ++ xs
|
||
|
||
/--
|
||
Reverses an array by repeatedly swapping elements.
|
||
|
||
The original array is modified in place if there are no other references to it.
|
||
|
||
Examples:
|
||
* `(#[] : Array Nat).reverse = #[]`
|
||
* `#[0, 1].reverse = #[1, 0]`
|
||
* `#[0, 1, 2].reverse = #[2, 1, 0]`
|
||
-/
|
||
@[expose]
|
||
def reverse (as : Array α) : Array α :=
|
||
if h : as.size ≤ 1 then
|
||
as
|
||
else
|
||
loop as 0 ⟨as.size - 1, Nat.pred_lt (mt (fun h : as.size = 0 => h ▸ by decide) h)⟩
|
||
where
|
||
termination {i j : Nat} (h : i < j) : j - 1 - (i + 1) < j - i := by
|
||
rw [Nat.sub_sub, Nat.add_comm]
|
||
exact Nat.lt_of_le_of_lt (Nat.pred_le _) (Nat.sub_succ_lt_self _ _ h)
|
||
loop (as : Array α) (i : Nat) (j : Fin as.size) :=
|
||
if h : i < j then
|
||
have := termination h
|
||
let as := as.swap i j (Nat.lt_trans h j.2)
|
||
have : j-1 < as.size := by rw [size_swap]; exact Nat.lt_of_le_of_lt (Nat.pred_le _) j.2
|
||
loop as (i+1) ⟨j-1, this⟩
|
||
else
|
||
as
|
||
|
||
/--
|
||
Returns the array of elements in `as` for which `p` returns `true`.
|
||
|
||
Only elements from `start` (inclusive) to `stop` (exclusive) are considered. Elements outside that
|
||
range are discarded. By default, the entire array is considered.
|
||
|
||
Examples:
|
||
* `#[1, 2, 5, 2, 7, 7].filter (· > 2) = #[5, 7, 7]`
|
||
* `#[1, 2, 5, 2, 7, 7].filter (fun _ => false) = #[]`
|
||
* `#[1, 2, 5, 2, 7, 7].filter (fun _ => true) = #[1, 2, 5, 2, 7, 7]`
|
||
* `#[1, 2, 5, 2, 7, 7].filter (· > 2) (start := 3) = #[7, 7]`
|
||
* `#[1, 2, 5, 2, 7, 7].filter (fun _ => true) (start := 3) = #[2, 7, 7]`
|
||
* `#[1, 2, 5, 2, 7, 7].filter (fun _ => true) (stop := 3) = #[1, 2, 5]`
|
||
-/
|
||
@[inline, expose]
|
||
def filter (p : α → Bool) (as : Array α) (start := 0) (stop := as.size) : Array α :=
|
||
as.foldl (init := #[]) (start := start) (stop := stop) fun acc a =>
|
||
if p a then acc.push a else acc
|
||
|
||
/--
|
||
Applies the monadic predicate `p` to every element in the array, in order from left to right, and
|
||
returns the array of elements for which `p` returns `true`.
|
||
|
||
Only elements from `start` (inclusive) to `stop` (exclusive) are considered. Elements outside that
|
||
range are discarded. By default, the entire array is checked.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #[1, 2, 5, 2, 7, 7].filterM fun x => do
|
||
IO.println s!"Checking {x}"
|
||
return x < 3
|
||
```
|
||
```output
|
||
Checking 1
|
||
Checking 2
|
||
Checking 5
|
||
Checking 2
|
||
Checking 7
|
||
Checking 7
|
||
```
|
||
```output
|
||
#[1, 2, 2]
|
||
```
|
||
-/
|
||
@[inline]
|
||
def filterM {α : Type} [Monad m] (p : α → m Bool) (as : Array α) (start := 0) (stop := as.size) : m (Array α) :=
|
||
as.foldlM (init := #[]) (start := start) (stop := stop) fun acc a => do
|
||
if (← p a) then return acc.push a else return acc
|
||
|
||
/--
|
||
Applies the monadic predicate `p` on every element in the array in reverse order, from right to
|
||
left, and returns those elements for which `p` returns `true`. The elements of the returned list are
|
||
in the same order as in the input list.
|
||
|
||
Only elements from `start` (exclusive) to `stop` (inclusive) are considered. Elements outside that
|
||
range are discarded. Because the array is examined in reverse order, elements are only examined when
|
||
`start > stop`. By default, the entire array is considered.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #[1, 2, 5, 2, 7, 7].filterRevM fun x => do
|
||
IO.println s!"Checking {x}"
|
||
return x < 3
|
||
```
|
||
```output
|
||
Checking 7
|
||
Checking 7
|
||
Checking 2
|
||
Checking 5
|
||
Checking 2
|
||
Checking 1
|
||
```
|
||
```output
|
||
#[1, 2, 2]
|
||
```
|
||
-/
|
||
@[inline]
|
||
def filterRevM {α : Type} [Monad m] (p : α → m Bool) (as : Array α) (start := as.size) (stop := 0) : m (Array α) :=
|
||
reverse <$> as.foldrM (init := #[]) (start := start) (stop := stop) fun a acc => do
|
||
if (← p a) then return acc.push a else return acc
|
||
|
||
/--
|
||
Applies a monadic function that returns an `Option` to each element of an array, collecting the
|
||
non-`none` values.
|
||
|
||
Only elements from `start` (inclusive) to `stop` (exclusive) are considered. Elements outside that
|
||
range are discarded. By default, the entire array is considered.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #[1, 2, 5, 2, 7, 7].filterMapM fun x => do
|
||
IO.println s!"Examining {x}"
|
||
if x > 2 then return some (2 * x)
|
||
else return none
|
||
```
|
||
```output
|
||
Examining 1
|
||
Examining 2
|
||
Examining 5
|
||
Examining 2
|
||
Examining 7
|
||
Examining 7
|
||
```
|
||
```output
|
||
#[10, 14, 14]
|
||
```
|
||
-/
|
||
@[specialize, expose]
|
||
def filterMapM [Monad m] (f : α → m (Option β)) (as : Array α) (start := 0) (stop := as.size) : m (Array β) :=
|
||
as.foldlM (init := #[]) (start := start) (stop := stop) fun bs a => do
|
||
match (← f a) with
|
||
| some b => pure (bs.push b)
|
||
| none => pure bs
|
||
|
||
/--
|
||
Applies a function that returns an `Option` to each element of an array, collecting the non-`none`
|
||
values.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #[1, 2, 5, 2, 7, 7].filterMap fun x =>
|
||
if x > 2 then some (2 * x) else none
|
||
```
|
||
```output
|
||
#[10, 14, 14]
|
||
```
|
||
-/
|
||
@[inline, expose]
|
||
def filterMap (f : α → Option β) (as : Array α) (start := 0) (stop := as.size) : Array β :=
|
||
Id.run <| as.filterMapM (pure <| f ·) (start := start) (stop := stop)
|
||
|
||
/--
|
||
Returns the largest element of the array, as determined by the comparison `lt`, or `none` if
|
||
the array is empty.
|
||
|
||
Examples:
|
||
* `(#[] : Array Nat).getMax? (· < ·) = none`
|
||
* `#["red", "green", "blue"].getMax? (·.length < ·.length) = some "green"`
|
||
* `#["red", "green", "blue"].getMax? (· < ·) = some "red"`
|
||
-/
|
||
@[specialize]
|
||
def getMax? (as : Array α) (lt : α → α → Bool) : Option α :=
|
||
if h : 0 < as.size then
|
||
let a0 := as[0]
|
||
some <| as.foldl (init := a0) (start := 1) fun best a =>
|
||
if lt best a then a else best
|
||
else
|
||
none
|
||
|
||
/--
|
||
Returns a pair of arrays that together contain all the elements of `as`. The first array contains
|
||
those elements for which `p` returns `true`, and the second contains those for which `p` returns
|
||
`false`.
|
||
|
||
`as.partition p` is equivalent to `(as.filter p, as.filter (not ∘ p))`, but it is
|
||
more efficient since it only has to do one pass over the array.
|
||
|
||
Examples:
|
||
* `#[1, 2, 5, 2, 7, 7].partition (· > 2) = (#[5, 7, 7], #[1, 2, 2])`
|
||
* `#[1, 2, 5, 2, 7, 7].partition (fun _ => false) = (#[], #[1, 2, 5, 2, 7, 7])`
|
||
* `#[1, 2, 5, 2, 7, 7].partition (fun _ => true) = (#[1, 2, 5, 2, 7, 7], #[])`
|
||
-/
|
||
@[inline]
|
||
def partition (p : α → Bool) (as : Array α) : Array α × Array α := Id.run do
|
||
let mut bs := #[]
|
||
let mut cs := #[]
|
||
for a in as do
|
||
if p a then
|
||
bs := bs.push a
|
||
else
|
||
cs := cs.push a
|
||
return (bs, cs)
|
||
|
||
/--
|
||
Removes all the elements that satisfy a predicate from the end of an array.
|
||
|
||
The longest contiguous sequence of elements that all satisfy the predicate is removed.
|
||
|
||
Examples:
|
||
* `#[0, 1, 2, 3, 4].popWhile (· > 2) = #[0, 1, 2]`
|
||
* `#[3, 2, 3, 4].popWhile (· > 2) = #[3, 2]`
|
||
* `(#[] : Array Nat).popWhile (· > 2) = #[]`
|
||
-/
|
||
def popWhile (p : α → Bool) (as : Array α) : Array α :=
|
||
if h : as.size > 0 then
|
||
if p (as[as.size - 1]'(Nat.sub_lt h (by decide))) then
|
||
popWhile p as.pop
|
||
else
|
||
as
|
||
else
|
||
as
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
|
||
@[simp, grind =] theorem popWhile_empty {p : α → Bool} :
|
||
popWhile p #[] = #[] := by
|
||
simp [popWhile]
|
||
|
||
/--
|
||
Returns a new array that contains the longest prefix of elements that satisfy the predicate `p` from
|
||
an array.
|
||
|
||
Examples:
|
||
* `#[0, 1, 2, 3, 2, 1].takeWhile (· < 2) = #[0, 1]`
|
||
* `#[0, 1, 2, 3, 2, 1].takeWhile (· < 20) = #[0, 1, 2, 3, 2, 1]`
|
||
* `#[0, 1, 2, 3, 2, 1].takeWhile (· < 0) = #[]`
|
||
-/
|
||
def takeWhile (p : α → Bool) (as : Array α) : Array α :=
|
||
let rec go (i : Nat) (acc : Array α) : Array α :=
|
||
if h : i < as.size then
|
||
let a := as[i]
|
||
if p a then
|
||
go (i+1) (acc.push a)
|
||
else
|
||
acc
|
||
else
|
||
acc
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
go 0 #[]
|
||
|
||
/--
|
||
Removes the element at a given index from an array without a run-time bounds check.
|
||
|
||
This function takes worst-case `O(n)` time because it back-shifts all elements at positions
|
||
greater than `i`.
|
||
|
||
Examples:
|
||
* `#["apple", "pear", "orange"].eraseIdx 0 = #["pear", "orange"]`
|
||
* `#["apple", "pear", "orange"].eraseIdx 1 = #["apple", "orange"]`
|
||
* `#["apple", "pear", "orange"].eraseIdx 2 = #["apple", "pear"]`
|
||
-/
|
||
def eraseIdx (xs : Array α) (i : Nat) (h : i < xs.size := by get_elem_tactic) : Array α :=
|
||
if h' : i + 1 < xs.size then
|
||
let xs' := xs.swap (i + 1) i
|
||
xs'.eraseIdx (i + 1) (by simp [xs', h'])
|
||
else
|
||
xs.pop
|
||
termination_by xs.size - i
|
||
decreasing_by simp_wf; exact Nat.sub_succ_lt_self _ _ h
|
||
|
||
-- This is required in `Lean.Data.PersistentHashMap`.
|
||
@[simp, grind =]
|
||
theorem size_eraseIdx {xs : Array α} (i : Nat) (h) : (xs.eraseIdx i h).size = xs.size - 1 := by
|
||
induction xs, i, h using Array.eraseIdx.induct with
|
||
| @case1 xs i h h' xs' ih =>
|
||
unfold eraseIdx
|
||
simp +zetaDelta [h', ih]
|
||
| case2 xs i h h' =>
|
||
unfold eraseIdx
|
||
simp [h']
|
||
|
||
/--
|
||
Removes the element at a given index from an array. Does nothing if the index is out of bounds.
|
||
|
||
This function takes worst-case `O(n)` time because it back-shifts all elements at positions greater
|
||
than `i`.
|
||
|
||
Examples:
|
||
* `#["apple", "pear", "orange"].eraseIdxIfInBounds 0 = #["pear", "orange"]`
|
||
* `#["apple", "pear", "orange"].eraseIdxIfInBounds 1 = #["apple", "orange"]`
|
||
* `#["apple", "pear", "orange"].eraseIdxIfInBounds 2 = #["apple", "pear"]`
|
||
* `#["apple", "pear", "orange"].eraseIdxIfInBounds 3 = #["apple", "pear", "orange"]`
|
||
* `#["apple", "pear", "orange"].eraseIdxIfInBounds 5 = #["apple", "pear", "orange"]`
|
||
-/
|
||
def eraseIdxIfInBounds (xs : Array α) (i : Nat) : Array α :=
|
||
if h : i < xs.size then xs.eraseIdx i h else xs
|
||
|
||
/--
|
||
Removes the element at a given index from an array. Panics if the index is out of bounds.
|
||
|
||
This function takes worst-case `O(n)` time because it back-shifts all elements at positions
|
||
greater than `i`.
|
||
-/
|
||
def eraseIdx! (xs : Array α) (i : Nat) : Array α :=
|
||
if h : i < xs.size then xs.eraseIdx i h else panic! "invalid index"
|
||
|
||
/--
|
||
Removes the first occurrence of a specified element from an array, or does nothing if it is not
|
||
present.
|
||
|
||
This function takes worst-case `O(n)` time because it back-shifts all later elements.
|
||
|
||
Examples:
|
||
* `#[1, 2, 3].erase 2 = #[1, 3]`
|
||
* `#[1, 2, 3].erase 5 = #[1, 2, 3]`
|
||
* `#[1, 2, 3, 2, 1].erase 2 = #[1, 3, 2, 1]`
|
||
* `(#[] : List Nat).erase 2 = #[]`
|
||
-/
|
||
def erase [BEq α] (as : Array α) (a : α) : Array α :=
|
||
match as.finIdxOf? a with
|
||
| none => as
|
||
| some i => as.eraseIdx i
|
||
|
||
/--
|
||
Removes the first element that satisfies the predicate `p`. If no element satisfies `p`, the array
|
||
is returned unmodified.
|
||
|
||
This function takes worst-case `O(n)` time because it back-shifts all later elements.
|
||
|
||
Examples:
|
||
* `#["red", "green", "", "blue"].eraseP (·.isEmpty) = #["red", "green", "blue"]`
|
||
* `#["red", "green", "", "blue", ""].eraseP (·.isEmpty) = #["red", "green", "blue", ""]`
|
||
* `#["red", "green", "blue"].eraseP (·.length % 2 == 0) = #["red", "green"]`
|
||
* `#["red", "green", "blue"].eraseP (fun _ => true) = #["green", "blue"]`
|
||
* `(#[] : Array String).eraseP (fun _ => true) = #[]`
|
||
-/
|
||
def eraseP (as : Array α) (p : α → Bool) : Array α :=
|
||
match as.findFinIdx? p with
|
||
| none => as
|
||
| some i => as.eraseIdx i
|
||
|
||
/--
|
||
Inserts an element into an array at the specified index. If the index is greater than the size of
|
||
the array, then the array is returned unmodified.
|
||
|
||
In other words, the new element is inserted into the array `as` after the first `i` elements of
|
||
`as`.
|
||
|
||
This function takes worst case `O(n)` time because it has to swap the inserted element into place.
|
||
|
||
Examples:
|
||
* `#["tues", "thur", "sat"].insertIdx 1 "wed" = #["tues", "wed", "thur", "sat"]`
|
||
* `#["tues", "thur", "sat"].insertIdx 2 "wed" = #["tues", "thur", "wed", "sat"]`
|
||
* `#["tues", "thur", "sat"].insertIdx 3 "wed" = #["tues", "thur", "sat", "wed"]`
|
||
-/
|
||
@[inline] def insertIdx (as : Array α) (i : Nat) (a : α) (_ : i ≤ as.size := by get_elem_tactic) : Array α :=
|
||
let rec loop (as : Array α) (j : Fin as.size) :=
|
||
if i < j then
|
||
let j' : Fin as.size := ⟨j-1, Nat.lt_of_le_of_lt (Nat.pred_le _) j.2⟩
|
||
let as := as.swap j' j
|
||
loop as ⟨j', by rw [size_swap]; exact j'.2⟩
|
||
else
|
||
as
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
let j := as.size
|
||
let as := as.push a
|
||
loop as ⟨j, size_push .. ▸ j.lt_succ_self⟩
|
||
|
||
/--
|
||
Inserts an element into an array at the specified index. Panics if the index is greater than the
|
||
size of the array.
|
||
|
||
In other words, the new element is inserted into the array `as` after the first `i` elements of
|
||
`as`.
|
||
|
||
This function takes worst case `O(n)` time because it has to swap the inserted element into place.
|
||
`Array.insertIdx` and `Array.insertIdxIfInBounds` are safer alternatives.
|
||
|
||
Examples:
|
||
* `#["tues", "thur", "sat"].insertIdx! 1 "wed" = #["tues", "wed", "thur", "sat"]`
|
||
* `#["tues", "thur", "sat"].insertIdx! 2 "wed" = #["tues", "thur", "wed", "sat"]`
|
||
* `#["tues", "thur", "sat"].insertIdx! 3 "wed" = #["tues", "thur", "sat", "wed"]`
|
||
-/
|
||
def insertIdx! (as : Array α) (i : Nat) (a : α) : Array α :=
|
||
if h : i ≤ as.size then
|
||
insertIdx as i a
|
||
else panic! "invalid index"
|
||
|
||
/--
|
||
Inserts an element into an array at the specified index. The array is returned unmodified if the
|
||
index is greater than the size of the array.
|
||
|
||
In other words, the new element is inserted into the array `as` after the first `i` elements of
|
||
`as`.
|
||
|
||
This function takes worst case `O(n)` time because it has to swap the inserted element into place.
|
||
|
||
Examples:
|
||
* `#["tues", "thur", "sat"].insertIdxIfInBounds 1 "wed" = #["tues", "wed", "thur", "sat"]`
|
||
* `#["tues", "thur", "sat"].insertIdxIfInBounds 2 "wed" = #["tues", "thur", "wed", "sat"]`
|
||
* `#["tues", "thur", "sat"].insertIdxIfInBounds 3 "wed" = #["tues", "thur", "sat", "wed"]`
|
||
* `#["tues", "thur", "sat"].insertIdxIfInBounds 4 "wed" = #["tues", "thur", "sat"]`
|
||
-/
|
||
@[grind]
|
||
def insertIdxIfInBounds (as : Array α) (i : Nat) (a : α) : Array α :=
|
||
if h : i ≤ as.size then
|
||
insertIdx as i a
|
||
else
|
||
as
|
||
|
||
def isPrefixOfAux [BEq α] (as bs : Array α) (hle : as.size ≤ bs.size) (i : Nat) : Bool :=
|
||
if h : i < as.size then
|
||
let a := as[i]
|
||
have : i < bs.size := Nat.lt_of_lt_of_le h hle
|
||
let b := bs[i]
|
||
if a == b then
|
||
isPrefixOfAux as bs hle (i+1)
|
||
else
|
||
false
|
||
else
|
||
true
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
|
||
/--
|
||
Return `true` if `as` is a prefix of `bs`, or `false` otherwise.
|
||
|
||
Examples:
|
||
* `#[0, 1, 2].isPrefixOf #[0, 1, 2, 3] = true`
|
||
* `#[0, 1, 2].isPrefixOf #[0, 1, 2] = true`
|
||
* `#[0, 1, 2].isPrefixOf #[0, 1] = false`
|
||
* `#[].isPrefixOf #[0, 1] = true`
|
||
-/
|
||
def isPrefixOf [BEq α] (as bs : Array α) : Bool :=
|
||
if h : as.size ≤ bs.size then
|
||
isPrefixOfAux as bs h 0
|
||
else
|
||
false
|
||
|
||
@[specialize]
|
||
def zipWithMAux {m : Type v → Type w} [Monad m] (as : Array α) (bs : Array β) (f : α → β → m γ) (i : Nat) (cs : Array γ) : m (Array γ) := do
|
||
if h : i < as.size then
|
||
let a := as[i]
|
||
if h : i < bs.size then
|
||
let b := bs[i]
|
||
zipWithMAux as bs f (i+1) <| cs.push (← f a b)
|
||
else
|
||
return cs
|
||
else
|
||
return cs
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
|
||
/--
|
||
Applies a function to the corresponding elements of two arrays, stopping at the end of the shorter
|
||
array.
|
||
|
||
Examples:
|
||
* `#[1, 2].zipWith (· + ·) #[5, 6] = #[6, 8]`
|
||
* `#[1, 2, 3].zipWith (· + ·) #[5, 6, 10] = #[6, 8, 13]`
|
||
* `#[].zipWith (· + ·) #[5, 6] = #[]`
|
||
* `#[x₁, x₂, x₃].zipWith f #[y₁, y₂, y₃, y₄] = #[f x₁ y₁, f x₂ y₂, f x₃ y₃]`
|
||
-/
|
||
@[inline] def zipWith (f : α → β → γ) (as : Array α) (bs : Array β) : Array γ :=
|
||
Id.run (zipWithMAux as bs (pure <| f · ·) 0 #[])
|
||
|
||
/--
|
||
Combines two arrays into an array of pairs in which the first and second components are the
|
||
corresponding elements of each input array. The resulting array is the length of the shorter of the
|
||
input arrays.
|
||
|
||
Examples:
|
||
* `#["Mon", "Tue", "Wed"].zip #[1, 2, 3] = #[("Mon", 1), ("Tue", 2), ("Wed", 3)]`
|
||
* `#["Mon", "Tue", "Wed"].zip #[1, 2] = #[("Mon", 1), ("Tue", 2)]`
|
||
* `#[x₁, x₂, x₃].zip #[y₁, y₂, y₃, y₄] = #[(x₁, y₁), (x₂, y₂), (x₃, y₃)]`
|
||
-/
|
||
def zip (as : Array α) (bs : Array β) : Array (α × β) :=
|
||
zipWith Prod.mk as bs
|
||
|
||
/--
|
||
Applies a function to the corresponding elements of both arrays, stopping when there are no more
|
||
elements in either array. If one array is shorter than the other, the function is passed `none` for
|
||
the missing elements.
|
||
|
||
Examples:
|
||
* `#[1, 6].zipWithAll min #[5, 2] = #[some 1, some 2]`
|
||
* `#[1, 2, 3].zipWithAll Prod.mk #[5, 6] = #[(some 1, some 5), (some 2, some 6), (some 3, none)]`
|
||
* `#[x₁, x₂].zipWithAll f #[y] = #[f (some x₁) (some y), f (some x₂) none]`
|
||
-/
|
||
def zipWithAll (f : Option α → Option β → γ) (as : Array α) (bs : Array β) : 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
|
||
|
||
/--
|
||
Applies a monadic function to the corresponding elements of two arrays, left-to-right, stopping at
|
||
the end of the shorter array. `zipWithM f as bs` is equivalent to `mapM id (zipWith f as bs)`.
|
||
-/
|
||
@[inline] def zipWithM {m : Type v → Type w} [Monad m] (f : α → β → m γ) (as : Array α) (bs : Array β) : m (Array γ) :=
|
||
zipWithMAux as bs f 0 #[]
|
||
|
||
/--
|
||
Separates an array of pairs into two arrays that contain the respective first and second components.
|
||
|
||
Examples:
|
||
* `#[("Monday", 1), ("Tuesday", 2)].unzip = (#["Monday", "Tuesday"], #[1, 2])`
|
||
* `#[(x₁, y₁), (x₂, y₂), (x₃, y₃)].unzip = (#[x₁, x₂, x₃], #[y₁, y₂, y₃])`
|
||
* `(#[] : Array (Nat × String)).unzip = ((#[], #[]) : List Nat × List String)`
|
||
-/
|
||
def unzip (as : Array (α × β)) : Array α × Array β :=
|
||
as.foldl (init := (#[], #[])) fun (as, bs) (a, b) => (as.push a, bs.push b)
|
||
|
||
/--
|
||
Replaces the first occurrence of `a` with `b` in an array. The modification is performed in-place
|
||
when the reference to the array is unique. Returns the array unmodified when `a` is not present.
|
||
|
||
Examples:
|
||
* `#[1, 2, 3, 2, 1].replace 2 5 = #[1, 5, 3, 2, 1]`
|
||
* `#[1, 2, 3, 2, 1].replace 0 5 = #[1, 2, 3, 2, 1]`
|
||
* `#[].replace 2 5 = #[]`
|
||
-/
|
||
def replace [BEq α] (xs : Array α) (a b : α) : Array α :=
|
||
match xs.finIdxOf? a with
|
||
| none => xs
|
||
| some i => xs.set i b
|
||
|
||
/-! ### Lexicographic ordering -/
|
||
|
||
instance instLT [LT α] : LT (Array α) := ⟨fun as bs => as.toList < bs.toList⟩
|
||
instance instLE [LT α] : LE (Array α) := ⟨fun as bs => as.toList ≤ bs.toList⟩
|
||
|
||
-- See `Init.Data.Array.Lex.Basic` for the boolean valued lexicographic comparator.
|
||
|
||
/-! ## Auxiliary functions used in metaprogramming.
|
||
|
||
We do not currently intend to provide verification theorems for these functions.
|
||
-/
|
||
|
||
/-! ### leftpad and rightpad -/
|
||
|
||
/--
|
||
Pads `xs : Array α` on the left with repeated occurrences of `a : α` until it is of size `n`. If `xs`
|
||
already has at least `n` elements, it is returned unmodified.
|
||
|
||
Examples:
|
||
* `#[1, 2, 3].leftpad 5 0 = #[0, 0, 1, 2, 3]`
|
||
* `#["red", "green", "blue"].leftpad 4 "blank" = #["blank", "red", "green", "blue"]`
|
||
* `#["red", "green", "blue"].leftpad 3 "blank" = #["red", "green", "blue"]`
|
||
* `#["red", "green", "blue"].leftpad 1 "blank" = #["red", "green", "blue"]`
|
||
-/
|
||
def leftpad (n : Nat) (a : α) (xs : Array α) : Array α := replicate (n - xs.size) a ++ xs
|
||
|
||
/--
|
||
Pads `xs : Array α` on the right with repeated occurrences of `a : α` until it is of length `n`. If
|
||
`l` already has at least `n` elements, it is returned unmodified.
|
||
|
||
Examples:
|
||
* `#[1, 2, 3].rightpad 5 0 = #[1, 2, 3, 0, 0]`
|
||
* `#["red", "green", "blue"].rightpad 4 "blank" = #["red", "green", "blue", "blank"]`
|
||
* `#["red", "green", "blue"].rightpad 3 "blank" = #["red", "green", "blue"]`
|
||
* `#["red", "green", "blue"].rightpad 1 "blank" = #["red", "green", "blue"]`
|
||
-/
|
||
def rightpad (n : Nat) (a : α) (xs : Array α) : Array α := xs ++ replicate (n - xs.size) a
|
||
|
||
/- ### reduceOption -/
|
||
|
||
/-- Drop `none`s from a Array, and replace each remaining `some a` with `a`. -/
|
||
@[inline] def reduceOption (as : Array (Option α)) : Array α :=
|
||
as.filterMap id
|
||
|
||
/-! ### eraseReps -/
|
||
|
||
/--
|
||
Erases repeated elements, keeping the first element of each run.
|
||
|
||
`O(|as|)`.
|
||
|
||
Example:
|
||
* `#[1, 3, 2, 2, 2, 3, 3, 5].eraseReps = #[1, 3, 2, 3, 5]`
|
||
-/
|
||
def eraseReps {α} [BEq α] (as : Array α) : Array α :=
|
||
if h : 0 < as.size then
|
||
let ⟨last, acc⟩ := as.foldl (init := (as[0], #[])) fun ⟨last, acc⟩ a =>
|
||
if a == last then ⟨last, acc⟩ else ⟨a, acc.push last⟩
|
||
acc.push last
|
||
else
|
||
#[]
|
||
|
||
/-! ### allDiff -/
|
||
|
||
private def allDiffAuxAux [BEq α] (as : Array α) (a : α) : forall (i : Nat), i < as.size → Bool
|
||
| 0, _ => true
|
||
| i+1, h =>
|
||
have : i < as.size := Nat.lt_trans (Nat.lt_succ_self _) h;
|
||
a != as[i] && allDiffAuxAux as a i this
|
||
|
||
private def allDiffAux [BEq α] (as : Array α) (i : Nat) : Bool :=
|
||
if h : i < as.size then
|
||
allDiffAuxAux as as[i] i h && allDiffAux as (i+1)
|
||
else
|
||
true
|
||
decreasing_by simp_wf; decreasing_trivial_pre_omega
|
||
|
||
/--
|
||
Returns `true` if no two elements of `as` are equal according to the `==` operator.
|
||
|
||
Examples:
|
||
* `#["red", "green", "blue"].allDiff = true`
|
||
* `#["red", "green", "red"].allDiff = false`
|
||
* `(#[] : Array Nat).allDiff = true`
|
||
-/
|
||
def allDiff [BEq α] (as : Array α) : Bool :=
|
||
allDiffAux as 0
|
||
|
||
/-! ### getEvenElems -/
|
||
|
||
/--
|
||
Returns a new array that contains the elements at even indices in `as`, starting with the element at
|
||
index `0`.
|
||
|
||
Examples:
|
||
* `#[0, 1, 2, 3, 4].getEvenElems = #[0, 2, 4]`
|
||
* `#[1, 2, 3, 4].getEvenElems = #[1, 3]`
|
||
* `#["red", "green", "blue"].getEvenElems = #["red", "blue"]`
|
||
* `(#[] : Array String).getEvenElems = #[]`
|
||
-/
|
||
@[inline] def getEvenElems (as : Array α) : Array α :=
|
||
(·.2) <| as.foldl (init := (true, Array.empty)) fun (even, acc) a =>
|
||
if even then
|
||
(false, acc.push a)
|
||
else
|
||
(true, acc)
|
||
|
||
/-! ### Repr and ToString -/
|
||
|
||
protected def Array.repr {α : Type u} [Repr α] (xs : Array α) : Std.Format :=
|
||
let _ : Std.ToFormat α := ⟨repr⟩
|
||
if xs.size == 0 then
|
||
"#[]"
|
||
else
|
||
Std.Format.bracketFill "#[" (Std.Format.joinSep (toList xs) ("," ++ Std.Format.line)) "]"
|
||
|
||
instance {α : Type u} [Repr α] : Repr (Array α) where
|
||
reprPrec xs _ := Array.repr xs
|
||
|
||
instance [ToString α] : ToString (Array α) where
|
||
toString xs := String.Internal.append "#" (toString xs.toList)
|
||
|
||
end Array
|