Merge pull request #89 from nshepperd/createT

Add 'createT', allowing monadic initialisation of any Traversable con…

Author: dolio

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,
@@ -700,6 +700,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

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,
@@ -716,6 +716,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
 -- ------------------------
 

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,
@@ -616,6 +616,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
 -- ------------------------
 

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,
@@ -626,6 +626,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
 -- ------------------------
 

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,
@@ -595,6 +595,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
 -- ------------------------
 

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)
@@ -219,6 +220,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 ->
@@ -660,4 +664,3 @@ tests = [
          testGroup "Data.Vector.Unboxed.Vector (Int,Bool,Int)" (testTupleUnboxedVector (undefined :: Data.Vector.Unboxed.Vector (Int,Bool,Int)))
 
     ]
-