A type class for generic Ranges

RangeAsClass.lean

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]

N.B., there is a slight tension at play here in the design space w.r.t. whether or not to have [START:STOP:STEP] treat STOP as an exclusive bound: if STOP is exclusive, then empty ranges can be specified using the syntax (yay!), but finite range-types cannot both include the highest value in the type and use a STOP with the same underlying type (since by definition, STOP would need to be MAX_VALUE+1 which would no longer be in the finite numerical type).

An alternative is to have [START:STOP] be inclusive on both START and STOP, but now this syntactic form cannot always specify empty ranges, since, e.g., [0,0] now includes 0. One solution to this is to have a syntactic form for ranges with an inclusive STOP and one for ranges with an exclusive STOP. E.g., in Cryptol [START .. STOP] is used to specify ranges with an inclusive STOP, while [START ..< STOP] represents ranges with an exclusive STOP. There would still be the challenge for some types that [START ..< STOP] may end up having slightly different types for START and STOP. In particular, for Fin n ranges, START : Fin n and STOP : Fin (n+1).

Another alternative I suppose is have a type class for heterogeneous ranges that is able to capture some of the difficulties in bounded numeric types. E.g., perhaps [START:STOP] where START : Fin n and STOP : Fin (n+1) is recognized explicitly as being the way to specify a range of Fin n values... I have not thought deeply about the consequences and possible complications of this "heterogenous range" approach.

Anyway, in the above implemented design, I tried to stick as close to Lean's current [START:STOP:STEP] semantics as possible by hiding some of the details in each range's typeclass' implementation and allowing some types to simply omit the START or the STOP. E.g., for Fin n, omitting STOP lets us sneakily insert Fin.ofNat n : Fin (n+1) in for STOP behind the scenes where the subtle details regarding the type of START and STOP are hidden. There are probably other reasonable designs as well, perhaps that are more preferable. Anyway... that's all I've got to say about that.

There is also possible use of mathematical notation for half-open ranges, e.g., [START..STOP) for an exclusive upper bound or [START, START+n .. STOP) for inferring the STEP. I'm not sure how well editors will do with the mismatched grouping symbols though.