Rename unsafeFreeze to basicUnsafeFreeze and add unsafeFreeze as a free function
[darcs-mirrors/vector.git] / Data / Vector / Generic.hs
index 8bff449..dfe41d0 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,
+  unsafeFreeze, thaw, copy, unsafeCopy,
 
   -- * Fusion support
 
@@ -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,
@@ -560,6 +564,25 @@ 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
 -- ----------------------
 
@@ -769,19 +792,37 @@ backpermute :: (Vector v a, Vector v Int)
 -- does not retain references to the original one even if it is lazy in its
 -- elements. This would not be the case if we simply used map (v!)
 backpermute v is = seq v
+                 $ seq n
                  $ unstream
                  $ Stream.unbox
-                 $ Stream.map (indexM v)
+                 $ Stream.map index
                  $ stream is
+  where
+    n = length v
+
+    {-# INLINE index #-}
+    -- NOTE: we do it this way to avoid triggering LiberateCase on n in
+    -- polymorphic code
+    index i = BOUNDS_CHECK(checkIndex) "backpermute" i n
+            $ basicUnsafeIndexM v i
 
 -- | Same as 'backpermute' but without bounds checking.
 unsafeBackpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
 {-# INLINE unsafeBackpermute #-}
 unsafeBackpermute v is = seq v
+                       $ seq n
                        $ unstream
                        $ Stream.unbox
-                       $ Stream.map (unsafeIndexM v)
+                       $ Stream.map index
                        $ stream is
+  where
+    n = length v
+
+    {-# INLINE index #-}
+    -- NOTE: we do it this way to avoid triggering LiberateCase on n in
+    -- polymorphic code
+    index i = UNSAFE_CHECK(checkIndex) "unsafeBackpermute" i n
+            $ basicUnsafeIndexM v i
 
 -- Safe destructive updates
 -- ------------------------
@@ -822,7 +863,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
 -- ---------------
@@ -1530,9 +1573,55 @@ 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(1)/ Unsafe convert a mutable vector to an immutable one without
+-- copying. The mutable vector may not be used after this operation.
+unsafeFreeze
+  :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m (v a)
+{-# INLINE unsafeFreeze #-}
+unsafeFreeze = basicUnsafeFreeze
+
+-- | /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
@@ -1557,7 +1646,7 @@ unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
 -- | /O(1)/ Convert a vector to a 'Stream'
 stream :: Vector v a => v a -> Stream a
 {-# INLINE_STREAM stream #-}
-stream v = v `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
+stream v = v `seq` n `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
   where
     n = length v
 
@@ -1596,7 +1685,7 @@ unstream s = new (New.unstream s)
 -- | /O(1)/ Convert a vector to a 'Stream', proceeding from right to left
 streamR :: Vector v a => v a -> Stream a
 {-# INLINE_STREAM streamR #-}
-streamR v = v `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
+streamR v = v `seq` n `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
   where
     n = length v