module StreamA
where

data Step s a = Done | Yield a s | Skip s
data Stream a = forall s. Stream (s -> Step s a) s

mkStream :: (s -> Step s a) -> s -> Stream a
mkStream = Stream

next :: Stream a -> Step (Stream a) a
{-# INLINE next #-}
next (Stream step s) = case step s of
                         Done       -> Done
                         Skip s'    -> Skip (Stream step s')
                         Yield x s' -> Yield x (Stream step s')

unstream :: Stream a -> [a]
{-# INLINE [1] unstream #-}
unstream s = go s
  where
    go s = case next s of
             Done       -> []
             Yield x s' -> x : go s'
             Skip s'    -> go s'

stream :: [a] -> Stream a
{-# INLINE [1] stream #-}
stream xs = mkStream step xs
  where
    step []     = Done
    step (x:xs) = Yield x xs

{-# RULES

"stream/unstream" forall s.
  stream (unstream s) = s

 #-}

emptyS :: Stream a
emptyS = mkStream (const Done) ()

mapS :: (a -> b) -> Stream a -> Stream b
{-# INLINE [1] mapS #-}
mapS f s = mkStream step s
  where
    step s = case next s of
               Done       -> Done
               Skip s'    -> Skip s'
               Yield x s' -> Yield (f x) s'

filterS :: (a -> Bool) -> Stream a -> Stream a
{-# INLINE [1] filterS #-}
filterS p s = mkStream step s
  where
    step s = case next s of
               Done                   -> Done
               Skip s'                -> Skip s'
               Yield x s' | p x       -> Yield x s'
                          | otherwise -> Skip s'

appendS :: Stream a -> Stream a -> Stream a
{-# INLINE [1] appendS #-}
appendS s1 s2 = mkStream step (Left s1)
  where
    step (Left s1)  = case next s1 of
                        Done        -> Skip (Right s2)
                        Skip s1'    -> Skip (Left s1')
                        Yield x s1' -> Yield x (Left s1')
    step (Right s2) = case next s2 of
                        Done        -> Done
                        Skip s2'    -> Skip (Right s2')
                        Yield x s2' -> Yield x (Right s2')

concatMapS :: (a -> Stream b) -> Stream a -> Stream b
{-# INLINE [1] concatMapS #-}
concatMapS f s = Stream step (s, emptyS)
  where
    step (s, t) =
      case next t of
        Done       -> case next s of
                        Done       -> Done
                        Skip s'    -> Skip (s', emptyS)
                        Yield x s' -> Skip (s', f x)
        Skip t'    -> Skip (s, t')
        Yield x t' -> Yield x (s, t')

foldlS :: (a -> b -> a) -> a -> Stream b -> a
{-# INLINE [1] foldlS #-}
foldlS f z s = go z s
  where
    go z s = case next s of
               Done       -> z
               Skip s'    -> go z s'
               Yield x s' -> go (f z x) s'

foldrS :: (a -> b -> b) -> b -> Stream a -> b
{-# INLINE [1] foldrS #-}
foldrS f z s = go s
  where
    go s = case next s of
             Done       -> z
             Skip s'    -> go s'
             Yield x s' -> f x (go s')

enumS :: Int -> Int -> Stream Int
{-# INLINE [1] enumS #-}
enumS m n = mkStream step m
  where
    step m | m > n     = Done
           | otherwise = Yield m (m+1)
               
mapL :: (a -> b) -> [a] -> [b]
{-# INLINE mapL #-}
mapL f = unstream . mapS f . stream

filterL :: (a -> Bool) -> [a] -> [a]
{-# INLINE filterL #-}
filterL f = unstream . filterS f . stream

appendL :: [a] -> [a] -> [a]
{-# INLINE appendL #-}
appendL xs ys = unstream (appendS (stream xs) (stream ys))

foldlL :: (a -> b -> a) -> a -> [b] -> a
{-# INLINE foldlL #-}
foldlL f z = foldlS f z . stream

foldrL :: (a -> b -> b) -> b -> [a] -> b
{-# INLINE foldrL #-}
foldrL f z = foldrS f z . stream

enumL :: Int -> Int -> [Int]
{-# INLINE enumL #-}
enumL m = unstream . enumS m

concatMapL :: (a -> [b]) -> [a] -> [b]
{-# INLINE concatMapL #-}
concatMapL f = unstream . concatMapS (stream . f) . stream