module Data.Stream
export  { streamOfList; listOfStream
        ; sgenerate; senumFrom; srepeat; scons
        ; smap; smapacc
        ; sfold; sany
        ; stake; stakeWhile; sfilter }
import Data.Numeric.Nat
import Data.List
import Data.Maybe
import Data.Tuple
import Data.Array
import Data.Numeric.Bool
import Data.Function
where


-- | Unbounded streams, 
--   wraps a function that produces elements on demand.
data Stream (s a: Data) where
        MkStream : (s -> Step s a) -> s -> Stream s a

data Step (s a: Data) where
        Yield   : a -> s -> Step s a
        Skip    : s -> Step s a 
        Done    : Step s a


-- Conversions ------------------------------------------------------------------------------------
-- | Convert a list to a stream.
streamOfList (xx: List a): Stream (List a) a
 = let  step (s1: List a)
         = case s1 of
                Nil       -> Done
                Cons x xs -> Yield x xs
   in   MkStream step xx


-- | Convert a stream to a list.
listOfStream  (ss: Stream s a): List a
 = case ss of
        MkStream f s0
         -> case f s0 of
                Yield x s1      -> Cons x (listOfStream (MkStream f s1))
                Skip  s1        -> listOfStream (MkStream f s1)
                Done            -> Nil


-- | Load the given number of elements from a stream and write them
--   into a freshly allocated array.
arrayOfStream 
        [r1: Region] 
        (n: Nat#)        -- ^ Length of result array.
        (d: a)           -- ^ Default element value.
        (ss: Stream s a) -- ^ Stream to evaluate.
        : S (Alloc r1) (Array r1 a)
 = extend r1 using r2 with { Write r2; Alloc r2 } in
   do
        -- Allocate an array of the given maximum size.
        arr     = allocArray [r2] n d

        -- Unstream elements into the array.
        unstreamToArray ss arr 0

        -- Return the completed array.
        arr


-- | Unstream all available elements into the given array.
unstreamToArray
        [r: Region] 
        (ss: Stream s a) (arr: Array r a) (ix: Nat#)
        : S (Write r) Unit
 = case ss of
        MkStream f s0
         -> case f s0 of
                Yield x s1
                 -> do  writeArray arr ix x
                        unstreamToArray (MkStream f s1) arr (ix + 1)

                Skip s1
                 ->     unstreamToArray (MkStream f s1)  arr ix

                Done
                 ->     ()


-- Constructors -----------------------------------------------------------------------------------
-- | Generate a stream, given a starting value and a stepper function.
sgenerate (x: s) (step: s -> Tup2 s a): Stream s a
 = let  step' sA
         = case step sA of
                T2 s' x -> Yield x s'
   in   MkStream step' x

senumFrom (x: Nat#): Stream Nat# Nat#
 = sgenerate x (λs: Nat#. T2 (s + 1) s)


-- | Create a stream that returns copies of the same value.
srepeat (x: a): Stream a a
 = sgenerate x (λs: a. T2 s s)


-- | Cons an element to the front of a stream.
scons (x: a) (ss: Stream s a): Stream (Tup2 s Bool#) a
 = case ss of
        MkStream stepA sA0
         -> let stepA2 q
                 = case q of
                        T2 sA1 b
                         -> case b of
                                True    -> Yield x (T2 sA1 False)
                                False   -> case stepA sA1 of
                                                Yield y sA2 -> Yield y (T2 sA2 False)
                                                Skip  sA2   -> Skip    (T2 sA2 False)
                                                Done        -> Done
            in  MkStream stepA2 (T2 sA0 True)


-- Maps -------------------------------------------------------------------------------------------
-- | Apply a function to every element of a stream.
smap (f: a -> b) (ss: Stream s a): Stream s b
 = case ss of  
        MkStream stepA sA0
         -> let stepB q 
                 = case stepA q of
                        Yield x sA1     -> Yield (f x) sA1
                        Skip sA2        -> Skip sA2
                        Done            -> Done
            in  MkStream stepB sA0


-- Scans ------------------------------------------------------------------------------------------
-- | Like `smap`, but keep a running accumulator as we walk along the stream.
smapacc (f: a -> b -> Tup2 a c) (z: a) 
        (ss: Stream s b): Stream (Tup2 s a) c
 = case ss of
        MkStream fB sB0
         -> let stepC q
                 = case q of
                        T2 sB1 xA1
                         -> case fB sB1 of
                                Yield xB1 sB2 
                                 -> case f xA1 xB1 of
                                        T2 xA2 xC2
                                         -> Yield xC2 (T2 sB2 xA2)
                                Skip sB2 -> Skip      (T2 sB2 xA1)
                                Done     -> Done
            in  MkStream stepC (T2 sB0 z)


-- Folds ------------------------------------------------------------------------------------------
-- | Fold all the elements from a stream.
sfold    (f: a -> b -> a) (acc: a) (ss: Stream s b): a
 = case ss of
        MkStream step s0 
         -> sconsume f acc step s0


sconsume (f:    a -> b -> a)   (acc:   a) 
         (step: s -> Step s b) (state: s) : a
 = case step state of
        Yield x s'      -> sconsume f (f acc x) step s'
        Skip    s'      -> sconsume f acc       step s'
        Done            -> acc


-- | Check if any of the elements of this stream are true,
--   demanding only the prefix of non-true elements from the stream.
sany [s: Data] (ss: Stream s Bool#): Bool#
 = sfold or False $ stakeWhile id ss


-- Projections ------------------------------------------------------------------------------------
-- | Take the given number of elements from a stream.
stake (n: Nat#) (ss: Stream s a): Stream (Tup2 s Nat#) a
 = case ss of
        MkStream fA sA0
         -> let stepB q
                 = case q of
                        T2 sA ix 
                         | ix >= n   -> Done
                         | otherwise 
                         -> case fA sA of
                                Yield x sA2 -> Yield x (T2 sA2 (ix + 1))
                                Skip sA3    -> Skip  (T2 sA3 ix)
                                Done        -> Done
            in   MkStream stepB (T2 sA0 0)


-- | Take elements from a strem while they match the given predicate.
stakeWhile (p: a -> Bool#) (ss: Stream s a): Stream s a
 = case ss of
        MkStream stepA sA0
         -> let stepB q
                 = case stepA q of
                        Yield x sA1     
                         | p x          -> Yield x sA1
                         | otherwise    -> Done [s] [a]

                        Skip sA1        -> Skip    sA1
                        Done            -> Done

            in  MkStream stepB sA0


-- | Take elements from a stream that match the given predicate.
sfilter (p: a -> Bool#) (ss: Stream s a): Stream s a
 = case ss of
        MkStream stepA sA0
         -> let stepB q
                 = case stepA q of
                        Yield x sA1     
                         | p x          -> Yield x sA1
                         | otherwise    -> Skip    sA1

                        Skip sA1        -> Skip    sA1
                        Done            -> Done

            in  MkStream stepB sA0