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
 = MkStream step xx
 where  
        step (s1: List a)
         = case s1 of
                Nil       -> Done
                Cons x xs -> Yield x xs


-- | Convert a stream to a list.
listOfStream (ss: Stream s a): List a
 | MkStream f s0 <- ss
 = 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 
        (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
        (ss: Stream s a) (arr: Array r a) (ix: Nat)
        : S (Write r) Unit
 | MkStream f s0  <- ss
 = go s0 0
 where  
        go (s: s) (ix: Nat): S (Write r) Unit
         = case f s of
                Yield x s1
                 -> do  writeArray arr ix x
                        go s1 (ix + 1)

                Skip s1
                 ->     go s1 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
 = MkStream step' x
 where  
        step' sA
         = case step sA of
                T2 s' x -> Yield x s'

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
 | MkStream stepA  sA0 <- ss
 = MkStream stepA2 (T2 sA0 True)
 where  
        stepA2 (T2 sA1 True)
         = Yield x (T2 sA1 False)

        stepA2 (T2 sA1 False)
         = case stepA sA1 of
                Yield y sA2 -> Yield y (T2 sA2 False)
                Skip  sA2   -> Skip    (T2 sA2 False)
                Done        -> Done


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


-- Scans ----------------------------------------------------------------------
-- | Like `smap`, but keep a running accumulator as we walk along the stream.
smapacc :  [a b c s: Data]
        .  (a -> b -> Tup2 a c)
        -> a
        -> Stream s b
        -> Stream (Tup2 s a) c

smapacc f z (MkStream fB sB0)
 = MkStream stepC (T2 sB0 z)
 where
        stepC (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


-- Folds ----------------------------------------------------------------------
-- | Fold all the elements from a stream.
sfold   (f: a -> b -> a) 
        (acc: a) 
        ((MkStream step s0): Stream s b)
        : a
 = 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)
        ((MkStream fA sA0): Stream s a)
        : Stream (Tup2 s Nat) a

 = MkStream stepB (T2 sA0 0)
 where  
        stepB (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


-- | Take elements from a stream while they match the given predicate.
stakeWhile 
        (p: a -> Bool) 
        ((MkStream stepA sA0): Stream s a)
        : Stream s a

 = MkStream stepB sA0
 where  
        stepB (q: s)
         = case stepA q of
                Yield x sA1     
                 | p x          -> Yield x sA1
                 | otherwise    -> Done

                Skip sA1        -> Skip    sA1
                Done            -> Done


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

                Skip sA1        -> Skip sA1
                Done            -> Done