dph-prim-seq: use ArrayArray for Vectors
[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 , copyArrayArray)
16 where
17 import GHC.Prim
18 import GHC.Base
19 import GHC.ST
20 import Data.Primitive.ByteArray
21
22
23 data MutableArrayArray s e
24 = MutableArrayArray (MutableArrayArray# s)
25
26 data ArrayArray e
27 = ArrayArray ArrayArray#
28
29
30 -- | Create an `ArrayArray` with the given number of elements.
31 newArrayArray :: Int -> ST s (MutableArrayArray s e)
32 newArrayArray (I# n#)
33 = ST $ \s# -> case newArrayArray# n# s# of
34 (# s'#, arr# #) -> (# s'#, MutableArrayArray arr# #)
35 {-# INLINE newArrayArray #-}
36
37
38 -- | Write a `MutableByteArray` to an `MutableArrayArray`.
39 writeArrayArrayMut :: MutableArrayArray s (MutableByteArray s) -> Int -> MutableByteArray s -> ST s ()
40 writeArrayArrayMut (MutableArrayArray arrs#) (I# i#) (MutableByteArray mba#)
41 = ST $ \s# -> case writeMutableByteArrayArray# arrs# i# mba# s# of
42 s'# -> (# s'#, () #)
43 {-# INLINE writeArrayArrayMut #-}
44
45
46 -- | Write a `ByteArray` to a `MutableArrayArray`.
47 writeArrayArray :: MutableArrayArray s ByteArray -> Int -> ByteArray -> ST s ()
48 writeArrayArray (MutableArrayArray arrs#) (I# i#) (ByteArray ba#)
49 = ST $ \s# -> case writeByteArrayArray# arrs# i# ba# s# of
50 s'# -> (# s'#, () #)
51 {-# INLINE writeArrayArray #-}
52
53
54 -- | Read a `MutableByteArray` from a `MutableArrayArray`.
55 readArrayArray :: MutableArrayArray s (MutableByteArray s) -> Int -> ST s (MutableByteArray s)
56 readArrayArray (MutableArrayArray arrs#) (I# i#)
57 = ST $ \s# -> case readMutableByteArrayArray# arrs# i# s# of
58 (# s'#, mba# #) -> (# s'#, MutableByteArray mba# #)
59 {-# INLINE readArrayArray #-}
60
61
62 -- | Index an `ArrayArray` of `ByteArray`s.
63 indexArrayArray :: ArrayArray ByteArray -> Int -> ByteArray
64 indexArrayArray (ArrayArray arrs#) (I# i#)
65 = ByteArray (indexByteArrayArray# arrs# i#)
66 {-# INLINE indexArrayArray #-}
67
68
69 -- | Freeze a `MutableArrayArray` into a plain `ArrayArray`.
70 unsafeFreezeArrayArray :: MutableArrayArray s e -> ST s (ArrayArray e)
71 unsafeFreezeArrayArray (MutableArrayArray marrs#)
72 = ST $ \s# -> case unsafeFreezeArrayArray# marrs# s# of
73 (# s'#, arrs# #) -> (# s'#, ArrayArray arrs# #)
74 {-# INLINE unsafeFreezeArrayArray #-}
75
76
77 -- | Freeze a nested `MutableArrayArray` into an `ArrayArray`.
78 unsafeDeepFreezeArrayArray
79 :: forall s
80 . MutableArrayArray s (MutableByteArray s)
81 -> ST s (ArrayArray ByteArray)
82
83 unsafeDeepFreezeArrayArray marrs@(MutableArrayArray marrs#)
84 = do let n = I# (sizeofMutableArrayArray# marrs#)
85 marrs_halfFrozen = MutableArrayArray marrs# -- :: MutableArrayArray s (ByteArray e)
86 mapM_ (freezeSubArray marrs_halfFrozen) [0..n - 1]
87 unsafeFreezeArrayArray marrs_halfFrozen
88
89 where
90 freezeSubArray marrs_halfFrozen i
91 = do mba <- readArrayArray marrs i
92 ba <- unsafeFreezeByteArray mba
93 writeArrayArray marrs_halfFrozen i ba
94 {-# INLINE unsafeDeepFreezeArrayArray #-}
95
96
97 -- | Copy an ArrayArray
98 copyArrayArray
99 :: MutableArrayArray s ByteArray -> Int
100 -> ArrayArray ByteArray -> Int
101 -> Int -> ST s ()
102
103 copyArrayArray dst startDst src startSrc len
104 = loop startDst startSrc len
105 where loop !ixDst !ixSrc !len'
106 | len' <= 0 = return ()
107 | otherwise
108 = do writeArrayArray dst ixDst $ indexArrayArray src ixSrc
109 loop (ixDst + 1) (ixSrc + 1) (len' - 1)
110