[commit: vector] : Faster concatMap (427c50b)

[Geoffrey Mainland]

Repository : http://darcs.haskell.org/ghc.git/

On branch  : 

http://hackage.haskell.org/trac/ghc/changeset/427c50b6a1c08744e0833004432a44ba588f81ef

>---------------------------------------------------------------

commit 427c50b6a1c08744e0833004432a44ba588f81ef
Author: Roman Leshchinskiy <rl at cse.unsw.edu.au>
Date:   Sun Jan 29 00:57:02 2012 +0000

    Faster concatMap

>---------------------------------------------------------------

 Data/Vector/Fusion/Stream.hs         |    6 +++++-
 Data/Vector/Fusion/Stream/Monadic.hs |   31 ++++++++++++++++++++++++++++++-
 Data/Vector/Generic.hs               |   10 +++++++++-
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/Data/Vector/Fusion/Stream.hs b/Data/Vector/Fusion/Stream.hs
index 64e64c7..96f5248 100644
--- a/Data/Vector/Fusion/Stream.hs
+++ b/Data/Vector/Fusion/Stream.hs
@@ -68,7 +68,7 @@ module Data.Vector.Fusion.Stream (
 
   -- * Conversions
   toList, fromList, fromListN, unsafeFromList, liftStream,
-  fromVector, reVector, fromVectors,
+  fromVector, reVector, fromVectors, fromVectorStream,
 
   -- * Monadic combinators
   mapM, mapM_, zipWithM, zipWithM_, filterM, foldM, fold1M, foldM', fold1M',
@@ -643,6 +643,10 @@ fromVectors :: Vector v a => [v a] -> Stream v a
 {-# INLINE fromVectors #-}
 fromVectors = M.fromVectors
 
+fromVectorStream :: Vector v a => Stream u (v a) -> Stream v a
+{-# INLINE fromVectorStream #-}
+fromVectorStream = M.fromVectorStream
+
 -- | Create a 'Stream' of values from a 'Stream' of streamable things
 flatten :: (a -> s) -> (s -> Step s b) -> Size -> Stream v a -> Stream v b
 {-# INLINE_STREAM flatten #-}
diff --git a/Data/Vector/Fusion/Stream/Monadic.hs b/Data/Vector/Fusion/Stream/Monadic.hs
index a754fb2..eb31806 100644
--- a/Data/Vector/Fusion/Stream/Monadic.hs
+++ b/Data/Vector/Fusion/Stream/Monadic.hs
@@ -73,7 +73,7 @@ module Data.Vector.Fusion.Stream.Monadic (
 
   -- * Conversions
   toList, fromList, fromListN, unsafeFromList,
-  fromVector, reVector, fromVectors
+  fromVector, reVector, fromVectors, fromVectorStream
 ) where
 
 import Data.Vector.Generic.Base
@@ -1622,6 +1622,35 @@ fromVectors vs = Stream (Unf pstep (Left vs))
     vstep (v:vs) = return $ Yield (Chunk (basicLength v)
                                          (\mv -> basicUnsafeCopy mv v)) vs
 
+
+fromVectorStream :: (Monad m, Vector v a) => Stream m u (v a) -> Stream m v a
+{-# INLINE_STREAM fromVectorStream #-}
+fromVectorStream Stream{sElems = Unf step s}
+  = Stream (Unf pstep (Left s))
+           (Unf vstep s)
+           Unknown
+  where
+    pstep (Left s) = do
+      r <- step s
+      case r of
+        Yield v s' -> basicLength v `seq` return (Skip (Right (v,0,s')))
+        Skip    s' -> return (Skip (Left s'))
+        Done       -> return Done
+
+    pstep (Right (v,i,s))
+      | i >= basicLength v = return (Skip (Left s))
+      | otherwise          = case basicUnsafeIndexM v i of
+                               Box x -> return (Yield x (Right (v,i+1,s)))
+
+
+    vstep s = do
+      r <- step s
+      case r of
+        Yield v s' -> return (Yield (Chunk (basicLength v)
+                                           (\mv -> basicUnsafeCopy mv v)) s')
+        Skip    s' -> return (Skip s')
+        Done       -> return Done
+
 reVector :: Monad m => Stream m u a -> Stream m v a
 {-# INLINE_STREAM reVector #-}
 reVector Stream{sElems = Unf step s, sSize = n} = simple step s n
diff --git a/Data/Vector/Generic.hs b/Data/Vector/Generic.hs
index 53e2daa..06c5492 100644
--- a/Data/Vector/Generic.hs
+++ b/Data/Vector/Generic.hs
@@ -994,7 +994,8 @@ concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
 -- Slowest
 -- concatMap f = unstream . Stream.concatMap (stream . f) . stream
 
--- Seems to be fastest
+-- Used to be fastest
+{-
 concatMap f = unstream
             . Stream.flatten mk step Unknown
             . stream
@@ -1010,6 +1011,13 @@ concatMap f = unstream
                k = length v
            in
            k `seq` (v,0,k)
+-}
+
+-- This seems to be fastest now
+concatMap f = unstream
+            . Stream.fromVectorStream
+            . Stream.map f
+            . stream
 
 -- Monadic mapping
 -- ---------------