# Range — inclusive/exclusive finite numeric ranges for Bend 2.0.2.
#
# Base has no numeric typeclass mechanism, so the public API is deliberately
# split into Range.Nat and Range.U32. Both use the same Bound values and the
# same endpoint semantics. No termination escapes are used.
import Base

type Range.Bound is Data:
  Inclusive{}
  Exclusive{}

type Range.NatRange is Data:
  NR{lo: Nat, lo_bound: Range.Bound, hi: Nat, hi_bound: Range.Bound}

type Range.U32Range is Data:
  UR{lo: U32, lo_bound: Range.Bound, hi: U32, hi_bound: Range.Bound}

def Range.Bound.is_inclusive(b: Range.Bound) -> Bool:
  match b:
    case Inclusive{}:
      True{}
    case Exclusive{}:
      False{}

def Range.Bound.intersection(a: Range.Bound, b: Range.Bound) -> Range.Bound:
  match a b:
    case Inclusive{} Inclusive{}:
      Inclusive{}
    case _ _:
      Exclusive{}

def Range.Bound.exclusions(a: Range.Bound, b: Range.Bound) -> Nat:
  match a b:
    case Inclusive{} Inclusive{}:
      0n
    case _ Inclusive{}:
      1n
    case Inclusive{} _:
      1n
    case Exclusive{} Exclusive{}:
      2n

def Range.Bound.contains_nat.lower(
  b: Range.Bound, x: Nat, lo: Nat
) -> Bool:
  match b:
    case Inclusive{}:
      Nat.is_ge(x, lo)
    case Exclusive{}:
      Nat.is_gt(x, lo)

def Range.Bound.contains_nat.upper(
  b: Range.Bound, x: Nat, hi: Nat
) -> Bool:
  match b:
    case Inclusive{}:
      Nat.is_le(x, hi)
    case Exclusive{}:
      Nat.is_lt(x, hi)

def Range.Bound.contains_u32.lower(
  b: Range.Bound, x: U32, lo: U32
) -> Bool:
  match b:
    case Inclusive{}:
      U32.is_ge(x, lo)
    case Exclusive{}:
      U32.is_gt(x, lo)

def Range.Bound.contains_u32.upper(
  b: Range.Bound, x: U32, hi: U32
) -> Bool:
  match b:
    case Inclusive{}:
      U32.is_le(x, hi)
    case Exclusive{}:
      U32.is_lt(x, hi)

# Constructors. Range.make is the Nat spelling; Bend has no overloads, so
# U32 callers use Range.make_u32 or the namespaced constructors.
def Range.Nat.make(
  lo: Nat, lo_bound: Range.Bound, hi: Nat, hi_bound: Range.Bound
) -> Range.NatRange:
  NR{lo, lo_bound, hi, hi_bound}

def Range.U32.make(
  lo: U32, lo_bound: Range.Bound, hi: U32, hi_bound: Range.Bound
) -> Range.U32Range:
  UR{lo, lo_bound, hi, hi_bound}

def Range.make(
  lo: Nat, lo_bound: Range.Bound, hi: Nat, hi_bound: Range.Bound
) -> Range.NatRange:
  Range.Nat.make(lo, lo_bound, hi, hi_bound)

def Range.make_u32(
  lo: U32, lo_bound: Range.Bound, hi: U32, hi_bound: Range.Bound
) -> Range.U32Range:
  Range.U32.make(lo, lo_bound, hi, hi_bound)

def Range.Nat.closed(lo: Nat, hi: Nat) -> Range.NatRange:
  Range.Nat.make(lo, Inclusive{}, hi, Inclusive{})

def Range.Nat.open(lo: Nat, hi: Nat) -> Range.NatRange:
  Range.Nat.make(lo, Exclusive{}, hi, Exclusive{})

def Range.Nat.closed_open(lo: Nat, hi: Nat) -> Range.NatRange:
  Range.Nat.make(lo, Inclusive{}, hi, Exclusive{})

def Range.Nat.open_closed(lo: Nat, hi: Nat) -> Range.NatRange:
  Range.Nat.make(lo, Exclusive{}, hi, Inclusive{})

def Range.U32.closed(lo: U32, hi: U32) -> Range.U32Range:
  Range.U32.make(lo, Inclusive{}, hi, Inclusive{})

def Range.U32.open(lo: U32, hi: U32) -> Range.U32Range:
  Range.U32.make(lo, Exclusive{}, hi, Exclusive{})

def Range.U32.closed_open(lo: U32, hi: U32) -> Range.U32Range:
  Range.U32.make(lo, Inclusive{}, hi, Exclusive{})

def Range.U32.open_closed(lo: U32, hi: U32) -> Range.U32Range:
  Range.U32.make(lo, Exclusive{}, hi, Inclusive{})

# Nat range queries.
def Range.Nat.is_empty.cmp(cmp: Cmp, both_inclusive: Bool) -> Bool:
  match cmp:
    case LT{}:
      False{}
    case GT{}:
      True{}
    case EQ{}:
      Bool.not(both_inclusive)

def Range.Nat.is_empty(r: Range.NatRange) -> Bool:
  match r:
    case NR{lo, lo_bound, hi, hi_bound}:
      Range.Nat.is_empty.cmp(
        Nat.cmp(lo, hi),
        Bool.and(
          Range.Bound.is_inclusive(lo_bound),
          Range.Bound.is_inclusive(hi_bound)))

def Range.Nat.contains(r: Range.NatRange, +x: Nat) -> Bool:
  match r:
    case NR{lo, lo_bound, hi, hi_bound}:
      Bool.and(
        Range.Bound.contains_nat.lower(lo_bound, x, lo),
        Range.Bound.contains_nat.upper(hi_bound, x, hi))

def Range.Nat.length.cmp(cmp: Cmp, exclusions: Nat, n: Nat) -> Nat:
  match cmp:
    case LT{}:
      Nat.sub(n, exclusions)
    case EQ{}:
      Bool.pick(Nat, Nat.is_eq(exclusions, 0n), n, 0n)
    case GT{}:
      0n

def Range.Nat.length(r: Range.NatRange) -> Nat:
  match r:
    case NR{+lo, lo_bound, +hi, hi_bound}:
      Range.Nat.length.cmp(
        Nat.cmp(lo, hi),
        Range.Bound.exclusions(lo_bound, hi_bound),
        Nat.add(1n, Nat.sub(hi, lo)))

def Range.Nat.lower_choice.cmp(
  cmp: Cmp, alo: Nat, alb: Range.Bound, blo: Nat, blb: Range.Bound
) -> Nat & Range.Bound:
  match cmp:
    case LT{}:
      (blo, blb)
    case GT{}:
      (alo, alb)
    case EQ{}:
      (alo, Range.Bound.intersection(alb, blb))

def Range.Nat.lower_choice(
  +alo: Nat, alb: Range.Bound, +blo: Nat, blb: Range.Bound
) -> Nat & Range.Bound:
  Range.Nat.lower_choice.cmp(Nat.cmp(alo, blo), alo, alb, blo, blb)

def Range.Nat.upper_choice.cmp(
  cmp: Cmp, ahi: Nat, ahb: Range.Bound, bhi: Nat, bhb: Range.Bound
) -> Nat & Range.Bound:
  match cmp:
    case LT{}:
      (ahi, ahb)
    case GT{}:
      (bhi, bhb)
    case EQ{}:
      (ahi, Range.Bound.intersection(ahb, bhb))

def Range.Nat.upper_choice(
  +ahi: Nat, ahb: Range.Bound, +bhi: Nat, bhb: Range.Bound
) -> Nat & Range.Bound:
  Range.Nat.upper_choice.cmp(Nat.cmp(ahi, bhi), ahi, ahb, bhi, bhb)

def Range.Nat.overlap.finish(
  cmp: Cmp, lo_bound: Range.Bound, hi_bound: Range.Bound
) -> Bool:
  match cmp:
    case LT{}:
      True{}
    case GT{}:
      False{}
    case EQ{}:
      match lo_bound hi_bound:
        case Inclusive{} Inclusive{}:
          True{}
        case _ _:
          False{}

def Range.Nat.overlap.choices(
  low: Nat & Range.Bound, high: Nat & Range.Bound
) -> Bool:
  (lo, lo_bound) = low
  (hi, hi_bound) = high
  Range.Nat.overlap.finish(Nat.cmp(lo, hi), lo_bound, hi_bound)

def Range.Nat.overlap(a: Range.NatRange, b: Range.NatRange) -> Bool:
  match a b:
    case NR{alo, alb, ahi, ahb} NR{blo, blb, bhi, bhb}:
      Range.Nat.overlap.choices(
        Range.Nat.lower_choice(alo, alb, blo, blb),
        Range.Nat.upper_choice(ahi, ahb, bhi, bhb))

def Range.Nat.overlaps(a: Range.NatRange, b: Range.NatRange) -> Bool:
  Range.Nat.overlap(a, b)

def Range.Nat.clamp.cmp(cmp: Cmp, x: Nat, lo: Nat, hi: Nat) -> Nat:
  match cmp:
    case LT{}:
      Nat.min(Nat.max(x, lo), hi)
    case EQ{}:
      lo
    case GT{}:
      lo

def Range.Nat.clamp(r: Range.NatRange, x: Nat) -> Nat:
  match r:
    case NR{+lo, lo_bound, +hi, hi_bound}:
      Range.Nat.clamp.cmp(Nat.cmp(lo, hi), x, lo, hi)

# The recursion is explicitly fueled by the computed length.
def Range.Nat.to_list.go(fuel: Nat, +current: Nat) -> List<&2, Nat>:
  match fuel:
    case 0n:
      Nil{}
    case 1n+f:
      current <> Range.Nat.to_list.go(f, 1n+current)

def Range.Nat.start(lo: Nat, lo_bound: Range.Bound) -> Nat:
  match lo_bound:
    case Inclusive{}:
      lo
    case Exclusive{}:
      1n+lo

def Range.Nat.to_list.parts(
  +lo: Nat, +lo_bound: Range.Bound, +hi: Nat, hi_bound: Range.Bound
) -> List<&2, Nat>:
  fuel = Range.Nat.length.cmp(
    Nat.cmp(lo, hi),
    Range.Bound.exclusions(lo_bound, hi_bound),
    Nat.add(1n, Nat.sub(hi, lo)))
  Range.Nat.to_list.go(fuel, Range.Nat.start(lo, lo_bound))

def Range.Nat.to_list(r: Range.NatRange) -> List<&2, Nat>:
  match r:
    case NR{lo, lo_bound, hi, hi_bound}:
      Range.Nat.to_list.parts(lo, lo_bound, hi, hi_bound)

# U32 range queries. Length is Nat so the full 0..U32.max range is representable.
def Range.U32.is_empty.cmp(cmp: Cmp, both_inclusive: Bool) -> Bool:
  match cmp:
    case LT{}:
      False{}
    case GT{}:
      True{}
    case EQ{}:
      Bool.not(both_inclusive)

def Range.U32.is_empty(r: Range.U32Range) -> Bool:
  match r:
    case UR{lo, lo_bound, hi, hi_bound}:
      Range.U32.is_empty.cmp(
        U32.cmp(lo, hi),
        Bool.and(
          Range.Bound.is_inclusive(lo_bound),
          Range.Bound.is_inclusive(hi_bound)))

def Range.U32.contains(r: Range.U32Range, +x: U32) -> Bool:
  match r:
    case UR{lo, lo_bound, hi, hi_bound}:
      Bool.and(
        Range.Bound.contains_u32.lower(lo_bound, x, lo),
        Range.Bound.contains_u32.upper(hi_bound, x, hi))

def Range.U32.length.cmp(cmp: Cmp, exclusions: Nat, n: Nat) -> Nat:
  match cmp:
    case LT{}:
      Nat.sub(n, exclusions)
    case EQ{}:
      Bool.pick(Nat, Nat.is_eq(exclusions, 0n), n, 0n)
    case GT{}:
      0n

def Range.U32.length(r: Range.U32Range) -> Nat:
  match r:
    case UR{+lo, lo_bound, +hi, hi_bound}:
      Range.U32.length.cmp(
        U32.cmp(lo, hi),
        Range.Bound.exclusions(lo_bound, hi_bound),
        Nat.add(1n, Nat.sub(U32.to_nat(hi), U32.to_nat(lo))))

def Range.U32.lower_choice.cmp(
  cmp: Cmp, alo: U32, alb: Range.Bound, blo: U32, blb: Range.Bound
) -> U32 & Range.Bound:
  match cmp:
    case LT{}:
      (blo, blb)
    case GT{}:
      (alo, alb)
    case EQ{}:
      (alo, Range.Bound.intersection(alb, blb))

def Range.U32.lower_choice(
  +alo: U32, alb: Range.Bound, +blo: U32, blb: Range.Bound
) -> U32 & Range.Bound:
  Range.U32.lower_choice.cmp(U32.cmp(alo, blo), alo, alb, blo, blb)

def Range.U32.upper_choice.cmp(
  cmp: Cmp, ahi: U32, ahb: Range.Bound, bhi: U32, bhb: Range.Bound
) -> U32 & Range.Bound:
  match cmp:
    case LT{}:
      (ahi, ahb)
    case GT{}:
      (bhi, bhb)
    case EQ{}:
      (ahi, Range.Bound.intersection(ahb, bhb))

def Range.U32.upper_choice(
  +ahi: U32, ahb: Range.Bound, +bhi: U32, bhb: Range.Bound
) -> U32 & Range.Bound:
  Range.U32.upper_choice.cmp(U32.cmp(ahi, bhi), ahi, ahb, bhi, bhb)

def Range.U32.overlap.finish(
  cmp: Cmp, lo_bound: Range.Bound, hi_bound: Range.Bound
) -> Bool:
  match cmp:
    case LT{}:
      True{}
    case GT{}:
      False{}
    case EQ{}:
      match lo_bound hi_bound:
        case Inclusive{} Inclusive{}:
          True{}
        case _ _:
          False{}

def Range.U32.overlap.choices(
  low: U32 & Range.Bound, high: U32 & Range.Bound
) -> Bool:
  (lo, lo_bound) = low
  (hi, hi_bound) = high
  Range.U32.overlap.finish(U32.cmp(lo, hi), lo_bound, hi_bound)

def Range.U32.overlap(a: Range.U32Range, b: Range.U32Range) -> Bool:
  match a b:
    case UR{alo, alb, ahi, ahb} UR{blo, blb, bhi, bhb}:
      Range.U32.overlap.choices(
        Range.U32.lower_choice(alo, alb, blo, blb),
        Range.U32.upper_choice(ahi, ahb, bhi, bhb))

def Range.U32.overlaps(a: Range.U32Range, b: Range.U32Range) -> Bool:
  Range.U32.overlap(a, b)

def Range.U32.clamp.cmp(cmp: Cmp, x: U32, lo: U32, hi: U32) -> U32:
  match cmp:
    case LT{}:
      U32.min(U32.max(x, lo), hi)
    case EQ{}:
      lo
    case GT{}:
      lo

def Range.U32.clamp(r: Range.U32Range, x: U32) -> U32:
  match r:
    case UR{+lo, lo_bound, +hi, hi_bound}:
      Range.U32.clamp.cmp(U32.cmp(lo, hi), x, lo, hi)

def Range.U32.to_list.go(fuel: Nat, +current: U32) -> List<&2, U32>:
  match fuel:
    case 0n:
      Nil{}
    case 1n+f:
      current <> Range.U32.to_list.go(f, U32.inc(current))

def Range.U32.start(lo: U32, lo_bound: Range.Bound) -> U32:
  match lo_bound:
    case Inclusive{}:
      lo
    case Exclusive{}:
      U32.inc(lo)

def Range.U32.to_list.parts(
  +lo: U32, +lo_bound: Range.Bound, +hi: U32, hi_bound: Range.Bound
) -> List<&2, U32>:
  fuel = Range.U32.length.cmp(
    U32.cmp(lo, hi),
    Range.Bound.exclusions(lo_bound, hi_bound),
    Nat.add(1n, Nat.sub(U32.to_nat(hi), U32.to_nat(lo))))
  Range.U32.to_list.go(fuel, Range.U32.start(lo, lo_bound))

def Range.U32.to_list(r: Range.U32Range) -> List<&2, U32>:
  match r:
    case UR{lo, lo_bound, hi, hi_bound}:
      Range.U32.to_list.parts(lo, lo_bound, hi, hi_bound)