#include "fusion-phases.h"
module Data.Array.Vector.Stream where
import Debug.Trace
import Data.Array.Vector.Prim.Hyperstrict
data Step s a = Done
| Skip !s
| Yield !a !s
instance Functor (Step s) where
fmap f Done = Done
fmap f (Skip s) = Skip s
fmap f (Yield x s) = Yield (f x) s
data Stream a = forall s. Stream (s -> Step s a) !s Int
newtype Box a = Box a
emptyS :: Stream a
emptyS = Stream (const Done) () 0
nullS :: Stream a -> Bool
nullS (Stream next s0 _) = loop_null s0
where
loop_null s = case next s of
Done -> True
Yield _ _ -> False
Skip s' -> s' `seq` loop_null s'
singletonS :: a -> Stream a
singletonS x = Stream next True 1
where
next True = Yield x False
next False = Done
consS :: a -> Stream a -> Stream a
consS x (Stream next s n) = Stream next' (JustS (Box x) :*: s) (n+1)
where
next' (JustS (Box x) :*: s) = Yield x (NothingS :*: s)
next' (NothingS :*: s) = case next s of
Yield y s' -> Yield y (NothingS :*: s')
Skip s' -> Skip (NothingS :*: s')
Done -> Done
snocS :: Stream a -> a -> Stream a
snocS (Stream next s n) x = Stream next' (JustS s) (n+1)
where
next' (JustS s) = case next s of
Yield y s' -> Yield y (JustS s')
Skip s' -> Skip (JustS s')
Done -> Yield x NothingS
next' NothingS = Done
replicateS :: Int -> a -> Stream a
replicateS n x = Stream next 0 n
where
next i | i == n = Done
| otherwise = Yield x (i+1)
replicateEachS :: Int -> Stream (Int :*: a) -> Stream a
replicateEachS n (Stream next s _) =
Stream next' (0 :*: NothingS :*: s) n
where
next' (0 :*: _ :*: s) =
case next s of
Done -> Done
Skip s' -> Skip (0 :*: NothingS :*: s')
Yield (k :*: x) s' -> Skip (k :*: JustS (Box x) :*: s')
next' (k :*: NothingS :*: s) = Done
next' (k :*: JustS (Box x) :*: s) =
Yield x (k1 :*: JustS (Box x) :*: s)
(+++) :: Stream a -> Stream a -> Stream a
Stream next1 s1 n1 +++ Stream next2 s2 n2 = Stream next (LeftS s1) (n1 + n2)
where
next (LeftS s1) =
case next1 s1 of
Done -> Skip (RightS s2)
Skip s1' -> Skip (LeftS s1')
Yield x s1' -> Yield x (LeftS s1')
next (RightS s2) =
case next2 s2 of
Done -> Done
Skip s2' -> Skip (RightS s2')
Yield x s2' -> Yield x (RightS s2')
indexS :: Stream a -> Int -> a
indexS (Stream next s0 _) n0
| n0 < 0 = error "Data.Array.Vector.Stream.indexS: negative index"
| otherwise = loop_index n0 s0
where
loop_index n s = case next s of
Yield x s' | n == 0 -> x
| otherwise -> s' `seq` loop_index (n1) s'
Skip s' -> s' `seq` loop_index n s'
Done -> error "Data.Array.Vector.Stream.indexS: index too large"
indexedS :: Stream a -> Stream (Int :*: a)
indexedS (Stream next s n) = Stream next' (0 :*: s) n
where
next' (i :*: s) = case next s of
Yield x s' -> Yield (i :*: x) ((i+1) :*: s')
Skip s' -> Skip (i :*: s')
Done -> Done
headS :: Stream a -> a
headS (Stream next s0 _) = loop_head s0
where
loop_head s = case next s of
Yield x _ -> x
Skip s' -> s' `seq` loop_head s'
Done -> errorEmptyStream "head"
tailS :: Stream a -> Stream a
tailS (Stream next s n) = Stream next' (False :*: s) (n1)
where
next' (False :*: s) = case next s of
Yield x s' -> Skip (True :*: s')
Skip s' -> Skip (False :*: s')
Done -> error "Stream.tailS: empty stream"
next' (True :*: s) = case next s of
Yield x s' -> Yield x (True :*: s')
Skip s' -> Skip (True :*: s')
Done -> Done
toStream :: [a] -> Stream a
toStream xs = Stream gen (Box xs) (length xs)
where
gen (Box []) = Done
gen (Box (x:xs)) = Yield x (Box xs)
fromStream :: Stream a -> [a]
fromStream (Stream next s _) = gen s
where
gen s = case next s of
Done -> []
Skip s' -> gen s'
Yield x s' -> x : gen s'
initS :: Stream a -> Stream a
initS (Stream next0 s0 n) = Stream next' (NothingS :*: s0) (n1)
where
next' (NothingS :*: s) = case next0 s of
Yield x s' -> Skip (JustS (Box x) :*: s')
Skip s' -> Skip (NothingS :*: s')
Done -> errorEmptyStream "init"
next' (JustS (Box x) :*: s) = case next0 s of
Yield x' s' -> Yield x (JustS (Box x') :*: s')
Skip s' -> Skip (JustS (Box x) :*: s')
Done -> Done
takeS :: Int -> Stream a -> Stream a
takeS n0 (Stream next0 s0 _) = Stream next' (n0 :*: s0) (max 0 n0)
where
next' (n :*: s)
| n <= 0 = Done
| otherwise = case next0 s of
Yield x s' -> Yield x ((n1) :*: s')
Skip s' -> Skip ( n :*: s')
Done -> Done
dropS :: Int -> Stream a -> Stream a
dropS n0 (Stream next0 s0 n) = Stream next' (JustS (max 0 n0) :*: s0) (max 0 (n n0))
where
next' (JustS n :*: s)
| n == 0 = Skip (NothingS :*: s)
| otherwise = case next0 s of
Yield _ s' -> Skip (JustS (n1) :*: s')
Skip s' -> Skip (JustS n :*: s')
Done -> Done
next' (NothingS :*: s) = case next0 s of
Yield x s' -> Yield x (NothingS :*: s')
Skip s' -> Skip (NothingS :*: s')
Done -> Done
elemS :: Eq a => a -> Stream a -> Bool
elemS x (Stream next s0 _) = loop_elem s0
where
loop_elem s = case next s of
Yield y s'
| x == y -> True
| otherwise -> s' `seq` loop_elem s'
Skip s' -> s' `seq` loop_elem s'
Done -> False
lookupS :: Eq a => a -> Stream (a :*: b) -> Maybe b
lookupS key (Stream next s0 _) = loop_lookup s0
where
loop_lookup s = case next s of
Yield (x :*: y) s'
| key == x -> Just y
| otherwise -> s' `seq` loop_lookup s'
Skip s' -> s' `seq` loop_lookup s'
Done -> Nothing
mapS :: (a -> b) -> Stream a -> Stream b
mapS f (Stream next s n) = Stream next' s n
where
next' 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
filterS f (Stream next s n) = Stream next' s n
where
next' s = case next s of
Done -> Done
Skip s' -> Skip s'
Yield x s' | f x -> Yield x s'
| otherwise -> Skip s'
foldS :: (b -> a -> b) -> b -> Stream a -> b
foldS f z (Stream next s _) = fold z s
where
fold z s = z `seq` case next s of
Yield x s' -> s' `seq` fold (f z x) s'
Skip s' -> s' `seq` fold z s'
Done -> z
foldl1S :: (a -> a -> a) -> Stream a -> a
foldl1S f (Stream next s0 _) = loop0_foldl1' s0
where
loop0_foldl1' s = case next s of
Yield x s' -> s' `seq` loop_foldl1' x s'
Skip s' -> s' `seq` loop0_foldl1' s'
Done -> errorEmptyStream "foldl1"
loop_foldl1' z s = z `seq` case next s of
Yield x s' -> s' `seq` loop_foldl1' (f z x) s'
Skip s' -> s' `seq` loop_foldl1' z s'
Done -> z
fold1MaybeS :: (a -> a -> a) -> Stream a -> MaybeS a
fold1MaybeS f (Stream next s _) = fold0 s
where
fold0 s = case next s of
Done -> NothingS
Skip s' -> s' `seq` fold0 s'
Yield x s' -> s' `seq` fold1 x s'
fold1 z s = z `seq` case next s of
Done -> JustS z
Skip s' -> s' `seq` fold1 z s'
Yield x s' -> s' `seq` fold1 (f z x) s'
foldrS :: (a -> b -> b) -> b -> Stream a -> b
foldrS f z (Stream next s0 _len) = loop_foldr s0
where
loop_foldr !s = case next s of
Done -> z
Skip s' -> loop_foldr s'
Yield x s' -> f x (loop_foldr s')
foldr1S :: (a -> a -> a) -> Stream a -> a
foldr1S f (Stream next s0 _len) = loop0_foldr1 s0
where
loop0_foldr1 !s = case next s of
Done -> error "foldr1S: empty stream"
Skip s' -> loop0_foldr1 s'
Yield x s' -> loop_foldr1 x s'
loop_foldr1 x !s = case next s of
Done -> x
Skip s' -> loop_foldr1 x s'
Yield x' s' -> f x (loop_foldr1 x' s')
scanS :: (b -> a -> b) -> b -> Stream a -> Stream b
scanS f z (Stream next s n) = Stream next' (Box z :*: s) n
where
next' (Box z :*: s) = case next s of
Done -> Done
Skip s' -> Skip (Box z :*: s')
Yield x s' -> Yield z (Box (f z x) :*: s')
scan1S :: (a -> a -> a) -> Stream a -> Stream a
scan1S f (Stream next s n) = Stream next' (NothingS :*: s) n
where
next' (NothingS :*: s) =
case next s of
Yield x s' -> Yield x (JustS (Box x) :*: s')
Skip s' -> Skip (NothingS :*: s')
Done -> Done
next' (JustS (Box z) :*: s) =
case next s of
Yield x s' -> let y = f z x
in
Yield y (JustS (Box y) :*: s')
Skip s' -> Skip (JustS (Box z) :*: s)
Done -> Done
mapAccumS :: (acc -> a -> acc :*: b) -> acc -> Stream a -> Stream b
mapAccumS f acc (Stream step s n) = Stream step' (s :*: Box acc) n
where
step' (s :*: Box acc) = case step s of
Done -> Done
Skip s' -> Skip (s' :*: Box acc)
Yield x s' -> let acc' :*: y = f acc x
in
Yield y (s' :*: Box acc')
combineS:: Stream Bool -> Stream a -> Stream a -> Stream a
combineS (Stream next1 s m) (Stream nextS1 t1 n1) (Stream nextS2 t2 n2) =
Stream next (s :*: t1 :*: t2) m
where
next (s :*: t1 :*: t2) =
case next1 s of
Done -> Done
Skip s' -> Skip (s' :*: t1 :*: t2 )
Yield c s' -> if c
then case nextS1 t1 of
Done -> error "combineS: stream 1 terminated unexpectedly"
Skip t1' -> Skip (s :*: t1' :*: t2)
Yield x t1' -> Yield x (s' :*: t1' :*: t2)
else case nextS2 t2 of
Done -> error "combineS: stream 2 terminated unexpectedly"
Skip t2' -> Skip (s :*: t1 :*: t2')
Yield x t2' -> Yield x (s' :*: t1 :*: t2')
zipWithS :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
zipWithS f (Stream next1 s m) (Stream next2 t n) =
Stream next (s :*: t) m
where
next (s :*: t) =
case next1 s of
Done -> Done
Skip s' -> Skip (s' :*: t)
Yield x s' -> case next2 t of
Done -> Done
Skip t' -> Skip (s :*: t')
Yield y t' -> Yield (f x y) (s' :*: t')
zipWith3S :: (a -> b -> c -> d) -> Stream a -> Stream b -> Stream c -> Stream d
zipWith3S f (Stream next1 s1 n) (Stream next2 s2 _) (Stream next3 s3 _) =
Stream next (s1 :*: s2 :*: s3) n
where
next (s1 :*: s2 :*: s3) =
case next1 s1 of
Done -> Done
Skip s1' -> Skip (s1' :*: s2 :*: s3)
Yield x s1' ->
case next2 s2 of
Done -> Done
Skip s2' -> Skip (s1 :*: s2' :*: s3)
Yield y s2' ->
case next3 s3 of
Done -> Done
Skip s3' -> Skip (s1 :*: s2 :*: s3')
Yield z s3' -> Yield (f x y z) (s1' :*: s2' :*: s3')
zipS :: Stream a -> Stream b -> Stream (a :*: b)
zipS = zipWithS (:*:)
enumFromToFracS :: (Ord a, RealFrac a) => a -> a -> Stream a
enumFromToFracS n m = Stream next n (truncate (m n) + 2)
where
lim = m + 1/2
next s | s > lim = Done
| otherwise = Yield s (s+1)
enumFromToS :: (Integral a, Ord a) => a -> a -> Stream a
enumFromToS start end
= Stream step start (max 0 (fromIntegral (end start + 1)))
where
step s | s > end = Done
| otherwise = Yield s (s+1)
enumFromThenToS :: Int -> Int -> Int -> Stream Int
enumFromThenToS start next end
= enumFromStepLenS start delta len
where
delta = next start
diff = end start
len | start < next && start <= end = ((endstart) `div` delta) + 1
| start > next && start >= end = ((startend) `div` (startnext)) + 1
| otherwise = 0
enumFromStepLenS :: Int -> Int -> Int -> Stream Int
enumFromStepLenS s d n = Stream step (s :*: n) n
where
step (s :*: 0) = Done
step (s :*: n) = Yield s ((s+d) :*: (n1))
enumFromToEachS :: Int -> Stream (Int :*: Int) -> Stream Int
enumFromToEachS n (Stream next s _) = Stream next' (NothingS :*: s) n
where
next' (NothingS :*: s)
= case next s of
Yield (k :*: m) s' -> Skip (JustS (k :*: m) :*: s')
Skip s' -> Skip (NothingS :*: s')
Done -> Done
next' (JustS (k :*: m) :*: s)
| k > m = Skip (NothingS :*: s)
| otherwise = Yield k (JustS (k+1 :*: m) :*: s)
findS :: (a -> Bool) -> Stream a -> Maybe a
findS p (Stream next s _) = go s
where
go s = case next s of
Yield x s' | p x -> Just x
| otherwise -> go s'
Skip s' -> go s'
Done -> Nothing
findIndexS :: (a -> Bool) -> Stream a -> Maybe Int
findIndexS p (Stream next s _) = go 0 s
where
go i s = case next s of
Yield x s' | p x -> Just i
| otherwise -> go (i+1) s'
Skip s' -> go i s'
Done -> Nothing
takeWhileS :: (a -> Bool) -> Stream a -> Stream a
takeWhileS p (Stream next0 s0 n) = Stream next s0 n
where
next s = case next0 s of
Done -> Done
Skip s' -> Skip s'
Yield x s' | p x -> Yield x s'
| otherwise -> Done
dropWhileS :: (a -> Bool) -> Stream a -> Stream a
dropWhileS p (Stream next0 s0 n) = Stream next (True :*: s0) n
where
next (True :*: s) = case next0 s of
Done -> Done
Skip s' -> Skip (True :*: s')
Yield x s' | p x -> Skip (True :*: s')
| otherwise -> Yield x (False :*: s')
next (False :*: s) = case next0 s of
Done -> Done
Skip s' -> Skip (False :*: s')
Yield x s' -> Yield x (False :*: s')
unfoldS :: Int -> (b -> MaybeS (a :*: b)) -> b -> Stream a
unfoldS 0 _ _ = emptyS
unfoldS n f s0 = Stream next (JustS (0 :*: s0)) n
where
next (JustS (i :*: s)) = case f s of
NothingS -> Done
JustS (w :*: s')
| (n1) == i-> Yield w NothingS
| otherwise -> Yield w (JustS (i+1 :*: s'))
next _ = Done
errorEmptyStream :: String -> a
errorEmptyStream fun = moduleError fun "empty vector"
moduleError :: String -> String -> a
moduleError fun msg = error ("Data.Array.Vector.Stream." ++ fun ++ ':':' ':msg)