Add 'createT', allowing monadic initialisation of any Traversable container of vectors

Author: nshepperd

Data/Vector.hs

@@ -56,7 +56,7 @@ module Data.Vector (
   empty, singleton, replicate, generate, iterateN,
 
   -- ** Monadic initialisation
-  replicateM, generateM, create,
+  replicateM, generateM, create, createT,
 
   -- ** Unfolding
   unfoldr, unfoldrN,
@@ -687,6 +687,11 @@ create :: (forall s. ST s (MVector s a)) -> Vector a
 -- NOTE: eta-expanded due to http://hackage.haskell.org/trac/ghc/ticket/4120
 create p = G.create p
 
+-- | Execute the monadic action and freeze the resulting vectors.
+createT :: Traversable f => (forall s. ST s (f (MVector s a))) -> f (Vector a)
+{-# INLINE createT #-}
+createT p = G.createT p
+
 
 
 -- Restricting memory usage
@@ -1572,4 +1577,3 @@ unsafeCopy = G.unsafeCopy
 copy :: PrimMonad m => MVector (PrimState m) a -> Vector a -> m ()
 {-# INLINE copy #-}
 copy = G.copy
-

Data/Vector/Generic.hs

@@ -39,7 +39,7 @@ module Data.Vector.Generic (
   empty, singleton, replicate, generate, iterateN,
 
   -- ** Monadic initialisation
-  replicateM, generateM, create,
+  replicateM, generateM, create, createT,
 
   -- ** Unfolding
   unfoldr, unfoldrN,
@@ -709,6 +709,11 @@ create :: Vector v a => (forall s. ST s (Mutable v s a)) -> v a
 {-# INLINE create #-}
 create p = new (New.create p)
 
+-- | Execute the monadic action and freeze the resulting vectors.
+createT :: (Traversable f, Vector v a) => (forall s. ST s (f (Mutable v s a))) -> f (v a)
+{-# INLINE createT #-}
+createT p = runST (p >>= traverse unsafeFreeze)
+
 -- Restricting memory usage
 -- ------------------------
 
@@ -2084,4 +2089,3 @@ dataCast :: (Vector v a, Data a, Typeable1 v, Typeable1 t)
          => (forall d. Data  d => c (t d)) -> Maybe  (c (v a))
 {-# INLINE dataCast #-}
 dataCast f = gcast1 f
-

Data/Vector/Primitive.hs

@@ -42,7 +42,7 @@ module Data.Vector.Primitive (
   empty, singleton, replicate, generate, iterateN,
 
   -- ** Monadic initialisation
-  replicateM, generateM, create,
+  replicateM, generateM, create, createT,
 
   -- ** Unfolding
   unfoldr, unfoldrN,
@@ -606,6 +606,11 @@ create :: Prim a => (forall s. ST s (MVector s a)) -> Vector a
 -- NOTE: eta-expanded due to http://hackage.haskell.org/trac/ghc/ticket/4120
 create p = G.create p
 
+-- | Execute the monadic action and freeze the resulting vectors.
+createT :: (Traversable f, Prim a) => (forall s. ST s (f (MVector s a))) -> f (Vector a)
+{-# INLINE createT #-}
+createT p = G.createT p
+
 -- Restricting memory usage
 -- ------------------------
 
@@ -1336,5 +1341,3 @@ unsafeCopy = G.unsafeCopy
 copy :: (Prim a, PrimMonad m) => MVector (PrimState m) a -> Vector a -> m ()
 {-# INLINE copy #-}
 copy = G.copy
-
-

Data/Vector/Storable.hs

@@ -39,7 +39,7 @@ module Data.Vector.Storable (
   empty, singleton, replicate, generate, iterateN,
 
   -- ** Monadic initialisation
-  replicateM, generateM, create,
+  replicateM, generateM, create, createT,
 
   -- ** Unfolding
   unfoldr, unfoldrN,
@@ -616,6 +616,11 @@ create :: Storable a => (forall s. ST s (MVector s a)) -> Vector a
 -- NOTE: eta-expanded due to http://hackage.haskell.org/trac/ghc/ticket/4120
 create p = G.create p
 
+-- | Execute the monadic action and freeze the resulting vectors.
+createT :: (Traversable f, Storable a) => (forall s. ST s (f (MVector s a))) -> f (Vector a)
+{-# INLINE createT #-}
+createT p = G.createT p
+
 -- Restricting memory usage
 -- ------------------------
 
@@ -1432,5 +1437,3 @@ unsafeToForeignPtr0 (Vector n fp) = (fp, n)
 unsafeWith :: Storable a => Vector a -> (Ptr a -> IO b) -> IO b
 {-# INLINE unsafeWith #-}
 unsafeWith (Vector _ fp) = withForeignPtr fp
-
-

Data/Vector/Unboxed.hs

@@ -62,7 +62,7 @@ module Data.Vector.Unboxed (
   empty, singleton, replicate, generate, iterateN,
 
   -- ** Monadic initialisation
-  replicateM, generateM, create,
+  replicateM, generateM, create, createT,
 
   -- ** Unfolding
   unfoldr, unfoldrN,
@@ -585,6 +585,11 @@ create :: Unbox a => (forall s. ST s (MVector s a)) -> Vector a
 -- NOTE: eta-expanded due to http://hackage.haskell.org/trac/ghc/ticket/4120
 create p = G.create p
 
+-- | Execute the monadic action and freeze the resulting vectors.
+createT :: (Traversable f, Unbox a) => (forall s. ST s (f (MVector s a))) -> f (Vector a)
+{-# INLINE createT #-}
+createT p = G.createT p
+
 -- Restricting memory usage
 -- ------------------------
 
@@ -1431,4 +1436,3 @@ copy = G.copy
 
 #define DEFINE_IMMUTABLE
 #include "unbox-tuple-instances"
-

tests/Tests/Vector.hs

@@ -105,6 +105,7 @@ testPolymorphicFunctions _ = $(testProperties [
         'prop_generate, 'prop_iterateN,
 
         -- Monadic initialisation (FIXME)
+        'prop_createT,
         {- 'prop_replicateM, 'prop_generateM, 'prop_create, -}
 
         -- Unfolding (FIXME)
@@ -217,6 +218,9 @@ testPolymorphicFunctions _ = $(testProperties [
     prop_iterateN  :: P (Int -> (a -> a) -> a -> v a)
               = (\n _ _ -> n < 1000) ===> V.iterateN `eq` (\n f -> take n . iterate f)
 
+    prop_createT :: P ((a, v a) -> (a, v a))
+    prop_createT = (\v -> V.createT (traverse V.thaw v)) `eq` id
+
     prop_head      :: P (v a -> a) = not . V.null ===> V.head `eq` head
     prop_last      :: P (v a -> a) = not . V.null ===> V.last `eq` last
     prop_index        = \xs ->
@@ -649,4 +653,3 @@ tests = [
          testGroup "Data.Vector.Unboxed.Vector (Int,Bool,Int)" (testTupleUnboxedVector (undefined :: Data.Vector.Unboxed.Vector (Int,Bool,Int)))
 
     ]
-