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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
|
-- |
-- Module : Basement.UArray.Mutable -- License : BSD-style
-- Maintainer : Vincent Hanquez <vincent@snarc.org>
-- Stability : experimental
-- Portability : portable
--
-- A simple array abstraction that allow to use typed
-- array of bytes where the array is pinned in memory
-- to allow easy use with Foreign interfaces, ByteString
-- and always aligned to 64 bytes.
--
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Basement.UArray.Mutable
( MUArray(..)
-- * Property queries
, sizeInMutableBytesOfContent
, mutableLength
, mutableOffset
, mutableSame
, onMutableBackend
-- * Allocation & Copy
, new
, newPinned
, newNative
, newNative_
, mutableForeignMem
, copyAt
, copyFromPtr
, copyToPtr
, sub
-- , copyAddr
-- * Reading and Writing cells
, unsafeWrite
, unsafeRead
, write
, read
, withMutablePtr
, withMutablePtrHint
) where
import GHC.Prim
import GHC.Exts
import GHC.Types
import GHC.Ptr
import Basement.Compat.Base
import Basement.Compat.Primitive
import Data.Proxy
import Basement.Types.OffsetSize
import Basement.Monad
import Basement.PrimType
import Basement.FinalPtr
import Basement.Exception
import qualified Basement.Block as BLK
import qualified Basement.Block.Mutable as MBLK
import Basement.Block (MutableBlock(..))
import Basement.UArray.Base hiding (empty)
import Basement.Numerical.Subtractive
import Foreign.Marshal.Utils (copyBytes)
sizeInMutableBytesOfContent :: forall ty s . PrimType ty => MUArray ty s -> CountOf Word8
sizeInMutableBytesOfContent _ = primSizeInBytes (Proxy :: Proxy ty)
{-# INLINE sizeInMutableBytesOfContent #-}
-- | read a cell in a mutable array.
--
-- If the index is out of bounds, an error is raised.
read :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Offset ty -> prim ty
read array n
| isOutOfBound n len = primOutOfBound OOB_Read n len
| otherwise = unsafeRead array n
where len = mutableLength array
{-# INLINE read #-}
-- | Write to a cell in a mutable array.
--
-- If the index is out of bounds, an error is raised.
write :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
write array n val
| isOutOfBound n len = primOutOfBound OOB_Write n len
| otherwise = unsafeWrite array n val
where
len = mutableLength array
{-# INLINE write #-}
empty :: (PrimType ty, PrimMonad prim) => prim (MUArray ty (PrimState prim))
empty = MUArray 0 0 . MUArrayMBA <$> MBLK.mutableEmpty
mutableSame :: MUArray ty st -> MUArray ty st -> Bool
mutableSame (MUArray sa ea (MUArrayMBA (MutableBlock ma))) (MUArray sb eb (MUArrayMBA (MutableBlock mb))) = (sa == sb) && (ea == eb) && bool# (sameMutableByteArray# ma mb)
mutableSame (MUArray s1 e1 (MUArrayAddr f1)) (MUArray s2 e2 (MUArrayAddr f2)) = (s1 == s2) && (e1 == e2) && finalPtrSameMemory f1 f2
mutableSame _ _ = False
mutableForeignMem :: (PrimMonad prim, PrimType ty)
=> FinalPtr ty -- ^ the start pointer with a finalizer
-> Int -- ^ the number of elements (in elements, not bytes)
-> prim (MUArray ty (PrimState prim))
mutableForeignMem fptr nb = pure $ MUArray (Offset 0) (CountOf nb) (MUArrayAddr fptr)
sub :: (PrimMonad prim, PrimType ty)
=> MUArray ty (PrimState prim)
-> Int -- The number of elements to drop ahead
-> Int -- Then the number of element to retain
-> prim (MUArray ty (PrimState prim))
sub (MUArray start sz back) dropElems' takeElems
| takeElems <= 0 = empty
| Just keepElems <- sz - dropElems, keepElems > 0
= pure $ MUArray (start `offsetPlusE` dropElems) (min (CountOf takeElems) keepElems) back
| otherwise = empty
where
dropElems = max 0 (CountOf dropElems')
-- | return the numbers of elements in a mutable array
mutableLength :: PrimType ty => MUArray ty st -> CountOf ty
mutableLength (MUArray _ end _) = end
withMutablePtrHint :: forall ty prim a . (PrimMonad prim, PrimType ty)
=> Bool
-> Bool
-> MUArray ty (PrimState prim)
-> (Ptr ty -> prim a)
-> prim a
withMutablePtrHint skipCopy skipCopyBack (MUArray start _ back) f =
case back of
MUArrayAddr fptr -> withFinalPtr fptr (\ptr -> f (ptr `plusPtr` os))
MUArrayMBA mb -> MBLK.withMutablePtrHint skipCopy skipCopyBack mb $ \ptr -> f (ptr `plusPtr` os)
where
sz = primSizeInBytes (Proxy :: Proxy ty)
!(Offset os) = offsetOfE sz start
-- | Create a pointer on the beginning of the mutable array
-- and call a function 'f'.
--
-- The mutable buffer can be mutated by the 'f' function
-- and the change will be reflected in the mutable array
--
-- If the mutable array is unpinned, a trampoline buffer
-- is created and the data is only copied when 'f' return.
withMutablePtr :: (PrimMonad prim, PrimType ty)
=> MUArray ty (PrimState prim)
-> (Ptr ty -> prim a)
-> prim a
withMutablePtr = withMutablePtrHint False False
-- | Copy from a pointer, @count@ elements, into the mutable array
copyFromPtr :: forall prim ty . (PrimMonad prim, PrimType ty)
=> Ptr ty -> CountOf ty -> MUArray ty (PrimState prim) -> prim ()
copyFromPtr src@(Ptr src#) count marr
| count > arrSz = primOutOfBound OOB_MemCopy (sizeAsOffset count) arrSz
| otherwise = onMutableBackend copyNative copyPtr marr
where
arrSz = mutableLength marr
ofs = mutableOffset marr
sz = primSizeInBytes (Proxy :: Proxy ty)
!count'@(CountOf bytes@(I# bytes#)) = sizeOfE sz count
!off'@(Offset od@(I# od#)) = offsetOfE sz ofs
copyNative mba = MBLK.unsafeCopyBytesPtr mba off' src count'
copyPtr fptr = withFinalPtr fptr $ \dst ->
unsafePrimFromIO $ copyBytes (dst `plusPtr` od) src bytes
-- | Copy all the block content to the memory starting at the destination address
copyToPtr :: forall ty prim . (PrimType ty, PrimMonad prim)
=> MUArray ty (PrimState prim) -- ^ the source mutable array to copy
-> Ptr ty -- ^ The destination address where the copy is going to start
-> prim ()
copyToPtr marr dst@(Ptr dst#) = onMutableBackend copyNative copyPtr marr
where
copyNative (MutableBlock mba) = primitive $ \s1 ->
case unsafeFreezeByteArray# mba s1 of
(# s2, ba #) -> (# copyByteArrayToAddr# ba os# dst# szBytes# s2, () #)
copyPtr fptr = unsafePrimFromIO $ withFinalPtr fptr $ \ptr ->
copyBytes dst (ptr `plusPtr` os) szBytes
!(Offset os@(I# os#)) = offsetInBytes $ mutableOffset marr
!(CountOf szBytes@(I# szBytes#)) = sizeInBytes $ mutableLength marr
mutableOffset :: MUArray ty st -> Offset ty
mutableOffset (MUArray ofs _ _) = ofs
|
