Remove thawMany
[darcs-mirrors/vector.git] / Data / Vector / Generic.hs
index d2de61a..1b795a9 100644 (file)
@@ -48,7 +48,7 @@ module Data.Vector.Generic (
   enumFromN, enumFromStepN, enumFromTo, enumFromThenTo,
 
   -- ** Concatenation
-  cons, snoc, (++),
+  cons, snoc, (++), concat,
 
   -- ** Restricting memory usage
   force,
@@ -126,8 +126,11 @@ module Data.Vector.Generic (
   -- ** Lists
   toList, fromList, fromListN,
 
+  -- ** Different vector types
+  convert,
+
   -- ** Mutable vectors
-  copy, unsafeCopy,
+  thaw, copy, unsafeCopy,
 
   -- * Fusion support
 
@@ -155,7 +158,7 @@ import qualified Data.Vector.Generic.New as New
 import           Data.Vector.Generic.New ( New )
 
 import qualified Data.Vector.Fusion.Stream as Stream
-import           Data.Vector.Fusion.Stream ( Stream, MStream, inplace )
+import           Data.Vector.Fusion.Stream ( Stream, MStream, inplace, liftStream )
 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
 import           Data.Vector.Fusion.Stream.Size
 import           Data.Vector.Fusion.Util
@@ -163,11 +166,12 @@ import           Data.Vector.Fusion.Util
 import Control.Monad.ST ( ST, runST )
 import Control.Monad.Primitive
 import qualified Control.Monad as Monad
+import qualified Data.List as List
 import Prelude hiding ( length, null,
-                        replicate, (++),
+                        replicate, (++), concat,
                         head, last,
                         init, tail, take, drop, reverse,
-                        map, concatMap,
+                        map, concat, concatMap,
                         zipWith, zipWith3, zip, zip3, unzip, unzip3,
                         filter, takeWhile, dropWhile, span, break,
                         elem, notElem,
@@ -324,27 +328,35 @@ unsafeLastM :: (Vector v a, Monad m) => v a -> m a
 {-# INLINE_STREAM unsafeLastM #-}
 unsafeLastM v = unsafeIndexM v (length v - 1)
 
--- FIXME: the rhs of these rules are lazy in the stream which is WRONG
-{- RULES
+{-# RULES
+
+"indexM/unstream [Vector]" forall s i.
+  indexM (new (New.unstream s)) i = liftStream s MStream.!! i
 
-"indexM/unstream [Vector]" forall v i s.
-  indexM (new' v (New.unstream s)) i = return (s Stream.!! i)
+"headM/unstream [Vector]" forall s.
+  headM (new (New.unstream s)) = MStream.head (liftStream s)
 
-"headM/unstream [Vector]" forall v s.
-  headM (new' v (New.unstream s)) = return (Stream.head s)
+"lastM/unstream [Vector]" forall s.
+  lastM (new (New.unstream s)) = MStream.last (liftStream s)
 
-"lastM/unstream [Vector]" forall v s.
-  lastM (new' v (New.unstream s)) = return (Stream.last s)
+"unsafeIndexM/unstream [Vector]" forall s i.
+  unsafeIndexM (new (New.unstream s)) i = liftStream s MStream.!! i
 
- -}
+"unsafeHeadM/unstream [Vector]" forall s.
+  unsafeHeadM (new (New.unstream s)) = MStream.head (liftStream s)
+
+"unsafeLastM/unstream [Vector]" forall s.
+  unsafeLastM (new (New.unstream s)) = MStream.last (liftStream s)
+
+  #-}
 
 -- Extracting subvectors (slicing)
 -- -------------------------------
 
 -- | /O(1)/ Yield a slice of the vector without copying it. The vector must
 -- contain at least @i+n@ elements.
-slice :: Vector v a => Int   -- ^ @i@, starting index
-                    -> Int   -- ^ @n@, length
+slice :: Vector v a => Int   -- ^ @i@ starting index
+                    -> Int   -- ^ @n@ length
                     -> v a
                     -> v a
 {-# INLINE_STREAM slice #-}
@@ -552,10 +564,29 @@ infixr 5 ++
 {-# INLINE (++) #-}
 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 = 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
 -- ----------------------
 
--- | /O(n)/ Perform the monadic action the given number of times and store the
+-- | /O(n)/ Execute the monadic action the given number of times and store the
 -- results in a vector.
 replicateM :: (Monad m, Vector v a) => Int -> m a -> m (v a)
 -- FIXME: specialise for ST and IO?
@@ -668,7 +699,7 @@ unsafeUpdate_stream = modifyWithStream M.unsafeUpdate
 --
 -- > accum (+) <5,9,2> [(2,4),(1,6),(0,3),(1,7)] = <5+3, 9+6+7, 2+4>
 accum :: Vector v a
-      => (a -> b -> a) -- ^ @f@ accumulating function
+      => (a -> b -> a) -- ^ accumulating function @f@
       -> v a           -- ^ initial vector (of length @m@)
       -> [(Int,b)]     -- ^ list of index/value pairs (of length @n@)
       -> v a
@@ -680,7 +711,7 @@ accum f v us = accum_stream f v (Stream.fromList us)
 --
 -- > accumulate (+) <5,9,2> <(2,4),(1,6),(0,3),(1,7)> = <5+3, 9+6+7, 2+4>
 accumulate :: (Vector v a, Vector v (Int, b))
-           => (a -> b -> a) -- ^ @f@ accumulating function
+           => (a -> b -> a) -- ^ accumulating function @f@
            -> v a           -- ^ initial vector (of length @m@)
            -> v (Int,b)     -- ^ vector of index/value pairs (of length @n@)
            -> v a
@@ -701,7 +732,7 @@ accumulate f v us = accum_stream f v (stream us)
 -- accumulate_ f as is bs = 'accumulate' f as ('zip' is bs)
 -- @
 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
-                => (a -> b -> a) -- ^ @f@ accumulating function
+                => (a -> b -> a) -- ^ accumulating function @f@
                 -> v a           -- ^ initial vector (of length @m@)
                 -> v Int         -- ^ index vector (of length @n1@)
                 -> v b           -- ^ value vector (of length @n2@)
@@ -814,7 +845,9 @@ imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
 -- | Map a function over a vector and concatenate the results.
 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
 {-# INLINE concatMap #-}
-concatMap f = unstream . Stream.concatMap (stream . f) . stream
+-- NOTE: We can't fuse concatMap anyway so don't pretend we do.
+-- concatMap f = unstream . Stream.concatMap (stream . f) . stream
+concatMap f = concat . Stream.toList . Stream.map f . stream
 
 -- Monadic mapping
 -- ---------------
@@ -833,13 +866,13 @@ mapM_ :: (Monad m, Vector v a) => (a -> m b) -> v a -> m ()
 mapM_ f = Stream.mapM_ f . stream
 
 -- | /O(n)/ Apply the monadic action to all elements of the vector, yielding a
--- vector of results. Equvalent to @flip mapM@.
+-- vector of results. Equvalent to @flip 'mapM'@.
 forM :: (Monad m, Vector v a, Vector v b) => v a -> (a -> m b) -> m (v b)
 {-# INLINE forM #-}
 forM as f = mapM f as
 
 -- | /O(n)/ Apply the monadic action to all elements of a vector and ignore the
--- results. Equivalent to @flip mapM_@.
+-- results. Equivalent to @flip 'mapM_'@.
 forM_ :: (Monad m, Vector v a) => v a -> (a -> m b) -> m ()
 {-# INLINE forM_ #-}
 forM_ as f = mapM_ f as
@@ -1522,9 +1555,48 @@ fromListN :: Vector v a => Int -> [a] -> v a
 {-# INLINE fromListN #-}
 fromListN n = unstream . Stream.fromListN n
 
+-- Conversions - Immutable vectors
+-- -------------------------------
+
+-- | /O(n)/ Convert different vector types
+convert :: (Vector v a, Vector w a) => v a -> w a
+{-# INLINE convert #-}
+convert = unstream . stream
+
 -- Conversions - Mutable vectors
 -- -----------------------------
 
+-- | /O(n)/ Yield a mutable copy of the immutable vector.
+thaw :: (PrimMonad m, Vector v a) => v a -> m (Mutable v (PrimState m) a)
+{-# INLINE_STREAM thaw #-}
+thaw v = do
+           mv <- M.unsafeNew (length v)
+           unsafeCopy mv v
+           return mv
+
+{-
+-- | /O(n)/ Yield a mutable vector containing copies of each vector in the
+-- list.
+thawMany :: (PrimMonad m, Vector v a) => [v a] -> m (Mutable v (PrimState m) a)
+{-# INLINE_STREAM thawMany #-}
+-- FIXME: add rule for (stream (new (New.create (thawMany vs))))
+-- NOTE: We don't try to consume the list lazily as this wouldn't significantly
+-- change the space requirements anyway.
+thawMany vs = do
+                mv <- M.new n
+                thaw_loop mv vs
+                return mv
+  where
+    n = List.foldl' (\k v -> k + length v) 0 vs
+
+    thaw_loop mv [] = mv `seq` return ()
+    thaw_loop mv (v:vs)
+      = do
+          let n = length v
+          unsafeCopy (M.unsafeTake n mv) v
+          thaw_loop (M.unsafeDrop n mv) vs
+-}
+
 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
 -- have the same length.
 copy