module Stdlib.Data.Range;

import Stdlib.Data.Fixity open;

import Stdlib.Data.Bool.Base open;
import Stdlib.Trait.Natural open;
import Stdlib.Trait.Ord open;
import Stdlib.Trait.Foldable open;
import Stdlib.Data.Nat open;

--- An inclusive range of naturals
type Range N :=
  mkRange {
    low : N;
    high : N
  };

type StepRange N :=
  mkStepRange {
    range : Range N;
    step : N
  };

syntax operator to range;

--- `x to y` is the range [x..y]
{-# inline: always #-}
to {N} {{Natural N}} (l h : N) : Range N := mkRange l h;

syntax operator step step;

--- `x to y step s` is the range [x, x + s, ..., y]
{-# inline: always #-}
step {N} (r : Range N) (s : N) : StepRange N := mkStepRange r s;

-- This instance assumes that `low <= high`.
{-# specialize: true, inline: case #-}
instance
foldableRangeI {N} {{Eq N}} {{Natural N}} : Foldable (Range N) N :=
  mkFoldable@{
    {-# specialize: [1, 3] #-}
    for {B : Type} (f : B -> N -> B) (ini : B) : Range N -> B
      | (mkRange low high) :=
        let
          {-# specialize-by: [f, high] #-}
          terminating
          go (acc : B) (next : N) : B :=
            if 
              | next == high := f acc next
              | else := go (f acc next) (next + 1);
        in go ini low;
    {-# specialize: [1, 3] #-}
    rfor {B : Type} (f : B -> N -> B) (ini : B) : Range N -> B
      | (mkRange low high) :=
        let
          {-# specialize-by: [f, high] #-}
          terminating
          go (next : N) : B :=
            if 
              | next == high := f ini next
              | else := f (go (next + 1)) next;
        in go low
  };

-- This instance assumes that (low + step*k > high) for some k.
{-# specialize: true, inline: case #-}
instance
foldableStepRangeI {N} {{Ord N}} {{Natural N}} : Foldable (StepRange N) N :=
  mkFoldable@{
    {-# specialize: [1, 3] #-}
    for {B : Type} (f : B -> N -> B) (ini : B) : StepRange N -> B
      | (mkStepRange (mkRange low high) step) :=
        let
          {-# specialize-by: [f, high, step] #-}
          terminating
          go (acc : B) (next : N) : B :=
            if 
              | next > high := acc
              | else := go (f acc next) (next + step);
        in go ini low;
    {-# specialize: [1, 3] #-}
    rfor {B : Type} (f : B -> N -> B) (ini : B) : StepRange N -> B
      | (mkStepRange (mkRange low high) step) :=
        let
          {-# specialize-by: [f, high, step] #-}
          terminating
          go (next : N) : B :=
            if 
              | next <= high := f (go (next + step)) next
              | else := ini;
        in go low
  };