{-# LANGUAGE CPP                 #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module HaskellWorks.Data.Streams.Stream where

import Data.Bool
import HaskellWorks.Data.Streams.Size

import Prelude hiding (drop, foldl, zipWith)

data Stream a where
  Stream :: ()
    => (s -> Step s a)
    -> s
    -> Size
    -> Stream a

instance Functor Stream where
  fmap :: forall a b. (a -> b) -> Stream a -> Stream b
fmap a -> b
f (Stream s -> Step s a
step s
s Size
i) = (s -> Step s b) -> s -> Size -> Stream b
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream ((a -> b) -> Step s a -> Step s b
forall a b. (a -> b) -> Step s a -> Step s b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (Step s a -> Step s b) -> (s -> Step s a) -> s -> Step s b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> Step s a
step) s
s Size
i
  {-# INLINE fmap #-}

data Step s a
  = Yield a s
  | Skip s
  | Done

instance Functor (Step s) where
  fmap :: forall a b. (a -> b) -> Step s a -> Step s b
fmap a -> b
f (Yield a
a s
s) = b -> s -> Step s b
forall s a. a -> s -> Step s a
Yield (a -> b
f a
a) s
s
  fmap a -> b
_ (Skip s
s)    = s -> Step s b
forall s a. s -> Step s a
Skip s
s
  fmap a -> b
_ Step s a
Done        = Step s b
forall s a. Step s a
Done
  {-# INLINE fmap #-}

zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
zipWith :: forall a b c. (a -> b -> c) -> Stream a -> Stream b -> Stream c
zipWith a -> b -> c
f (Stream s -> Step s a
stepa s
sa Size
na) (Stream s -> Step s b
stepb s
sb Size
nb) = ((s, s, Maybe a) -> Step (s, s, Maybe a) c)
-> (s, s, Maybe a) -> Size -> Stream c
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream (s, s, Maybe a) -> Step (s, s, Maybe a) c
step (s
sa, s
sb, Maybe a
forall a. Maybe a
Nothing) (Size -> Size -> Size
smaller Size
na Size
nb)
  where step :: (s, s, Maybe a) -> Step (s, s, Maybe a) c
step (s
ta, s
tb, Maybe a
Nothing) = case s -> Step s a
stepa s
ta of
          Yield a
xa s
ta0 -> (s, s, Maybe a) -> Step (s, s, Maybe a) c
forall s a. s -> Step s a
Skip (s
ta0, s
tb, a -> Maybe a
forall a. a -> Maybe a
Just a
xa)
          Skip s
ta0     -> (s, s, Maybe a) -> Step (s, s, Maybe a) c
forall s a. s -> Step s a
Skip (s
ta0, s
tb, Maybe a
forall a. Maybe a
Nothing)
          Step s a
Done         -> Step (s, s, Maybe a) c
forall s a. Step s a
Done
        step (s
ta, s
tb, Just a
xa) = case s -> Step s b
stepb s
tb of
          Yield b
y s
tb0 -> c -> (s, s, Maybe a) -> Step (s, s, Maybe a) c
forall s a. a -> s -> Step s a
Yield (a -> b -> c
f a
xa b
y) (s
ta, s
tb0, Maybe a
forall a. Maybe a
Nothing)
          Skip s
tb0    -> (s, s, Maybe a) -> Step (s, s, Maybe a) c
forall s a. s -> Step s a
Skip (s
ta, s
tb0, a -> Maybe a
forall a. a -> Maybe a
Just a
xa)
          Step s b
Done        -> Step (s, s, Maybe a) c
forall s a. Step s a
Done
        {-# INLINE [0] step #-}
{-# INLINE [1] zipWith #-}

zipWithState :: (a -> b -> s -> (c, s)) -> s -> Stream a -> Stream b -> Stream c
zipWithState :: forall a b s c.
(a -> b -> s -> (c, s)) -> s -> Stream a -> Stream b -> Stream c
zipWithState a -> b -> s -> (c, s)
f s
state (Stream s -> Step s a
stepa s
sa Size
na) (Stream s -> Step s b
stepb s
sb Size
nb) = ((s, s, Maybe a, s) -> Step (s, s, Maybe a, s) c)
-> (s, s, Maybe a, s) -> Size -> Stream c
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream (s, s, Maybe a, s) -> Step (s, s, Maybe a, s) c
step (s
sa, s
sb, Maybe a
forall a. Maybe a
Nothing, s
state) (Size -> Size -> Size
smaller Size
na Size
nb)
  where step :: (s, s, Maybe a, s) -> Step (s, s, Maybe a, s) c
step (s
ta, s
tb, Maybe a
Nothing, s
oldState) = case s -> Step s a
stepa s
ta of
          Yield a
xa s
ta0 -> (s, s, Maybe a, s) -> Step (s, s, Maybe a, s) c
forall s a. s -> Step s a
Skip (s
ta0, s
tb, a -> Maybe a
forall a. a -> Maybe a
Just a
xa, s
oldState)
          Skip s
ta0     -> (s, s, Maybe a, s) -> Step (s, s, Maybe a, s) c
forall s a. s -> Step s a
Skip (s
ta0, s
tb, Maybe a
forall a. Maybe a
Nothing, s
oldState)
          Step s a
Done         -> Step (s, s, Maybe a, s) c
forall s a. Step s a
Done
        step (s
ta, s
tb, Just a
xa, s
oldState) = case s -> Step s b
stepb s
tb of
          Yield b
y s
tb0 -> let (c
newValue, s
newState) = a -> b -> s -> (c, s)
f a
xa b
y s
oldState in c -> (s, s, Maybe a, s) -> Step (s, s, Maybe a, s) c
forall s a. a -> s -> Step s a
Yield c
newValue (s
ta, s
tb0, Maybe a
forall a. Maybe a
Nothing, s
newState)
          Skip s
tb0    -> (s, s, Maybe a, s) -> Step (s, s, Maybe a, s) c
forall s a. s -> Step s a
Skip (s
ta, s
tb0, a -> Maybe a
forall a. a -> Maybe a
Just a
xa, s
oldState)
          Step s b
Done        -> Step (s, s, Maybe a, s) c
forall s a. Step s a
Done
        {-# INLINE [0] step #-}
{-# INLINE [1] zipWithState #-}

enumFromStepN :: Num a => a -> a -> Int -> Stream a
enumFromStepN :: forall a. Num a => a -> a -> Int -> Stream a
enumFromStepN a
x a
y Int
n = a
x a -> Stream a -> Stream a
forall a b. a -> b -> b
`seq` a
y a -> Stream a -> Stream a
forall a b. a -> b -> b
`seq` Int
n Int -> Stream a -> Stream a
forall a b. a -> b -> b
`seq` ((a, Int) -> Step (a, Int) a) -> (a, Int) -> Size -> Stream a
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream (a, Int) -> Step (a, Int) a
step (a
x, Int
n) (Int -> Size
Exact Int
n)
  where step :: (a, Int) -> Step (a, Int) a
step (a
w, Int
m) | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0     = a -> (a, Int) -> Step (a, Int) a
forall s a. a -> s -> Step s a
Yield a
w (a
w a -> a -> a
forall a. Num a => a -> a -> a
+ a
y, Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                    | Bool
otherwise = Step (a, Int) a
forall s a. Step s a
Done
        {-# INLINE [0] step #-}
{-# INLINE [1] enumFromStepN #-}

foldl :: (a -> b -> a) -> a -> Stream b -> a
foldl :: forall a b. (a -> b -> a) -> a -> Stream b -> a
foldl a -> b -> a
f a
z (Stream s -> Step s b
step s
s Size
_) = a -> s -> a
loop a
z s
s
  where loop :: a -> s -> a
loop a
za s
sa = a
za a -> a -> a
forall a b. a -> b -> b
`seq` case s -> Step s b
step s
sa of
          Yield b
x s
sb -> a -> s -> a
loop (a -> b -> a
f a
za b
x) s
sb
          Skip s
sb    -> a -> s -> a
loop a
za s
sb
          Step s b
Done       -> a
za
{-# INLINE [1] foldl #-}

drop :: Int -> Stream a -> Stream a
drop :: forall a. Int -> Stream a -> Stream a
drop Int
n s :: Stream a
s@(Stream s -> Step s a
step s
state Size
size) = if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
  then case s -> Step s a
step s
state of
    Yield a
_ s
newState -> Int -> Stream a -> Stream a
forall a. Int -> Stream a -> Stream a
drop (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)  (Stream a -> Stream a) -> Stream a -> Stream a
forall a b. (a -> b) -> a -> b
$ (s -> Step s a) -> s -> Size -> Stream a
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream s -> Step s a
step s
newState (Size
size Size -> Size -> Size
forall a. Num a => a -> a -> a
- Size
1)
    Skip s
newState    -> Int -> Stream a -> Stream a
forall a. Int -> Stream a -> Stream a
drop  Int
n       (Stream a -> Stream a) -> Stream a -> Stream a
forall a b. (a -> b) -> a -> b
$ (s -> Step s a) -> s -> Size -> Stream a
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream s -> Step s a
step s
newState Size
size
    Step s a
Done             -> Stream a
s
  else Stream a
s

{-# RULES
  "zipWith xs xs [Vector.Stream]" forall f xs. zipWith f xs xs = fmap (\x -> f x x) xs   #-}

append :: Stream a -> Stream a -> Stream a
append :: forall a. Stream a -> Stream a -> Stream a
append (Stream s -> Step s a
stepa s
ta Size
na) (Stream s -> Step s a
stepb s
tb Size
nb) = (Either s s -> Step (Either s s) a)
-> Either s s -> Size -> Stream a
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream Either s s -> Step (Either s s) a
step (s -> Either s s
forall a b. a -> Either a b
Left s
ta) (Size
na Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Size
nb)
  where step :: Either s s -> Step (Either s s) a
step (Left s
sa) = case s -> Step s a
stepa s
sa of
          Yield a
x s
sa' -> a -> Either s s -> Step (Either s s) a
forall s a. a -> s -> Step s a
Yield a
x (s -> Either s s
forall a b. a -> Either a b
Left  s
sa')
          Skip    s
sa' -> Either s s -> Step (Either s s) a
forall s a. s -> Step s a
Skip    (s -> Either s s
forall a b. a -> Either a b
Left  s
sa')
          Step s a
Done        -> Either s s -> Step (Either s s) a
forall s a. s -> Step s a
Skip    (s -> Either s s
forall a b. b -> Either a b
Right s
tb)
        step (Right s
sb) = case s -> Step s a
stepb s
sb of
          Yield a
x s
sb' -> a -> Either s s -> Step (Either s s) a
forall s a. a -> s -> Step s a
Yield a
x (s -> Either s s
forall a b. b -> Either a b
Right s
sb')
          Skip    s
sb' -> Either s s -> Step (Either s s) a
forall s a. s -> Step s a
Skip    (s -> Either s s
forall a b. b -> Either a b
Right s
sb')
          Step s a
Done        -> Step (Either s s) a
forall s a. Step s a
Done
        {-# INLINE step #-}
{-# INLINE append #-}

singleton :: a -> Stream a
singleton :: forall a. a -> Stream a
singleton a
a = (Bool -> Step Bool a) -> Bool -> Size -> Stream a
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream (Step Bool a -> Step Bool a -> Bool -> Step Bool a
forall a. a -> a -> Bool -> a
bool Step Bool a
forall s a. Step s a
Done (a -> Bool -> Step Bool a
forall s a. a -> s -> Step s a
Yield a
a Bool
False)) Bool
True Size
1

repeat :: Int -> a -> Stream a
repeat :: forall a. Int -> a -> Stream a
repeat Int
n a
a = (Int -> Step Int a) -> Int -> Size -> Stream a
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream Int -> Step Int a
step Int
n Size
1
  where step :: Int -> Step Int a
step Int
i = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then a -> Int -> Step Int a
forall s a. a -> s -> Step s a
Yield a
a (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) else Step Int a
forall s a. Step s a
Done

transcribe :: (s -> a -> (b, s)) -> s -> Stream a -> Stream b
transcribe :: forall s a b. (s -> a -> (b, s)) -> s -> Stream a -> Stream b
transcribe s -> a -> (b, s)
f s
w (Stream s -> Step s a
step s
state Size
size) = ((s, s) -> Step (s, s) b) -> (s, s) -> Size -> Stream b
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream (s, s) -> Step (s, s) b
step' (s
state, s
w) Size
size
  where step' :: (s, s) -> Step (s, s) b
step' (s
t, s
x) = case s -> Step s a
step s
t of
          Yield a
a s
s' -> let (b
z, s
y) = s -> a -> (b, s)
f s
x a
a in b -> (s, s) -> Step (s, s) b
forall s a. a -> s -> Step s a
Yield b
z (s
s', s
y)
          Skip    s
s' -> (s, s) -> Step (s, s) b
forall s a. s -> Step s a
Skip (s
s', s
x)
          Step s a
Done       -> Step (s, s) b
forall s a. Step s a
Done
        {-# INLINE step' #-}
{-# INLINE transcribe #-}

concatMap :: (a -> Stream b) -> Stream a -> Stream b
concatMap :: forall a b. (a -> Stream b) -> Stream a -> Stream b
concatMap a -> Stream b
f (Stream s -> Step s a
stepA s
stateA Size
_) = ((s, Maybe (Stream b)) -> Step (s, Maybe (Stream b)) b)
-> (s, Maybe (Stream b)) -> Size -> Stream b
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream (s, Maybe (Stream b)) -> Step (s, Maybe (Stream b)) b
stepC (s
stateA, Maybe (Stream b)
forall a. Maybe a
Nothing) Size
Unknown
  where stepC :: (s, Maybe (Stream b)) -> Step (s, Maybe (Stream b)) b
stepC (s
stateA0, Maybe (Stream b)
Nothing) = case s -> Step s a
stepA s
stateA0 of
          Yield a
a s
stateA1 -> (s, Maybe (Stream b)) -> Step (s, Maybe (Stream b)) b
forall s a. s -> Step s a
Skip (s
stateA1, Stream b -> Maybe (Stream b)
forall a. a -> Maybe a
Just (a -> Stream b
f a
a))
          Skip    s
stateA1 -> (s, Maybe (Stream b)) -> Step (s, Maybe (Stream b)) b
forall s a. s -> Step s a
Skip (s
stateA1, Maybe (Stream b)
forall a. Maybe a
Nothing)
          Step s a
Done            -> Step (s, Maybe (Stream b)) b
forall s a. Step s a
Done
        stepC (s
stateA0, Just (Stream s -> Step s b
stepB s
stateB Size
_)) = case s -> Step s b
stepB s
stateB of
          Yield b
b s
stateB1 -> b -> (s, Maybe (Stream b)) -> Step (s, Maybe (Stream b)) b
forall s a. a -> s -> Step s a
Yield b
b  (s
stateA0, Stream b -> Maybe (Stream b)
forall a. a -> Maybe a
Just ((s -> Step s b) -> s -> Size -> Stream b
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream s -> Step s b
stepB s
stateB1 Size
Unknown))
          Skip    s
stateB1 -> (s, Maybe (Stream b)) -> Step (s, Maybe (Stream b)) b
forall s a. s -> Step s a
Skip     (s
stateA0, Stream b -> Maybe (Stream b)
forall a. a -> Maybe a
Just ((s -> Step s b) -> s -> Size -> Stream b
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream s -> Step s b
stepB s
stateB1 Size
Unknown))
          Step s b
Done            -> (s, Maybe (Stream b)) -> Step (s, Maybe (Stream b)) b
forall s a. s -> Step s a
Skip     (s
stateA0, Maybe (Stream b)
forall a. Maybe a
Nothing)
{-# INLINE concatMap #-}

dupMap :: forall a b. (a -> b) -> (a -> b) -> Stream a -> Stream b
dupMap :: forall a b. (a -> b) -> (a -> b) -> Stream a -> Stream b
dupMap a -> b
f a -> b
g (Stream s -> Step s a
stepA s
stateA Size
_) = ((s, Maybe b) -> Step (s, Maybe b) b)
-> (s, Maybe b) -> Size -> Stream b
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream ((s -> Step s a) -> (s, Maybe b) -> Step (s, Maybe b) b
forall s. (s -> Step s a) -> (s, Maybe b) -> Step (s, Maybe b) b
mkStepB s -> Step s a
stepA) (s
stateA, Maybe b
forall a. Maybe a
Nothing) Size
Unknown
  where mkStepB :: (s -> Step s a) -> (s, Maybe b) -> Step (s, Maybe b) b
        mkStepB :: forall s. (s -> Step s a) -> (s, Maybe b) -> Step (s, Maybe b) b
mkStepB s -> Step s a
step (s
s, Maybe b
mw) = case Maybe b
mw of
          Maybe b
Nothing -> case s -> Step s a
step s
s of
            Yield a
a s
t -> b -> (s, Maybe b) -> Step (s, Maybe b) b
forall s a. a -> s -> Step s a
Yield (a -> b
f a
a) (s
t, b -> Maybe b
forall a. a -> Maybe a
Just (a -> b
g a
a))
            Skip s
t    -> (s, Maybe b) -> Step (s, Maybe b) b
forall s a. s -> Step s a
Skip (s
t, Maybe b
forall a. Maybe a
Nothing)
            Step s a
Done      -> Step (s, Maybe b) b
forall s a. Step s a
Done
          Just b
w -> b -> (s, Maybe b) -> Step (s, Maybe b) b
forall s a. a -> s -> Step s a
Yield b
w (s
s, Maybe b
forall a. Maybe a
Nothing)
{-# INLINE dupMap #-}

fromList :: [a] -> Stream a
fromList :: forall a. [a] -> Stream a
fromList [a]
as = ([a] -> Step [a] a) -> [a] -> Size -> Stream a
forall s a. (s -> Step s a) -> s -> Size -> Stream a
Stream [a] -> Step [a] a
forall {a}. [a] -> Step [a] a
step [a]
as Size
Unknown
  where step :: [a] -> Step [a] a
step []     = Step [a] a
forall s a. Step s a
Done
        step (a
b:[a]
bs) = a -> [a] -> Step [a] a
forall s a. a -> s -> Step s a
Yield a
b [a]
bs
{-# INLINE [1] fromList #-}

foldMap :: Monoid a => (b -> a) -> Stream b -> a
foldMap :: forall a b. Monoid a => (b -> a) -> Stream b -> a
foldMap b -> a
f = (a -> b -> a) -> a -> Stream b -> a
forall a b. (a -> b -> a) -> a -> Stream b -> a
foldl (\a
a b
b -> a
a a -> a -> a
forall a. Semigroup a => a -> a -> a
<> b -> a
f b
b) a
forall a. Monoid a => a
mempty
{-# INLINE [1] foldMap #-}