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 :=
  mk@{
    low : N;
    high : N;
  };

type StepRange N :=
  mk@{
    range : Range N;
    step : N;
  };

syntax operator to range;

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

syntax operator step step;

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

-- This instance assumes that `low <= high`.
{-# specialize: true, inline: case #-}
instance
foldableRangeI {N} {{Eq N}} {{Natural N}} : Foldable (Range N) N :=
  Foldable.mk@{
    {-# specialize: [1, 3] #-}
    for {B : Type} (fun : B -> N -> B) (initialValue : B) : Range N -> B
      | (Range.mk low high) :=
        let
          {-# specialize-by: [fun, high] #-}
          terminating
          go (acc : B) (next : N) : B :=
            if 
              | next == high := fun acc next
              | else := go (fun acc next) (next + 1);
        in go initialValue low;
    {-# specialize: [1, 3] #-}
    rfor {B : Type} (fun : B -> N -> B) (initialValue : B) : Range N -> B
      | (Range.mk low high) :=
        let
          {-# specialize-by: [fun, high] #-}
          terminating
          go (next : N) : B :=
            if 
              | next == high := fun initialValue next
              | else := fun (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 :=
  Foldable.mk@{
    {-# specialize: [1, 3] #-}
    for {B : Type} (fun : B -> N -> B) (initialValue : B) : StepRange N -> B
      | (StepRange.mk (Range.mk low high) step) :=
        let
          {-# specialize-by: [fun, high, step] #-}
          terminating
          go (acc : B) (next : N) : B :=
            if 
              | next > high := acc
              | else := go (fun acc next) (next + step);
        in go initialValue low;
    {-# specialize: [1, 3] #-}
    rfor {B : Type} (fun : B -> N -> B) (initialValue : B) : StepRange N -> B
      | (StepRange.mk (Range.mk low high) step) :=
        let
          {-# specialize-by: [fun, high, step] #-}
          terminating
          go (next : N) : B :=
            if 
              | next <= high := fun (go (next + step)) next
              | else := initialValue;
        in go low;
  };