-
-
Save JasonGross/d852d6ca1bd7f31737edb2334164f613 to your computer and use it in GitHub Desktop.
A type class for generic Ranges
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import Init.Meta | |
-- This isn't strictly necessary of course - reverse could be a member of `Range`. | |
class Reverse (α : Type u) where | |
reverse : α → α | |
open Reverse (reverse) | |
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- | |
-- Range class definition | |
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- | |
/-- A finite range of values describable by some `[START:STOP:STEP]`. -/ | |
class Range (rng : Type u) (val : outParam (Type v)) extends Reverse rng : Type (max u v) where | |
/-- Number of values in the range. -/ | |
steps : rng → Nat | |
/-- Returns the value in the range after the given value. If the given value is not | |
in the range or if it is the last value in the range, the returned value is arbitrary. -/ | |
next : rng → val → val | |
/-- The first value in the range, i.e., either the lowest or highest value depending on | |
the ranges ordering. If the range is empty, the returned value is arbitrary. -/ | |
first : rng → val | |
-- AMK: I think there's an argument that `next` and `first` should be | |
-- `next!` and `first!` since they have non-trivial unchecked constraints | |
-- and will return garbage values if misused... ¯\_(ツ)_/¯ | |
namespace Range | |
universes u v | |
@[inline] def forIn [Range ρ α] {β : Type u} {m : Type u → Type v} [Monad m] (range : ρ) (init : β) (f : α → β → m (ForInStep β)) : m β := | |
let rec @[specialize] loop (n : Nat) (a : α) (b : β) : m β := do | |
match n with | |
| 0 => pure b | |
| n+1 => match ← f a b with | |
| ForInStep.done b => pure b | |
| ForInStep.yield b => loop n (next range a) b | |
loop (steps range) (first range) init | |
instance [Range ρ α] : ForIn m ρ α where | |
forIn := forIn | |
@[inline] protected def forM [Range ρ α] {m : Type u → Type v} [Monad m] (range : ρ) (f : α → m PUnit) : m PUnit := | |
let rec @[specialize] loop (n : Nat) (a : α) : m PUnit := do | |
match n with | |
| 0 => pure ⟨⟩ | |
| n+1 => f a; loop n (next range a) | |
loop (steps range) (first range) | |
def toArray [Range ρ α] (r : ρ) : Array α := do | |
let mut arr := #[] | |
for a in r do | |
arr := arr.push a | |
arr | |
-- N.B., not all types have inherent upper and/or lower bounds, so there | |
-- is a typeclass for each way constructing a range might be supported. | |
-- The class names... may not be the best. I couldn't decide. | |
/-- Class for creating a range from an inclusive start and exlusive stop. -/ | |
class FromStartStop (val : Type u) (rng : outParam (Type u)) : Type u where | |
rangeWithStep : Nat → val → val → rng | |
range : val → val → rng := rangeWithStep 1 | |
/-- Class for creating a range from just an inclusive start, i.e. `[START:]` | |
or `[START: : STEP]`. -/ | |
class FromStart (val : Type u) (rng : outParam (Type u)) : Type u where | |
rangeWithStep : Nat → val → rng | |
range : val → rng := rangeWithStep 1 | |
/-- Class for creating a range from just an exlusive stop, i.e. `[:STOP]` | |
or `[:STOP: STEP]`. -/ | |
class FromStop (val : Type u) (rng : outParam (Type u)) : Type u where | |
rangeWithStep : Nat → val → rng | |
range : val → rng := rangeWithStep 1 | |
-- `%` prefix to avoid collision with builtin Range syntax | |
syntax:max "%[" term ":" "]" : term | |
syntax:max "%[" ":" term "]" : term | |
syntax:max "%[" term ":" term "]" : term | |
syntax:max "%[" term ":" ":" term "]" : term | |
syntax:max "%[" ":" term ":" term "]" : term | |
syntax:max "%[" term ":" term ":" term "]" : term | |
macro_rules | |
| `(%[$start : ]) => `(FromStart.range $start) | |
| `(%[ : $stop]) => `(FromStop.range $stop) | |
| `(%[ $start : $stop ]) => `(FromStartStop.range $start $stop) | |
| `(%[$start : : $step]) => `(FromStart.rangeWithStep $step $start) | |
| `(%[ : $stop : $step ]) => `(FromStop.rangeWithStep $step $stop) | |
| `(%[ $start : $stop : $step ]) => `(FromStartStop.rangeWithStep $step $start $stop) | |
end Range | |
structure RangeStream (ρ α : Type u) [Range ρ α] where | |
steps : Nat | |
first : α | |
range : ρ | |
instance [Range ρ α] : ToStream ρ (RangeStream ρ α) where | |
toStream r := ⟨(Range.steps r), (Range.first r), r⟩ | |
instance [Range ρ α] : Stream (RangeStream ρ α) α where | |
next? | ⟨n, x, xs⟩ => | |
match n with | |
| 0 => none | |
| n+1 => (x, ⟨n, Range.next xs x, xs⟩) | |
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- | |
-- Nat Ranges | |
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- | |
namespace Nat | |
structure NatRange where | |
start : Nat := 0 | |
stop : Nat | |
step : Nat := 1 | |
ascending : Bool := true | |
namespace NatRange | |
@[inline] def steps (xs : NatRange) : Nat := | |
let size := xs.stop - xs.start | |
let step := xs.step | |
size / step + (if size % step = 0 then 0 else 1) | |
instance : Range NatRange Nat where | |
steps := steps | |
next xs x := | |
if xs.ascending | |
then x + xs.step | |
else x - xs.step | |
first xs := | |
if xs.ascending | |
then xs.start | |
else xs.start + (xs.step * (xs.steps - 1)) | |
reverse xs := | |
{xs with ascending := !xs.ascending} | |
instance : Range.FromStop Nat NatRange where | |
rangeWithStep step stop := {step := step, stop := stop} | |
instance : Range.FromStartStop Nat NatRange where | |
rangeWithStep step start stop := {step := step, start := start, stop := stop} | |
end NatRange | |
end Nat | |
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- | |
-- Integer Ranges | |
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- | |
namespace Int | |
structure IntRange where | |
start : Int | |
stop : Int | |
step : Nat := 1 | |
ascending : Bool := true | |
namespace IntRange | |
@[inline] def steps (xs : IntRange) : Nat := | |
let size := match xs.stop - xs.start with | |
| ofNat n => n | |
| _ => 0 | |
let step := xs.step | |
size / step + (if size % step = 0 then 0 else 1) | |
instance : Range IntRange Int where | |
steps := steps | |
next xs x := | |
if xs.ascending | |
then x + xs.step | |
else x - xs.step | |
first xs := | |
if xs.ascending | |
then xs.start | |
else xs.start + (xs.step * (xs.steps - 1)) | |
reverse r := | |
{r with ascending := !r.ascending} | |
instance : Range.FromStartStop Int IntRange where | |
rangeWithStep step start stop := {step := step, start := start, stop := stop} | |
end IntRange | |
end Int | |
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- | |
-- Fin Ranges | |
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- | |
-- Lemma to help with FinRange implementation | |
theorem Nat.ltAddRight (a b c : Nat) (h : a < b) : a < b + c := | |
match c with | |
| 0 => h | |
| c'+1 => | |
have h₁ : a < (b + c') + 1 := Nat.lt.step (ltAddRight a b c' h) | |
have h₂ : (b + c') + 1 = b + (c' + 1) from Nat.add_assoc .. | |
cast (congrArg (λ x => a < x) h₂) h₁ | |
namespace Fin | |
structure FinRange (n : Nat) where | |
start : Fin n | |
stop : Fin (n+1) | |
step : Nat := 1 | |
ascending : Bool := true | |
-- Fin helpers for FinRange implementation | |
def nonZero : {n : Nat} → Fin n → n > 0 | |
| 0, silly => Fin.elim0 silly | |
| n+1, _ => Nat.succPos n | |
def lift {n : Nat} (m : Nat) : Fin n → Fin (n+m) | |
| ⟨v, h⟩ => ⟨v, Nat.ltAddRight v n m h⟩ | |
namespace FinRange | |
@[inline] def steps (xs : FinRange n) : Nat := | |
let size := xs.stop.val - xs.start.val | |
let step := xs.step | |
size / step + (if size % step = 0 then 0 else 1) | |
instance : Range (FinRange n) (Fin n) where | |
steps := steps | |
next r x := | |
if r.ascending | |
then x + (Fin.ofNat' r.step (nonZero x)) | |
else x - (Fin.ofNat' r.step (nonZero x)) | |
first xs := | |
if xs.ascending | |
then xs.start | |
else (Fin.ofNat' (xs.start.val + (xs.step * (xs.steps - 1))) (nonZero xs.start)) | |
reverse xs := {xs with ascending := !xs.ascending} | |
instance : Range.FromStart (Fin n) (FinRange n) where | |
rangeWithStep step start := | |
{step := step, | |
start := start, | |
stop := Fin.ofNat n} | |
instance : Range.FromStop (Fin n) (FinRange n) where | |
rangeWithStep step stop := | |
{step := step, | |
start := Fin.ofNat' 0 (nonZero stop), | |
stop := stop.lift 1} | |
instance : Range.FromStartStop (Fin n) (FinRange n) where | |
rangeWithStep step start stop := | |
{step := step, | |
start := start, | |
stop := stop.lift 1} | |
end FinRange | |
end Fin | |
-- Nat range examples | |
#eval Range.toArray %[:5] | |
-- #[0, 1, 2, 3, 4] | |
#eval Range.toArray $ reverse %[:5] | |
-- #[4, 3, 2, 1, 0] | |
-- Int range examples | |
#eval Range.toArray %[0:(5 : Int)] | |
-- #[0, 1, 2, 3, 4] | |
#eval Range.toArray $ %[-5:5] | |
-- #[-5, -4, -3, -2, -1, 0, 1, 2, 3, 4] | |
#eval Range.toArray $ reverse %[-5:5] | |
-- #[4, 3, 2, 1, 0, -1, -2, -3, -4, -5] | |
-- Fin range examples | |
#eval Range.toArray %[:(5 : Fin 6)] | |
-- #[0, 1, 2, 3, 4] | |
#eval Range.toArray %[(5 : Fin 6):] | |
-- #[5] | |
#eval Range.toArray $ %[(0 : Fin 6):] | |
-- #[0, 1, 2, 3, 4, 5] | |
#eval Range.toArray $ reverse %[(0 : Fin 6):] | |
-- #[5, 4, 3, 2, 1, 0] | |
def Array.sum (a₁ a₂ : Array Int) (pf : a₁.size = a₂.size) : Array Int := | |
if h : a₁.size > 0 then do | |
let mut res := #[] | |
for i in %[Fin.ofNat' 0 h:] do | |
res := res.push (a₁.get i + a₂.get (cast (congrArg Fin pf) i)) | |
res | |
else | |
#[] | |
#eval Array.sum #[1,2,3] #[4,5,6] rfl | |
-- #[5, 7, 9] | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment