9903774deec7405bc5b2404a756da2d458340c1a
[packages/dph.git] / dph-prim-seq / Data / Array / Parallel / Unlifted / ArrayArray.hs
1
2 -- | Nested unboxed arrays.
3 --
4 -- TODO: move this into the Data.Primitive library.
5 module Data.Array.Parallel.Unlifted.ArrayArray
6 ( MutableArrayArray (..)
7 , ArrayArray(..)
8 , newArrayArray
9 , writeArrayArrayMut
10 , writeArrayArray
11 , readArrayArray
12 , indexArrayArray
13 , unsafeFreezeArrayArray
14 , unsafeDeepFreezeArrayArray)
15 where
16 import GHC.Prim
17 import GHC.Base
18 import GHC.ST
19 import Data.Primitive.ByteArray
20
21
22 data MutableArrayArray s e
23 = MutableArrayArray (MutableArrayArray# s)
24
25 data ArrayArray e
26 = ArrayArray ArrayArray#
27
28
29 -- | Create an `ArrayArray` with the given number of elements.
30 newArrayArray :: Int -> ST s (MutableArrayArray s e)
31 newArrayArray (I# n#)
32 = ST $ \s# -> case newArrayArray# n# s# of
33 (# s'#, arr# #) -> (# s'#, MutableArrayArray arr# #)
34 {-# INLINE newArrayArray #-}
35
36
37 -- | Write a `MutableByteArray` to an `MutableArrayArray`.
38 writeArrayArrayMut :: MutableArrayArray s (MutableByteArray s) -> Int -> MutableByteArray s -> ST s ()
39 writeArrayArrayMut (MutableArrayArray arrs#) (I# i#) (MutableByteArray mba#)
40 = ST $ \s# -> case writeMutableByteArrayArray# arrs# i# mba# s# of
41 s'# -> (# s'#, () #)
42 {-# INLINE writeArrayArrayMut #-}
43
44
45 -- | Write a `ByteArray` to a `MutableArrayArray`.
46 writeArrayArray :: MutableArrayArray s ByteArray -> Int -> ByteArray -> ST s ()
47 writeArrayArray (MutableArrayArray arrs#) (I# i#) (ByteArray ba#)
48 = ST $ \s# -> case writeByteArrayArray# arrs# i# ba# s# of
49 s'# -> (# s'#, () #)
50 {-# INLINE writeArrayArray #-}
51
52
53 -- | Read a `MutableByteArray` from a `MutableArrayArray`.
54 readArrayArray :: MutableArrayArray s (MutableByteArray s) -> Int -> ST s (MutableByteArray s)
55 readArrayArray (MutableArrayArray arrs#) (I# i#)
56 = ST $ \s# -> case readMutableByteArrayArray# arrs# i# s# of
57 (# s'#, mba# #) -> (# s'#, MutableByteArray mba# #)
58 {-# INLINE readArrayArray #-}
59
60
61 -- | Index an `ArrayArray` of `ByteArray`s.
62 indexArrayArray :: ArrayArray ByteArray -> Int -> ByteArray
63 indexArrayArray (ArrayArray arrs#) (I# i#)
64 = ByteArray (indexByteArrayArray# arrs# i#)
65 {-# INLINE indexArrayArray #-}
66
67
68 -- | Freeze a `MutableArrayArray` into a plain `ArrayArray`.
69 unsafeFreezeArrayArray :: MutableArrayArray s e -> ST s (ArrayArray e)
70 unsafeFreezeArrayArray (MutableArrayArray marrs#)
71 = ST $ \s# -> case unsafeFreezeArrayArray# marrs# s# of
72 (# s'#, arrs# #) -> (# s'#, ArrayArray arrs# #)
73 {-# INLINE unsafeFreezeArrayArray #-}
74
75
76 -- | Freeze a nested `MutableArrayArray` into an `ArrayArray`.
77 unsafeDeepFreezeArrayArray
78 :: forall s
79 . MutableArrayArray s (MutableByteArray s)
80 -> ST s (ArrayArray ByteArray)
81
82 unsafeDeepFreezeArrayArray marrs@(MutableArrayArray marrs#)
83 = do let n = I# (sizeofMutableArrayArray# marrs#)
84 marrs_halfFrozen = MutableArrayArray marrs# -- :: MutableArrayArray s (ByteArray e)
85 mapM_ (freezeSubArray marrs_halfFrozen) [0..n - 1]
86 unsafeFreezeArrayArray marrs_halfFrozen
87
88 where
89 freezeSubArray marrs_halfFrozen i
90 = do mba <- readArrayArray marrs i
91 ba <- unsafeFreezeByteArray mba
92 writeArrayArray marrs_halfFrozen i ba
93 {-# INLINE unsafeDeepFreezeArrayArray #-}
94