Heterogeneous array creation

Data/Primitive/Array.hs

@@ -1,6 +1,9 @@
 {-# LANGUAGE CPP, MagicHash, UnboxedTuples, DeriveDataTypeable, BangPatterns #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE TypeFamilies #-}
+#if __GLASGOW_HASKELL__ >= 708
+{-# LANGUAGE PolyKinds #-}
+#endif
 
 -- |
 -- Module      : Data.Primitive.Array
@@ -17,7 +20,7 @@ module Data.Primitive.Array (
   Array(..), MutableArray(..),
 
   newArray, readArray, writeArray, indexArray, indexArrayM,
-  freezeArray, thawArray, runArray,
+  freezeArray, thawArray, runArray, runArrays, runArraysHetOf, runArraysHetOfThen,
   unsafeFreezeArray, unsafeThawArray, sameMutableArray,
   copyArray, copyMutableArray,
   cloneArray, cloneMutableArray,
@@ -26,6 +29,7 @@ module Data.Primitive.Array (
   unsafeTraverseArray
 ) where
 
+import Data.Functor.Compose
 import Control.Monad.Primitive
 
 import GHC.Base  ( Int(..) )
@@ -341,7 +345,6 @@ emptyArray# _ = case emptyArray of Array ar -> ar
 {-# NOINLINE emptyArray# #-}
 #endif
 
-
 die :: String -> String -> a
 die fun problem = error $ "Data.Primitive.Array." ++ fun ++ ": " ++ problem
 
@@ -798,3 +801,74 @@ instance (Typeable s, Typeable a) => Data (MutableArray s a) where
   toConstr _ = error "toConstr"
   gunfold _ _ = error "gunfold"
   dataTypeOf _ = mkNoRepType "Data.Primitive.Array.MutableArray"
+
+-- | Create any number of arrays of the same type within an arbitrary
+-- 'Traversable' context. This will often be useful with traversables
+-- like  @(c,)@, @'Either' e@, @'Compose' (c,) ('Either' e)@, and
+-- @'Compose' ('Either' e) (c,)@. For a more general version, see
+-- 'runArraysHetOf'.
+runArrays
+  :: Traversable t
+  => (forall s. ST s (t (MutableArray s a)))
+  -> t (Array a)
+runArrays m = runST $ m >>= traverse unsafeFreezeArray
+
+-- | Create arbitrarily many arrays that may have different types.
+-- For a simpler but less general version, see 'runArrays'.
+--
+-- === __Examples__
+--
+-- ==== @'runArrays'@
+--
+-- @
+-- newtype Ha t a v = Ha {unHa :: t (v a)}
+-- runArrays m = unHa $ runArraysHetOf (\f (Ha t) -> Ha <$> traverse f t) (Ha <$> m)
+-- @
+--
+-- ==== @unzipArray@
+--
+-- @
+-- unzipArray :: Array (a, b) -> (Array a, Array b)
+-- unzipArray ar =
+--   unPair $ runArraysHetOf traversePair $ do
+--          xs <- newArray sz undefined
+--          ys <- newArray sz undefined
+--          let go k
+--                | k == sz = pure (Pair (xs, ys))
+--                | otherwise = do
+--                    (x,y) <- indexArrayM ar k
+--                    writeArray xs k x
+--                    writeArray ys k y
+--                    go (k + 1)
+--          go 0
+--   where sz = sizeofArray ar
+--
+-- data Pair ab v where
+--   Pair :: {unPair :: (v a, v b)} -> Pair (a,b) v
+--
+-- traversePair :: Applicative h => (forall x. f x -> h (g x)) -> Pair ab f -> h (Pair ab g)
+-- traversePair f (Pair (xs, ys)) = liftA2 (\x y -> Pair (x,y)) (f xs) (f ys)
+-- @
+runArraysHetOf
+  :: (forall h f g.
+       (Applicative h => (forall x. f x -> h (g x)) -> t f -> h (u g)))
+     -- ^ A rank-2 traversal
+  -> (forall s. ST s (t (MutableArray s)))
+     -- ^ An 'ST' action producing a rank-2 container of 'MutableArray's.
+  -> u Array
+runArraysHetOf trav m = runST $ m >>= trav unsafeFreezeArray
+
+-- | Similar to 'runArraysHetOf', but takes a function to traverse over the
+-- generated 'Array's as it freezes them.
+runArraysHetOfThen
+  :: Applicative q
+  => (forall h f g.
+       (Applicative h => (forall x. f x -> h (g x)) -> t f -> h (u g)))
+     -- ^ A rank-2 traversal
+  -> (forall x. Array x -> q (r x))
+     -- ^ A function to traverse over the container of 'Array's
+  -> (forall s. ST s (t (MutableArray s)))
+     -- ^ An 'ST' action producing a rank-2 container of 'MutableArray's.
+  -> q (u r)
+runArraysHetOfThen trav post m =
+    runST $ m >>= getCompose . trav (Compose . fmap post . unsafeFreezeArray)

Data/Primitive/SmallArray.hs

@@ -7,6 +7,9 @@
 {-# LANGUAGE DeriveDataTypeable #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE BangPatterns #-}
+#if __GLASGOW_HASKELL__ >= 708
+{-# LANGUAGE PolyKinds #-}
+#endif
 
 -- |
 -- Module : Data.Primitive.SmallArray
@@ -52,6 +55,9 @@ module Data.Primitive.SmallArray
   , unsafeFreezeSmallArray
   , thawSmallArray
   , runSmallArray
+  , runSmallArrays
+  , runSmallArraysHetOf
+  , runSmallArraysHetOfThen
   , unsafeThawSmallArray
   , sizeofSmallArray
   , sizeofSmallMutableArray
@@ -102,6 +108,7 @@ import qualified Data.Primitive.Array as Array
 #if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,4,0)
 import Data.Functor.Classes (Eq1(..),Ord1(..),Show1(..),Read1(..))
 #endif
+import Data.Functor.Compose
 
 #if HAVE_SMALL_ARRAY
 data SmallArray a = SmallArray (SmallArray# a)
@@ -940,3 +947,74 @@ smallArrayFromListN n l = SmallArray (Array.fromListN n l)
 -- | Create a 'SmallArray' from a list.
 smallArrayFromList :: [a] -> SmallArray a
 smallArrayFromList l = smallArrayFromListN (length l) l
+
+-- | Create any number of arrays of the same type within an arbitrary
+-- 'Traversable' context. This will often be useful with traversables
+-- like  @(c,)@, @'Either' e@, @'Compose' (c,) ('Either' e)@, and
+-- @'Compose' ('Either' e) (c,)@. For a more general version, see
+-- 'runArraysHetOf'.
+runSmallArrays
+  :: Traversable t
+  => (forall s. ST s (t (SmallMutableArray s a)))
+  -> t (SmallArray a)
+runSmallArrays m = runST $ m >>= traverse unsafeFreezeSmallArray
+
+-- | Create arbitrarily many arrays that may have different types. For
+-- a simpler but less general version, see 'runArrays'.
+--
+-- === __Examples__
+--
+-- ==== @'runSmallArrays'@
+--
+-- @
+-- newtype Ha t a v = Ha {unHa :: t (v a)}
+-- runSmallArrays m = unHa $ runSmallArraysHetOf (\f (Ha t) -> Ha <$> traverse f t) (Ha <$> m)
+-- @
+--
+-- ==== @unzipSmallArray@
+--
+-- @
+-- unzipSmallArray :: Array (a, b) -> (Array a, Array b)
+-- unzipSmallArray ar =
+--   unPair $ runSmallArraysHetOf traversePair $ do
+--          xs <- newSmallArray sz undefined
+--          ys <- newSmallArray sz undefined
+--          let go k
+--                | k == sz = pure (Pair (xs, ys))
+--                | otherwise = do
+--                    (x,y) <- indexSmallArrayM ar k
+--                    writeSmallArray xs k x
+--                    writeSmallArray ys k y
+--                    go (k + 1)
+--          go 0
+--   where sz = sizeofSmallArray ar
+--
+-- data Pair ab v where
+--   Pair :: {unPair :: (v a, v b)} -> Pair (a,b) v
+--
+-- traversePair :: Applicative h => (forall x. f x -> h (g x)) -> Pair ab f -> h (Pair ab g)
+-- traversePair f (Pair (xs, ys)) = liftA2 (\x y -> Pair (x,y)) (f xs) (f ys)
+-- @
+runSmallArraysHetOf
+  :: (forall h f g.
+       (Applicative h => (forall x. f x -> h (g x)) -> t f -> h (u g)))
+     -- ^ A rank-2 traversal
+  -> (forall s. ST s (t (SmallMutableArray s)))
+     -- ^ An 'ST' action producing a rank-2 container of 'MutableArray's.
+  -> u SmallArray
+runSmallArraysHetOf f m = runST $ m >>= f unsafeFreezeSmallArray
+
+-- | Similar to 'runSmallArraysHetOf', but takes a function to traverse over the
+-- generated 'SmallArray's as it freezes them.
+runSmallArraysHetOfThen
+  :: Applicative q
+  => (forall h f g.
+       (Applicative h => (forall x. f x -> h (g x)) -> t f -> h (u g)))
+     -- ^ A rank-2 traversal
+  -> (forall x. SmallArray x -> q (r x))
+     -- ^ A function to traverse over a container of 'Array's
+  -> (forall s. ST s (t (SmallMutableArray s)))
+     -- ^ An 'ST' action producing a rank-2 container of 'MutableArray's.
+  -> q (u r)
+runSmallArraysHetOfThen trav post m =
+  runST $ m >>= getCompose . trav (Compose . fmap post . unsafeFreezeSmallArray)