Add the dreaded concatMap
authorMax Bolingbroke <batterseapower@hotmail.com>
Sun, 8 Feb 2009 11:17:05 +0000 (11:17 +0000)
committerMax Bolingbroke <batterseapower@hotmail.com>
Sun, 8 Feb 2009 11:17:05 +0000 (11:17 +0000)
Data/Vector/Fusion/Stream.hs
Data/Vector/Fusion/Stream/Monadic.hs
Data/Vector/IVector.hs

index 661fe33..ba3e24c 100644 (file)
@@ -46,7 +46,7 @@ module Data.Vector.Fusion.Stream (
   foldl, foldl1, foldl', foldl1', foldr, foldr1,
 
   -- * Specialised folds
-  and, or,
+  and, or, concatMap,
 
   -- * Unfolding
   unfoldr,
@@ -73,7 +73,7 @@ import Prelude hiding ( length, null,
                         filter, takeWhile, dropWhile,
                         elem, notElem,
                         foldl, foldl1, foldr, foldr1,
-                        and, or,
+                        and, or, concatMap,
                         mapM_ )
 
 
@@ -303,6 +303,10 @@ and = unId . M.and
 or :: Stream Bool -> Bool
 or = unId . M.or
 
+concatMap :: (a -> Stream b) -> Stream a -> Stream b
+{-# INLINE concatMap #-}
+concatMap = M.concatMap
+
 -- Unfolding
 -- ---------
 
index b149caa..ade0c0c 100644 (file)
@@ -47,7 +47,7 @@ module Data.Vector.Fusion.Stream.Monadic (
   foldr, foldrM, foldr1, foldr1M,
 
   -- * Specialised folds
-  and, or,
+  and, or, concatMap, concatMapM,
 
   -- * Unfolding
   unfoldr, unfoldrM,
@@ -70,7 +70,7 @@ import Prelude hiding ( length, null,
                         filter, takeWhile, dropWhile,
                         elem, notElem,
                         foldl, foldl1, foldr, foldr1,
-                        and, or )
+                        and, or, concatMap )
 import qualified Prelude
 
 -- | Result of taking a single step in a stream
@@ -653,6 +653,29 @@ or (Stream step s _) = or_go s
                   Skip        s' -> or_go s'
                   Done           -> return False
 
+concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
+{-# INLINE concatMap #-}
+concatMap f = concatMapM (return . f)
+
+concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
+{-# INLINE_STREAM concatMapM #-}
+concatMapM f (Stream step s _) = Stream concatMap_go (Left s) Unknown
+  where
+    concatMap_go (Left s) = do
+        r <- step s
+        case r of
+            Yield a s' -> do
+                b_stream <- f a
+                return $ Skip (Right (b_stream, s'))
+            Skip    s' -> return $ Skip (Left s')
+            Done       -> return Done
+    concatMap_go (Right (Stream inner_step inner_s sz, s)) = do
+        r <- inner_step inner_s
+        case r of
+            Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s' sz, s))
+            Skip    inner_s' -> return $ Skip (Right (Stream inner_step inner_s' sz, s))
+            Done             -> return $ Skip (Left s)
+
 -- Unfolding
 -- ---------
 
index 0c9a59a..97e426d 100644 (file)
@@ -49,7 +49,7 @@ module Data.Vector.IVector (
   foldl, foldl1, foldl', foldl1', foldr, foldr1,
  
   -- * Specialised folds
-  and, or,
+  and, or, concatMap,
   
   -- * Enumeration
   enumFromTo,
@@ -97,7 +97,7 @@ import Prelude hiding ( length, null,
                         filter, takeWhile, dropWhile,
                         elem, notElem,
                         foldl, foldl1, foldr, foldr1,
-                        and, or,
+                        and, or, concatMap,
                         enumFromTo )
 
 -- | Class of immutable vectors.
@@ -502,6 +502,10 @@ or :: IVector v Bool => v Bool -> Bool
 {-# INLINE or #-}
 or = Stream.or . stream
 
+concatMap :: (IVector v a, IVector v b) => (a -> v b) -> v a -> v b
+{-# INLINE concatMap #-}
+concatMap f = unstream . Stream.concatMap (stream . f) . stream
+
 -- Enumeration
 -- -----------