summaryrefslogtreecommitdiff
path: root/testsuite/tests/array/should_run/arr020.hs
blob: bb025eff03943351b6dcfaf2c45d9b41a20258c4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
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)
       }