Add Stream.flatten and use it to implement concat
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Tue, 18 May 2010 05:48:42 +0000 (05:48 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Tue, 18 May 2010 05:48:42 +0000 (05:48 +0000)
Data/Vector/Fusion/Stream.hs
Data/Vector/Fusion/Stream/Monadic.hs
Data/Vector/Generic.hs

index bba2dcb..077a0f3 100644 (file)
@@ -35,7 +35,7 @@ module Data.Vector.Fusion.Stream (
   slice, init, tail, take, drop,
 
   -- * Mapping
-  map, concatMap, unbox,
+  map, concatMap, flatten, unbox,
   
   -- * Zipping
   indexed, indexedR,
@@ -613,4 +613,8 @@ unsafeFromList :: Size -> [a] -> Stream a
 {-# INLINE unsafeFromList #-}
 unsafeFromList = M.unsafeFromList
 
+-- | Create a 'Stream' of values from a 'Stream' of streamable things
+flatten :: (a -> s) -> (s -> Step s b) -> Size -> Stream a -> Stream b
+{-# INLINE_STREAM flatten #-}
+flatten mk istep sz = M.flatten (return . mk) (return . istep) sz . liftStream
 
index 08cd814..d6c0d8e 100644 (file)
@@ -31,7 +31,7 @@ module Data.Vector.Fusion.Stream.Monadic (
   slice, init, tail, take, drop,
 
   -- * Mapping
-  map, mapM, mapM_, trans, unbox, concatMap,
+  map, mapM, mapM_, trans, unbox, concatMap, flatten,
   
   -- * Zipping
   indexed, indexedR, zipWithM_,
@@ -931,6 +931,30 @@ concatMapM f (Stream step s _) = Stream concatMap_go (Left s) Unknown
             Skip    inner_s' -> return $ Skip (Right (Stream inner_step inner_s' sz, s))
             Done             -> return $ Skip (Left s)
 
+-- | Create a 'Stream' of values from a 'Stream' of streamable things
+flatten :: Monad m => (a -> m s) -> (s -> m (Step s b)) -> Size
+                   -> Stream m a -> Stream m b
+{-# INLINE_STREAM flatten #-}
+flatten mk istep sz (Stream ostep t _) = Stream step (Left t) sz
+  where
+    {-# INLINE_INNER step #-}
+    step (Left t) = do
+                      r <- ostep t
+                      case r of
+                        Yield a t' -> do
+                                        s <- mk a
+                                        return $ Skip (Right (s,t'))
+                        Skip    t' -> return $ Skip (Left t')
+                        Done       -> return $ Done
+
+    
+    step (Right (s,t)) = do
+                           r <- istep s
+                           case r of
+                             Yield x s' -> return $ Yield x (Right (s',t))
+                             Skip    s' -> return $ Skip    (Right (s',t))
+                             Done       -> return $ Skip    (Left t)
+
 -- Unfolding
 -- ---------
 
index 184ed58..bd316cd 100644 (file)
@@ -567,7 +567,21 @@ v ++ w = unstream (stream v Stream.++ stream w)
 -- | /O(n)/ Concatenate all vectors in the list
 concat :: Vector v a => [v a] -> v a
 {-# INLINE concat #-}
-concat vs = create (thawMany vs)
+-- concat vs = create (thawMany vs)
+concat vs = unstream (Stream.flatten mk step (Exact n) (Stream.fromList vs))
+  where
+    n = List.foldl' (\k v -> k + length v) 0 vs
+
+    {-# INLINE_INNER step #-}
+    step (v,i,k)
+      | i < k = case unsafeIndexM v i of
+                  Box x -> Stream.Yield x (v,i+1,k)
+      | otherwise = Stream.Done
+
+    {-# INLINE mk #-}
+    mk v = let k = length v
+           in
+           k `seq` (v,0,k)
 
 -- Monadic initialisation
 -- ----------------------