Reimplement concat
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sun, 8 Jan 2012 10:59:16 +0000 (10:59 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sun, 8 Jan 2012 10:59:16 +0000 (10:59 +0000)
Data/Vector/Fusion/Stream.hs
Data/Vector/Fusion/Stream/Monadic.hs
Data/Vector/Generic.hs

index 2c079cd..0b2ae0d 100644 (file)
@@ -68,7 +68,7 @@ module Data.Vector.Fusion.Stream (
 
   -- * Conversions
   toList, fromList, fromListN, unsafeFromList, liftStream,
-  fromVector, reVector,
+  fromVector, reVector, fromVectors,
 
   -- * Monadic combinators
   mapM, mapM_, zipWithM, zipWithM_, filterM, foldM, fold1M, foldM', fold1M',
@@ -637,6 +637,10 @@ reVector :: Stream u a -> Stream v a
 {-# INLINE reVector #-}
 reVector = M.reVector
 
+fromVectors :: Vector v a => [v a] -> Stream v a
+{-# INLINE fromVectors #-}
+fromVectors = M.fromVectors
+
 -- | 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 #-}
index 135360e..312ecd4 100644 (file)
@@ -73,13 +73,14 @@ module Data.Vector.Fusion.Stream.Monadic (
 
   -- * Conversions
   toList, fromList, fromListN, unsafeFromList,
-  fromVector, reVector
+  fromVector, reVector, fromVectors
 ) where
 
 import Data.Vector.Generic.Base
 import Data.Vector.Fusion.Stream.Size
 import Data.Vector.Fusion.Util ( Box(..), delay_inline )
 
+import qualified Data.List as List
 import Data.Char      ( ord )
 import GHC.Base       ( unsafeChr )
 import Control.Monad  ( liftM )
@@ -130,6 +131,24 @@ instance Functor (Step s) where
 
 data Unf m a = forall s. Unf (s -> m (Step s a)) s
 
+unvector :: (Monad m, Vector v a) => Unf m (Either a (v a)) -> Unf m a
+{-# INLINE unvector #-}
+unvector (Unf step s) = Unf step' (Left s)
+  where
+    step' (Left s) = do
+                       r <- step s
+                       case r of
+                         Yield (Left  x) s' -> return $ Yield x (Left s')
+                         Yield (Right v) s' -> basicLength v `seq`
+                                               return (Skip (Right (v,0,s')))
+                         Skip            s' -> return $ Skip (Left s')
+                         Done               -> return Done
+
+    step' (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))
+
 instance Monad m => Functor (Unf m) where
   {-# INLINE fmap #-}
   fmap f (Unf step s) = Unf step' s
@@ -1557,6 +1576,15 @@ fromVector v = v `seq` n `seq` Stream (Unf step 0) (Unf vstep True) (Exact n)
     vstep True  = return (Yield (Right v) False)
     vstep False = return Done
 
+fromVectors :: (Monad m, Vector v a) => [v a] -> Stream m v a
+{-# INLINE_STREAM fromVectors #-}
+fromVectors vs = Stream (unvector $ Unf step vs) (Unf step vs) (Exact n)
+  where
+    n = List.foldl' (\k v -> k + basicLength v) 0 vs
+
+    step [] = return Done
+    step (v:vs) = return $ Yield (Right v) vs
+
 reVector :: Monad m => Stream m u a -> Stream m v a
 {-# INLINE_STREAM reVector #-}
 reVector (Stream (Unf step s) _ n) = simple step s n
index 263c4e3..35acc3d 100644 (file)
@@ -676,6 +676,8 @@ 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 = unstream . Stream.fromVectors
+{-
 concat vs = unstream (Stream.flatten mk step (Exact n) (Stream.fromList vs))
   where
     n = List.foldl' (\k v -> k + length v) 0 vs
@@ -690,6 +692,7 @@ concat vs = unstream (Stream.flatten mk step (Exact n) (Stream.fromList vs))
     mk v = let k = length v
            in
            k `seq` (v,0,k)
+-}
 
 -- Monadic initialisation
 -- ----------------------