Skip to content

Commit

Permalink
Add Maybe (v a) to Stream representations
Browse files Browse the repository at this point in the history
  • Loading branch information
Roman Leshchinskiy committed Jan 29, 2012
1 parent 8a0a0c0 commit f375eff
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 9 deletions.
4 changes: 2 additions & 2 deletions Data/Vector/Fusion/Stream.hs
Expand Up @@ -124,9 +124,9 @@ inplace f s = s `seq` f s
-- | Convert a pure stream to a monadic stream
liftStream :: Monad m => Stream v a -> M.Stream m v a
{-# INLINE_STREAM liftStream #-}
liftStream (M.Stream (M.Unf step s) (M.Unf vstep t) sz)
liftStream (M.Stream (M.Unf step s) (M.Unf vstep t) v sz)
= M.Stream (M.Unf (return . unId . step) s)
(M.Unf (return . unId . vstep) t) sz
(M.Unf (return . unId . vstep) t) v sz

-- | 'Size' hint of a 'Stream'
size :: Stream v a -> Size
Expand Down
20 changes: 13 additions & 7 deletions Data/Vector/Fusion/Stream/Monadic.hs
Expand Up @@ -144,12 +144,13 @@ instance Monad m => Functor (Unf m) where
-- | Monadic streams
data Stream m v a = Stream { sElems :: Unf m a
, sChunks :: Unf m (Chunk v a)
, sVector :: Maybe (v a)
, sSize :: Size
}

simple :: Monad m => (s -> m (Step s a)) -> s -> Size -> Stream m v a
{-# INLINE simple #-}
simple step s sz = Stream (Unf step s) (Unf step' s) sz
simple step s sz = Stream (Unf step s) (Unf step' s) Nothing sz
where
step' s = do r <- step s
return $ fmap (\x -> Chunk 1 (\v -> M.basicUnsafeWrite v 0 x)) r
Expand Down Expand Up @@ -198,7 +199,7 @@ singleton x = simple (return . step) True (Exact 1)
-- | Replicate a value to a given length
replicate :: Monad m => Int -> a -> Stream m v a
{-# INLINE_STREAM replicate #-}
replicate n x = Stream (Unf pstep n) (Unf vstep True) (Exact len)
replicate n x = Stream (Unf pstep n) (Unf vstep True) Nothing (Exact len)
where
len = delay_inline max n 0

Expand Down Expand Up @@ -253,9 +254,9 @@ infixr 5 ++
-- | Concatenate two 'Stream's
(++) :: Monad m => Stream m v a -> Stream m v a -> Stream m v a
{-# INLINE_STREAM (++) #-}
Stream (Unf stepa sa) (Unf vstepa vsa) na
++ Stream (Unf stepb sb) (Unf vstepb vsb) nb
= Stream (Unf step (Left sa)) (Unf vstep (Left vsa)) (na + nb)
Stream{sElems = Unf stepa sa, sChunks = Unf vstepa vsa, sSize = na}
++ Stream{sElems = Unf stepb sb, sChunks = Unf vstepb vsb, sSize = nb}
= Stream (Unf step (Left sa)) (Unf vstep (Left vsa)) Nothing (na + nb)
where
{-# INLINE_INNER step #-}
step (Left sa) = do
Expand Down Expand Up @@ -1040,7 +1041,7 @@ concatMapM f Stream{sElems = Unf step s} = simple concatMap_go (Left s) Unknown
return $ Skip (Right (b_stream, s'))
Skip s' -> return $ Skip (Left s')
Done -> return Done
concatMap_go (Right (Stream (Unf inner_step inner_s) _ sz, s)) = do
concatMap_go (Right (Stream{sElems = Unf inner_step inner_s, sSize = sz}, s)) = do
r <- inner_step inner_s
case r of
Yield b inner_s' -> return $ Yield b (Right (simple inner_step inner_s' sz, s))
Expand Down Expand Up @@ -1587,7 +1588,10 @@ unsafeFromList sz xs = simple step xs sz

fromVector :: (Monad m, Vector v a) => v a -> Stream m v a
{-# INLINE_STREAM fromVector #-}
fromVector v = v `seq` n `seq` Stream (Unf step 0) (Unf vstep True) (Exact n)
fromVector v = v `seq` n `seq` Stream (Unf step 0)
(Unf vstep True)
(Just v)
(Exact n)
where
n = basicLength v

Expand All @@ -1605,6 +1609,7 @@ fromVectors :: (Monad m, Vector v a) => [v a] -> Stream m v a
{-# INLINE_STREAM fromVectors #-}
fromVectors vs = Stream (Unf pstep (Left vs))
(Unf vstep vs)
Nothing
(Exact n)
where
n = List.foldl' (\k v -> k + basicLength v) 0 vs
Expand All @@ -1628,6 +1633,7 @@ fromVectorStream :: (Monad m, Vector v a) => Stream m u (v a) -> Stream m v a
fromVectorStream Stream{sElems = Unf step s}
= Stream (Unf pstep (Left s))
(Unf vstep s)
Nothing
Unknown
where
pstep (Left s) = do
Expand Down

0 comments on commit f375eff

Please sign in to comment.