Add thaw and thawMany
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Wed, 5 May 2010 14:10:28 +0000 (14:10 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Wed, 5 May 2010 14:10:28 +0000 (14:10 +0000)
Data/Vector.hs
Data/Vector/Generic.hs
Data/Vector/Primitive.hs
Data/Vector/Storable.hs
Data/Vector/Unboxed.hs

index d7022a2..ec67c15 100644 (file)
@@ -137,7 +137,7 @@ module Data.Vector (
   toList, fromList, fromListN,
 
   -- ** Mutable vectors
-  copy, unsafeCopy
+  thaw, thawMany, copy, unsafeCopy
 ) where
 
 import qualified Data.Vector.Generic as G
@@ -1277,6 +1277,17 @@ fromListN = G.fromListN
 -- Conversions - Mutable vectors
 -- -----------------------------
 
+-- | /O(n)/ Yield a mutable copy of the immutable vector.
+thaw :: PrimMonad m => Vector a -> m (MVector (PrimState m) a)
+{-# INLINE thaw #-}
+thaw = G.thaw
+
+-- | /O(n)/ Yield a mutable vector containing copies of each vector in the
+-- list.
+thawMany :: PrimMonad m => [Vector a] -> m (MVector (PrimState m) a)
+{-# INLINE thawMany #-}
+thawMany = G.thawMany
+
 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
 -- have the same length.
 unsafeCopy :: PrimMonad m => MVector (PrimState m) a -> Vector a -> m ()
index 8bff449..9c9fbe0 100644 (file)
@@ -127,7 +127,7 @@ module Data.Vector.Generic (
   toList, fromList, fromListN,
 
   -- ** Mutable vectors
-  copy, unsafeCopy,
+  thaw, thawMany, copy, unsafeCopy,
 
   -- * Fusion support
 
@@ -163,6 +163,7 @@ 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, (++),
                         head, last,
@@ -1533,6 +1534,35 @@ fromListN n = unstream . Stream.fromListN n
 -- 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
index 7c0a97e..4694e12 100644 (file)
@@ -126,7 +126,7 @@ module Data.Vector.Primitive (
   toList, fromList, fromListN,
 
   -- ** Mutable vectors
-  copy, unsafeCopy
+  thaw, thawMany, copy, unsafeCopy
 ) where
 
 import qualified Data.Vector.Generic           as G
@@ -1190,6 +1190,17 @@ fromListN = G.fromListN
 -- Conversions - Mutable vectors
 -- -----------------------------
 
+-- | /O(n)/ Yield a mutable copy of the immutable vector.
+thaw :: (Prim a, PrimMonad m) => Vector a -> m (MVector (PrimState m) a)
+{-# INLINE thaw #-}
+thaw = G.thaw
+
+-- | /O(n)/ Yield a mutable vector containing copies of each vector in the
+-- list.
+thawMany :: (Prim a, PrimMonad m) => [Vector a] -> m (MVector (PrimState m) a)
+{-# INLINE thawMany #-}
+thawMany = G.thawMany
+
 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
 -- have the same length.
 unsafeCopy
index 4b88abe..4da3532 100644 (file)
@@ -123,7 +123,7 @@ module Data.Vector.Storable (
   toList, fromList, fromListN,
 
   -- ** Mutable vectors
-  copy, unsafeCopy,
+  thaw, thawMany, copy, unsafeCopy,
 
   -- * Raw pointers
   unsafeFromForeignPtr, unsafeToForeignPtr, unsafeWith
@@ -1232,6 +1232,18 @@ fromListN = G.fromListN
 -- Conversions - Mutable vectors
 -- -----------------------------
 
+-- | /O(n)/ Yield a mutable copy of the immutable vector.
+thaw :: (Storable a, PrimMonad m) => Vector a -> m (MVector (PrimState m) a)
+{-# INLINE thaw #-}
+thaw = G.thaw
+
+-- | /O(n)/ Yield a mutable vector containing copies of each vector in the
+-- list.
+thawMany
+  :: (Storable a, PrimMonad m) => [Vector a] -> m (MVector (PrimState m) a)
+{-# INLINE thawMany #-}
+thawMany = G.thawMany
+
 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
 -- have the same length.
 unsafeCopy
index 6345b58..ef9b0a5 100644 (file)
@@ -150,7 +150,7 @@ module Data.Vector.Unboxed (
   toList, fromList, fromListN,
 
   -- ** Mutable vectors
-  copy, unsafeCopy
+  thaw, thawMany, copy, unsafeCopy
 ) where
 
 import Data.Vector.Unboxed.Base
@@ -1222,6 +1222,17 @@ fromListN = G.fromListN
 -- Conversions - Mutable vectors
 -- -----------------------------
 
+-- | /O(n)/ Yield a mutable copy of the immutable vector.
+thaw :: (Unbox a, PrimMonad m) => Vector a -> m (MVector (PrimState m) a)
+{-# INLINE thaw #-}
+thaw = G.thaw
+
+-- | /O(n)/ Yield a mutable vector containing copies of each vector in the
+-- list.
+thawMany :: (Unbox a, PrimMonad m) => [Vector a] -> m (MVector (PrimState m) a)
+{-# INLINE thawMany #-}
+thawMany = G.thawMany
+
 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
 -- have the same length.
 unsafeCopy