412 lines
13 KiB
Text
412 lines
13 KiB
Text
/-
|
||
Copyright (c) 2020 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.Data.Array.Basic
|
||
import Init.Data.Array.GetLit
|
||
public import Init.Data.Slice.Basic
|
||
|
||
public section
|
||
|
||
set_option linter.indexVariables true -- Enforce naming conventions for index variables.
|
||
set_option linter.missingDocs true
|
||
|
||
universe u v w
|
||
|
||
/--
|
||
Internal representation of `Subarray`, which is an abbreviation for `Slice SubarrayData`.
|
||
-/
|
||
structure Std.Slice.Internal.SubarrayData (α : Type u) where
|
||
/-- The underlying array. -/
|
||
array : Array α
|
||
/-- The starting index of the region of interest (inclusive). -/
|
||
start : Nat
|
||
/-- The ending index of the region of interest (exclusive). -/
|
||
stop : Nat
|
||
/--
|
||
The starting index is no later than the ending index.
|
||
|
||
The ending index is exclusive. If the starting and ending indices are equal, then the subarray is
|
||
empty.
|
||
-/
|
||
start_le_stop : start ≤ stop
|
||
/-- The stopping index is no later than the end of the array.
|
||
|
||
The ending index is exclusive. If it is equal to the size of the array, then the last element of
|
||
the array is in the subarray.
|
||
-/
|
||
stop_le_array_size : stop ≤ array.size
|
||
|
||
open Std.Slice
|
||
|
||
/--
|
||
A region of some underlying array.
|
||
|
||
A subarray contains an array together with the start and end indices of a region of interest.
|
||
Subarrays can be used to avoid copying or allocating space, while being more convenient than
|
||
tracking the bounds by hand. The region of interest consists of every index that is both greater
|
||
than or equal to `start` and strictly less than `stop`.
|
||
-/
|
||
abbrev Subarray (α : Type u) := Std.Slice (Internal.SubarrayData α)
|
||
|
||
instance {α : Type u} : Self (Std.Slice (Internal.SubarrayData α)) (Subarray α) where
|
||
|
||
@[always_inline, inline, expose, inherit_doc Internal.SubarrayData.array]
|
||
def Subarray.array (xs : Subarray α) : Array α :=
|
||
xs.internalRepresentation.array
|
||
|
||
@[always_inline, inline, expose, inherit_doc Internal.SubarrayData.start]
|
||
def Subarray.start (xs : Subarray α) : Nat :=
|
||
xs.internalRepresentation.start
|
||
|
||
@[always_inline, inline, expose, inherit_doc Internal.SubarrayData.stop]
|
||
def Subarray.stop (xs : Subarray α) : Nat :=
|
||
xs.internalRepresentation.stop
|
||
|
||
@[always_inline, inline, expose, inherit_doc Internal.SubarrayData.start_le_stop]
|
||
def Subarray.start_le_stop (xs : Subarray α) : xs.start ≤ xs.stop :=
|
||
xs.internalRepresentation.start_le_stop
|
||
|
||
@[always_inline, inline, expose, inherit_doc Internal.SubarrayData.stop_le_array_size]
|
||
def Subarray.stop_le_array_size (xs : Subarray α) : xs.stop ≤ xs.array.size :=
|
||
xs.internalRepresentation.stop_le_array_size
|
||
|
||
namespace Subarray
|
||
|
||
/--
|
||
Computes the size of the subarray.
|
||
-/
|
||
def size (s : Subarray α) : Nat :=
|
||
s.stop - s.start
|
||
|
||
theorem size_le_array_size {s : Subarray α} : s.size ≤ s.array.size := by
|
||
let ⟨{array, start, stop, start_le_stop, stop_le_array_size}⟩ := s
|
||
simp only [size, ge_iff_le]
|
||
apply Nat.le_trans (Nat.sub_le stop start)
|
||
assumption
|
||
|
||
/--
|
||
Extracts an element from the subarray.
|
||
|
||
The index is relative to the start of the subarray, rather than the underlying array.
|
||
-/
|
||
def get (s : Subarray α) (i : Fin s.size) : α :=
|
||
have : s.start + i.val < s.array.size := by
|
||
apply Nat.lt_of_lt_of_le _ s.stop_le_array_size
|
||
have := i.isLt
|
||
simp only [size] at this
|
||
rw [Nat.add_comm]
|
||
exact Nat.add_lt_of_lt_sub this
|
||
s.array[s.start + i.val]
|
||
|
||
instance : GetElem (Subarray α) Nat α fun xs i => i < xs.size where
|
||
getElem xs i h := xs.get ⟨i, h⟩
|
||
|
||
/--
|
||
Extracts an element from the subarray, or returns a default value `v₀` when the index is out of
|
||
bounds.
|
||
|
||
The index is relative to the start and end of the subarray, rather than the underlying array.
|
||
-/
|
||
@[inline] def getD (s : Subarray α) (i : Nat) (v₀ : α) : α :=
|
||
if h : i < s.size then s[i] else v₀
|
||
|
||
/--
|
||
Extracts an element from the subarray, or returns a default value when the index is out of bounds.
|
||
|
||
The index is relative to the start and end of the subarray, rather than the underlying array. The
|
||
default value is that provided by the `Inhabited α` instance.
|
||
-/
|
||
abbrev get! [Inhabited α] (s : Subarray α) (i : Nat) : α :=
|
||
getD s i default
|
||
|
||
/--
|
||
Shrinks the subarray by incrementing its starting index if possible, returning it unchanged if not.
|
||
|
||
Examples:
|
||
* `#[1,2,3].toSubarray.popFront.toArray = #[2, 3]`
|
||
* `#[1,2,3].toSubarray.popFront.popFront.toArray = #[3]`
|
||
* `#[1,2,3].toSubarray.popFront.popFront.popFront.toArray = #[]`
|
||
* `#[1,2,3].toSubarray.popFront.popFront.popFront.popFront.toArray = #[]`
|
||
-/
|
||
def popFront (s : Subarray α) : Subarray α :=
|
||
if h : s.start < s.stop then
|
||
⟨{ s.internalRepresentation with
|
||
start := s.start + 1,
|
||
start_le_stop := Nat.le_of_lt_succ (Nat.add_lt_add_right h 1) }⟩
|
||
else
|
||
s
|
||
|
||
/--
|
||
The empty subarray.
|
||
|
||
This empty subarray is backed by an empty array.
|
||
-/
|
||
protected def empty : Subarray α := ⟨{
|
||
array := #[]
|
||
start := 0
|
||
stop := 0
|
||
start_le_stop := Nat.le_refl 0
|
||
stop_le_array_size := Nat.le_refl 0
|
||
}⟩
|
||
|
||
instance : EmptyCollection (Subarray α) :=
|
||
⟨Subarray.empty⟩
|
||
|
||
instance : Inhabited (Subarray α) :=
|
||
⟨{}⟩
|
||
|
||
/-!
|
||
`ForIn`, `foldlM`, `foldl` and other operations are implemented in `Init.Data.Slice.Array.Iterator`
|
||
using the slice iterator.
|
||
-/
|
||
|
||
/--
|
||
Folds a monadic operation from right to left over the elements in a subarray.
|
||
|
||
An accumulator of type `β` is constructed by starting with `init` and monadically combining each
|
||
element of the subarray with the current accumulator value in turn, moving from the end to the
|
||
start. The monad in question may permit early termination or repetition.
|
||
|
||
Examples:
|
||
```lean example
|
||
#eval #["red", "green", "blue"].toSubarray.foldrM (init := "") fun x acc => do
|
||
let l ← Option.guard (· ≠ 0) x.length
|
||
return s!"{acc}({l}){x} "
|
||
```
|
||
```output
|
||
some "(4)blue (5)green (3)red "
|
||
```
|
||
```lean example
|
||
#eval #["red", "green", "blue"].toSubarray.foldrM (init := 0) fun x acc => do
|
||
let l ← Option.guard (· ≠ 5) x.length
|
||
return s!"{acc}({l}){x} "
|
||
```
|
||
```output
|
||
none
|
||
```
|
||
-/
|
||
@[inline]
|
||
def foldrM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α → β → m β) (init : β) (as : Subarray α) : m β :=
|
||
as.array.foldrM f (init := init) (start := as.stop) (stop := as.start)
|
||
|
||
/--
|
||
Checks whether any of the elements in a subarray satisfy a monadic Boolean predicate.
|
||
|
||
The elements are tested starting at the lowest index and moving up. The search terminates as soon as
|
||
an element that satisfies the predicate is found.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #["red", "green", "blue", "orange"].toSubarray.popFront.anyM fun x => do
|
||
IO.println x
|
||
pure (x == "blue")
|
||
```
|
||
```output
|
||
green
|
||
blue
|
||
```
|
||
```output
|
||
true
|
||
```
|
||
-/
|
||
@[inline]
|
||
def anyM {α : Type u} {m : Type → Type w} [Monad m] (p : α → m Bool) (as : Subarray α) : m Bool :=
|
||
as.array.anyM p (start := as.start) (stop := as.stop)
|
||
|
||
/--
|
||
Checks whether all of the elements in a subarray satisfy a monadic Boolean predicate.
|
||
|
||
The elements are tested starting at the lowest index and moving up. The search terminates as soon as
|
||
an element that does not satisfy the predicate is found.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #["red", "green", "blue", "orange"].toSubarray.popFront.allM fun x => do
|
||
IO.println x
|
||
pure (x.length == 5)
|
||
```
|
||
```output
|
||
green
|
||
blue
|
||
```
|
||
```output
|
||
false
|
||
```
|
||
-/
|
||
@[inline]
|
||
def allM {α : Type u} {m : Type → Type w} [Monad m] (p : α → m Bool) (as : Subarray α) : m Bool :=
|
||
as.array.allM p (start := as.start) (stop := as.stop)
|
||
|
||
/--
|
||
Runs a monadic action on each element of a subarray.
|
||
|
||
The elements are processed starting at the lowest index and moving up.
|
||
-/
|
||
@[inline]
|
||
def forM {α : Type u} {m : Type v → Type w} [Monad m] (f : α → m PUnit) (as : Subarray α) : m PUnit :=
|
||
as.array.forM f (start := as.start) (stop := as.stop)
|
||
|
||
/--
|
||
Runs a monadic action on each element of a subarray, in reverse order.
|
||
|
||
The elements are processed starting at the highest index and moving down.
|
||
-/
|
||
@[inline]
|
||
def forRevM {α : Type u} {m : Type v → Type w} [Monad m] (f : α → m PUnit) (as : Subarray α) : m PUnit :=
|
||
as.array.forRevM f (start := as.stop) (stop := as.start)
|
||
|
||
/--
|
||
Folds an operation from right to left over the elements in a subarray.
|
||
|
||
An accumulator of type `β` is constructed by starting with `init` and combining each element of the
|
||
subarray with the current accumulator value in turn, moving from the end to the start.
|
||
|
||
Examples:
|
||
* `#["red", "green", "blue"].toSubarray.foldr (·.length + ·) 0 = 12`
|
||
* `#["red", "green", "blue"].toSubarray.popFront.foldr (·.length + ·) 0 = 9`
|
||
-/
|
||
@[inline]
|
||
def foldr {α : Type u} {β : Type v} (f : α → β → β) (init : β) (as : Subarray α) : β :=
|
||
Id.run <| as.foldrM (pure <| f · ·) (init := init)
|
||
|
||
/--
|
||
Checks whether any of the elements in a subarray satisfy a Boolean predicate.
|
||
|
||
The elements are tested starting at the lowest index and moving up. The search terminates as soon as
|
||
an element that satisfies the predicate is found.
|
||
-/
|
||
@[inline]
|
||
def any {α : Type u} (p : α → Bool) (as : Subarray α) : Bool :=
|
||
Id.run <| as.anyM (pure <| p ·)
|
||
|
||
/--
|
||
Checks whether all of the elements in a subarray satisfy a Boolean predicate.
|
||
|
||
The elements are tested starting at the lowest index and moving up. The search terminates as soon as
|
||
an element that does not satisfy the predicate is found.
|
||
-/
|
||
@[inline]
|
||
def all {α : Type u} (p : α → Bool) (as : Subarray α) : Bool :=
|
||
Id.run <| as.allM (pure <| p ·)
|
||
|
||
/--
|
||
Applies a monadic function to each element in a subarray in reverse order, stopping at the first
|
||
element for which the function succeeds by returning a value other than `none`. The succeeding value
|
||
is returned, or `none` if there is no success.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #["red", "green", "blue"].toSubarray.findSomeRevM? fun x => do
|
||
IO.println x
|
||
return Option.guard (· = 5) x.length
|
||
```
|
||
```output
|
||
blue
|
||
green
|
||
```
|
||
```output
|
||
some 5
|
||
```
|
||
-/
|
||
@[inline]
|
||
def findSomeRevM? {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Subarray α) (f : α → m (Option β)) : 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 _)
|
||
|
||
/--
|
||
Applies a monadic Boolean predicate to each element in a subarray in reverse order, stopping at the
|
||
first element that satisfies the predicate. The element that satisfies the predicate is returned, or
|
||
`none` if no element satisfies it.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #["red", "green", "blue"].toSubarray.findRevM? fun x => do
|
||
IO.println x
|
||
return (x.length = 5)
|
||
```
|
||
```output
|
||
blue
|
||
green
|
||
```
|
||
```output
|
||
some 5
|
||
```
|
||
-/
|
||
@[inline]
|
||
def findRevM? {α : Type} {m : Type → Type w} [Monad m] (as : Subarray α) (p : α → m Bool) : m (Option α) :=
|
||
as.findSomeRevM? fun a => return if (← p a) then some a else none
|
||
|
||
/--
|
||
Tests each element in a subarray with a Boolean predicate in reverse order, stopping at the first
|
||
element that satisfies the predicate. The element that satisfies the predicate is returned, or
|
||
`none` if no element satisfies the predicate.
|
||
|
||
Examples:
|
||
* `#["red", "green", "blue"].toSubarray.findRev? (·.length ≠ 4) = some "green"`
|
||
* `#["red", "green", "blue"].toSubarray.findRev? (fun _ => true) = some "blue"`
|
||
* `#["red", "green", "blue"].toSubarray 0 0 |>.findRev? (fun _ => true) = none`
|
||
-/
|
||
@[inline]
|
||
def findRev? {α : Type} (as : Subarray α) (p : α → Bool) : Option α :=
|
||
Id.run <| as.findRevM? (pure <| p ·)
|
||
|
||
end Subarray
|
||
|
||
namespace Array
|
||
variable {α : Type u}
|
||
|
||
/--
|
||
Returns a subarray of an array, with the given bounds.
|
||
|
||
If `start` or `stop` are not valid bounds for a subarray, then they are clamped to array's size.
|
||
Additionally, the starting index is clamped to the ending index.
|
||
-/
|
||
def toSubarray (as : Array α) (start : Nat := 0) (stop : Nat := as.size) : Subarray α :=
|
||
if h₂ : stop ≤ as.size then
|
||
if h₁ : start ≤ stop then
|
||
⟨{ array := as, start := start, stop := stop,
|
||
start_le_stop := h₁, stop_le_array_size := h₂ }⟩
|
||
else
|
||
⟨{ array := as, start := stop, stop := stop,
|
||
start_le_stop := Nat.le_refl _, stop_le_array_size := h₂ }⟩
|
||
else
|
||
if h₁ : start ≤ as.size then
|
||
⟨{ array := as,
|
||
start := start,
|
||
stop := as.size,
|
||
start_le_stop := h₁,
|
||
stop_le_array_size := Nat.le_refl _ }⟩
|
||
else
|
||
⟨{ array := as,
|
||
start := as.size,
|
||
stop := as.size,
|
||
start_le_stop := Nat.le_refl _,
|
||
stop_le_array_size := Nat.le_refl _ }⟩
|
||
|
||
/-- A subarray with the provided bounds.-/
|
||
syntax:max term noWs "[" withoutPosition(term ":" term) "]" : term
|
||
/-- A subarray with the provided lower bound that extends to the rest of the array. -/
|
||
syntax:max term noWs "[" withoutPosition(term ":") "]" : term
|
||
/-- A subarray with the provided upper bound, starting at the index 0. -/
|
||
syntax:max term noWs "[" withoutPosition(":" term) "]" : term
|
||
|
||
macro_rules
|
||
| `($a[$start : $stop]) => `(Array.toSubarray $a $start $stop)
|
||
| `($a[ : $stop]) => `(Array.toSubarray $a 0 $stop)
|
||
| `($a[$start : ]) => `(let a := $a; Array.toSubarray a $start a.size)
|
||
|
||
end Array
|