content: {-# LANGUAGE MagicHash, UnboxedTuples, ScopedTypeVariables #-}
  
  module Main where
  
  import GHC.Exts
  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)
         }
  