Tests for primtypes 'ArrayArray#' and 'MutableArrayArray#'

3 files changed, 135 insertions, 0 deletions

diff --git a/testsuite/tests/array/should_run/all.T b/testsuite/tests/array/should_run/all.T
index 8b1ab2dd4d..eee3b123d2 100644
--- a/testsuite/tests/array/should_run/all.T
+++ b/testsuite/tests/array/should_run/all.T
@@ -23,3 +23,4 @@ test('arr016', reqlib('random'), compile_and_run, [''])
 test('arr017', skip_if_fast, compile_and_run, [''])
 test('arr018', skip_if_fast, compile_and_run, [''])
 test('arr019', normal, compile_and_run, [''])
+test('arr020', normal, compile_and_run, [''])
diff --git a/testsuite/tests/array/should_run/arr020.hs b/testsuite/tests/array/should_run/arr020.hs
new file mode 100644
index 0000000000..bb025eff03
--- /dev/null
+++ b/testsuite/tests/array/should_run/arr020.hs
@@ -0,0 +1,132 @@
+{-# LANGUAGE MagicHash, UnboxedTuples, ScopedTypeVariables #-}
+
+module Main where
+
+import GHC.Prim
+import GHC.Base
+import GHC.ST
+import GHC.Word
+import Control.Monad
+import System.Mem
+
+data MutableByteArray s = MutableByteArray (MutableByteArray# s)
+
+data ByteArray e = ByteArray ByteArray#
+
+newByteArray :: Int -> ST s (MutableByteArray s)
+newByteArray (I# n#)
+  = ST $ \s# -> case newByteArray# n# s# of
+           (# s'#, arr# #) -> (# s'#, MutableByteArray arr# #)
+
+writeByteArray :: MutableByteArray s -> Int -> Word32 -> ST s ()
+writeByteArray (MutableByteArray mba#) (I# i#) (W32# w#)
+  = ST $ \s# -> case writeWord32Array# mba# i# w# s# of
+           s'# -> (# s'#, () #)
+
+indexArray :: ByteArray Word32 -> Int -> Word32
+indexArray (ByteArray arr#) (I# i#)
+  = W32# (indexWord32Array# arr# i#)
+
+unsafeFreezeByteArray :: MutableByteArray s -> ST s (ByteArray e)
+unsafeFreezeByteArray (MutableByteArray mba#)
+  = ST $ \s# -> case unsafeFreezeByteArray# mba# s# of
+           (# s'#, ba# #)  -> (# s'#, ByteArray ba# #)
+
+data MutableArrayArray s e = MutableArrayArray (MutableArrayArray# s)
+
+data ArrayArray e = ArrayArray ArrayArray#
+
+newArrayArray :: Int -> ST s (MutableArrayArray s e)
+newArrayArray (I# n#)
+  = ST $ \s# -> case newArrayArray# n# s# of
+           (# s'#, arr# #) -> (# s'#, MutableArrayArray arr# #)
+
+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'#, () #)
+
+writeArrayArray :: MutableArrayArray s (ByteArray s) -> Int -> ByteArray s
+                -> ST s ()
+writeArrayArray (MutableArrayArray arrs#) (I# i#) (ByteArray ba#)
+  = ST $ \s# -> case writeByteArrayArray# arrs# i# ba# s# of
+           s'# -> (# s'#, () #)
+
+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# #)
+
+indexArrayArray :: ArrayArray (ByteArray e) -> Int -> ByteArray e
+indexArrayArray (ArrayArray arrs#) (I# i#)
+  = ByteArray (indexByteArrayArray# arrs# i#)
+
+unsafeFreezeArrayArray :: MutableArrayArray s e -> ST s (ArrayArray e)
+unsafeFreezeArrayArray (MutableArrayArray marrs#)
+  = ST $ \s# -> case unsafeFreezeArrayArray# marrs# s# of
+           (# s'#, arrs# #)  -> (# s'#, ArrayArray arrs# #)
+
+unsafeDeepFreezeArrayArray :: forall s e
+                           .  MutableArrayArray s (MutableByteArray s) 
+                           -> ST s (ArrayArray (ByteArray e))
+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
+           }
+
+newByteArrays :: [Int] -> ST s (MutableArrayArray s (MutableByteArray s))
+newByteArrays ns
+  = do { arrs <- newArrayArray (length ns)
+       ; zipWithM_ (writeNewByteArray arrs) ns [0..]
+       ; return arrs
+       }
+  where
+    writeNewByteArray arrs n i
+      = do { mba <- newByteArray (n * 4)    -- we store 32-bit words
+           ; writeArrayArrayMut arrs i mba
+           }
+
+type UnboxedArray2D e = ArrayArray (ByteArray e)
+
+newUnboxedArray2D :: [[Word32]] -> UnboxedArray2D Word32
+newUnboxedArray2D values
+  = runST $
+    do { marrs <- newByteArrays (map length values)
+       ; zipWithM_ (fill marrs) values [0..]
+       ; arrs <- unsafeDeepFreezeArrayArray marrs
+       ; return arrs
+       }
+  where
+    fill marrs vs i
+      = do { mba <- readArrayArray marrs i
+           ; zipWithM_ (writeByteArray mba) [0..] vs
+           }
+
+unboxedArray2D :: UnboxedArray2D Word32
+unboxedArray2D 
+  = newUnboxedArray2D
+    [ [1..10]
+    , [11..200]
+    , []
+    , [1..1000] ++ [42] ++ [1001..2000]
+    , [1..100000]
+    ]
+
+indexUnboxedArray2D :: UnboxedArray2D Word32 -> (Int, Int) -> Word32
+indexUnboxedArray2D arr (i, j)
+  = indexArrayArray arr i `indexArray` j
+
+main 
+  = do { print $ unboxedArray2D `indexUnboxedArray2D` (3, 1000)
+       ; performGC
+       ; print $ unboxedArray2D `indexUnboxedArray2D` (3, 1000)
+       }
diff --git a/testsuite/tests/array/should_run/arr020.stdout b/testsuite/tests/array/should_run/arr020.stdout
new file mode 100644
index 0000000000..daaac9e303
--- /dev/null
+++ b/testsuite/tests/array/should_run/arr020.stdout
@@ -0,0 +1,2 @@
+42
+42