dph-prim-seq: add ArrayArray module
authorBen Lippmeier <benl@ouroborus.net>
Fri, 16 Dec 2011 01:45:42 +0000 (12:45 +1100)
committerBen Lippmeier <benl@ouroborus.net>
Fri, 16 Dec 2011 01:45:42 +0000 (12:45 +1100)
dph-prim-seq/Data/Array/Parallel/Unlifted/ArrayArray.hs [new file with mode: 0644]
dph-prim-seq/dph-prim-seq.cabal

diff --git a/dph-prim-seq/Data/Array/Parallel/Unlifted/ArrayArray.hs b/dph-prim-seq/Data/Array/Parallel/Unlifted/ArrayArray.hs
new file mode 100644 (file)
index 0000000..9903774
--- /dev/null
@@ -0,0 +1,94 @@
+
+-- | Nested unboxed arrays.
+--  
+--   TODO: move this into the Data.Primitive library.
+module Data.Array.Parallel.Unlifted.ArrayArray
+        ( MutableArrayArray (..)
+        , ArrayArray(..)
+        , newArrayArray
+        , writeArrayArrayMut
+        , writeArrayArray
+        , readArrayArray
+        , indexArrayArray
+        , unsafeFreezeArrayArray
+        , unsafeDeepFreezeArrayArray)
+where  
+import GHC.Prim
+import GHC.Base
+import GHC.ST
+import Data.Primitive.ByteArray
+
+
+data MutableArrayArray s e 
+        = MutableArrayArray (MutableArrayArray# s)
+
+data ArrayArray e 
+        = ArrayArray ArrayArray#
+
+
+-- | Create an `ArrayArray` with the given number of elements.
+newArrayArray :: Int -> ST s (MutableArrayArray s e)
+newArrayArray (I# n#)
+  = ST $ \s# -> case newArrayArray# n# s# of
+           (# s'#, arr# #) -> (# s'#, MutableArrayArray arr# #)
+{-# INLINE newArrayArray #-}
+
+
+-- | Write a `MutableByteArray` to an `MutableArrayArray`.
+writeArrayArrayMut :: MutableArrayArray s (MutableByteArray s) -> Int -> MutableByteArray s -> ST s ()
+writeArrayArrayMut (MutableArrayArray arrs#) (I# i#) (MutableByteArray mba#)
+  = ST $ \s# -> case writeMutableByteArrayArray# arrs# i# mba# s# of
+           s'# -> (# s'#, () #)
+{-# INLINE writeArrayArrayMut #-}
+
+
+-- | Write a `ByteArray` to a `MutableArrayArray`.
+writeArrayArray :: MutableArrayArray s ByteArray -> Int -> ByteArray -> ST s ()
+writeArrayArray (MutableArrayArray arrs#) (I# i#) (ByteArray ba#)
+  = ST $ \s# -> case writeByteArrayArray# arrs# i# ba# s# of
+           s'# -> (# s'#, () #)
+{-# INLINE writeArrayArray #-}
+
+
+-- | Read a `MutableByteArray` from a `MutableArrayArray`.
+readArrayArray :: MutableArrayArray s (MutableByteArray s) -> Int -> ST s (MutableByteArray s)
+readArrayArray (MutableArrayArray arrs#) (I# i#)
+  = ST $ \s# -> case readMutableByteArrayArray# arrs# i# s# of
+           (# s'#, mba# #) -> (# s'#, MutableByteArray mba# #)
+{-# INLINE readArrayArray #-}
+
+
+-- | Index an `ArrayArray` of `ByteArray`s.
+indexArrayArray :: ArrayArray ByteArray -> Int -> ByteArray
+indexArrayArray (ArrayArray arrs#) (I# i#)
+  = ByteArray (indexByteArrayArray# arrs# i#)
+{-# INLINE indexArrayArray #-}
+
+
+-- | Freeze a `MutableArrayArray` into a plain `ArrayArray`.
+unsafeFreezeArrayArray :: MutableArrayArray s e -> ST s (ArrayArray e)
+unsafeFreezeArrayArray (MutableArrayArray marrs#)
+  = ST $ \s# -> case unsafeFreezeArrayArray# marrs# s# of
+           (# s'#, arrs# #)  -> (# s'#, ArrayArray arrs# #)
+{-# INLINE unsafeFreezeArrayArray #-}
+
+
+-- | Freeze a nested `MutableArrayArray` into an `ArrayArray`.
+unsafeDeepFreezeArrayArray 
+        :: forall s
+        .  MutableArrayArray s (MutableByteArray s) 
+        -> ST s (ArrayArray ByteArray)
+
+unsafeDeepFreezeArrayArray marrs@(MutableArrayArray marrs#)
+ = do   let n = I# (sizeofMutableArrayArray# marrs#)
+            marrs_halfFrozen = MutableArrayArray marrs#  -- :: MutableArrayArray s (ByteArray e)
+        mapM_ (freezeSubArray marrs_halfFrozen) [0..n - 1]
+        unsafeFreezeArrayArray marrs_halfFrozen
+       
+  where
+    freezeSubArray marrs_halfFrozen i
+      = do  mba <- readArrayArray marrs i
+            ba  <- unsafeFreezeByteArray mba
+            writeArrayArray marrs_halfFrozen i ba
+{-# INLINE unsafeDeepFreezeArrayArray #-}
+
index b46ac3b..73441e5 100644 (file)
@@ -24,6 +24,7 @@ Library
         Data.Array.Parallel.Unlifted.Stream.Segments
         Data.Array.Parallel.Unlifted.Stream
         Data.Array.Parallel.Unlifted.Vectors
+        Data.Array.Parallel.Unlifted.ArrayArray
         Data.Array.Parallel.Unlifted
         
   Other-Modules:
@@ -46,5 +47,6 @@ Library
         random    == 1.0.*,
         vector    == 0.9.*,
         primitive == 0.4.*,
+        ghc-prim  == 0.2.*,
         dph-base  == 0.5.*,
         dph-prim-interface == 0.5.*