1-- |
2-- Module      : Basement.UArray.Mutable -- License     : BSD-style
3-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
4-- Stability   : experimental
5-- Portability : portable
6--
7-- A simple array abstraction that allow to use typed
8-- array of bytes where the array is pinned in memory
9-- to allow easy use with Foreign interfaces, ByteString
10-- and always aligned to 64 bytes.
11--
12{-# LANGUAGE MagicHash #-}
13{-# LANGUAGE UnboxedTuples #-}
14{-# LANGUAGE ScopedTypeVariables #-}
15module Basement.UArray.Mutable
16    ( MUArray(..)
17    -- * Property queries
18    , sizeInMutableBytesOfContent
19    , mutableLength
20    , mutableOffset
21    , mutableSame
22    , onMutableBackend
23    -- * Allocation & Copy
24    , new
25    , newPinned
26    , newNative
27    , newNative_
28    , mutableForeignMem
29    , copyAt
30    , copyFromPtr
31    , copyToPtr
32    , sub
33    -- , copyAddr
34    -- * Reading and Writing cells
35    , unsafeWrite
36    , unsafeRead
37    , write
38    , read
39    , withMutablePtr
40    , withMutablePtrHint
41    ) where
42
43import           GHC.Prim
44import           GHC.Types
45import           GHC.Ptr
46import           Basement.Compat.Base
47import           Basement.Compat.Primitive
48import           Data.Proxy
49import           Basement.Types.OffsetSize
50import           Basement.Monad
51import           Basement.PrimType
52import           Basement.FinalPtr
53import           Basement.Exception
54import qualified Basement.Block         as BLK
55import qualified Basement.Block.Mutable as MBLK
56import           Basement.Block         (MutableBlock(..))
57import           Basement.UArray.Base hiding (empty)
58import           Basement.Numerical.Subtractive
59import           Foreign.Marshal.Utils (copyBytes)
60
61sizeInMutableBytesOfContent :: forall ty s . PrimType ty => MUArray ty s -> CountOf Word8
62sizeInMutableBytesOfContent _ = primSizeInBytes (Proxy :: Proxy ty)
63{-# INLINE sizeInMutableBytesOfContent #-}
64
65-- | read a cell in a mutable array.
66--
67-- If the index is out of bounds, an error is raised.
68read :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Offset ty -> prim ty
69read array n
70    | isOutOfBound n len = primOutOfBound OOB_Read n len
71    | otherwise          = unsafeRead array n
72  where len = mutableLength array
73{-# INLINE read #-}
74
75-- | Write to a cell in a mutable array.
76--
77-- If the index is out of bounds, an error is raised.
78write :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
79write array n val
80    | isOutOfBound n len = primOutOfBound OOB_Write n len
81    | otherwise          = unsafeWrite array n val
82  where
83    len = mutableLength array
84{-# INLINE write #-}
85
86empty :: (PrimType ty, PrimMonad prim) => prim (MUArray ty (PrimState prim))
87empty = MUArray 0 0 . MUArrayMBA <$> MBLK.mutableEmpty
88
89mutableSame :: MUArray ty st -> MUArray ty st -> Bool
90mutableSame (MUArray sa ea (MUArrayMBA (MutableBlock ma))) (MUArray sb eb (MUArrayMBA (MutableBlock mb))) = (sa == sb) && (ea == eb) && bool# (sameMutableByteArray# ma mb)
91mutableSame (MUArray s1 e1 (MUArrayAddr f1)) (MUArray s2 e2 (MUArrayAddr f2)) = (s1 == s2) && (e1 == e2) && finalPtrSameMemory f1 f2
92mutableSame _ _ = False
93
94mutableForeignMem :: (PrimMonad prim, PrimType ty)
95                  => FinalPtr ty -- ^ the start pointer with a finalizer
96                  -> Int         -- ^ the number of elements (in elements, not bytes)
97                  -> prim (MUArray ty (PrimState prim))
98mutableForeignMem fptr nb = pure $ MUArray (Offset 0) (CountOf nb) (MUArrayAddr fptr)
99
100sub :: (PrimMonad prim, PrimType ty)
101    => MUArray ty (PrimState prim)
102    -> Int -- The number of elements to drop ahead
103    -> Int -- Then the number of element to retain
104    -> prim (MUArray ty (PrimState prim))
105sub (MUArray start sz back) dropElems' takeElems
106    | takeElems <= 0 = empty
107    | Just keepElems <- sz - dropElems, keepElems > 0
108                     = pure $ MUArray (start `offsetPlusE` dropElems) (min (CountOf takeElems) keepElems) back
109    | otherwise      = empty
110  where
111    dropElems = max 0 (CountOf dropElems')
112
113
114-- | return the numbers of elements in a mutable array
115mutableLength :: PrimType ty => MUArray ty st -> CountOf ty
116mutableLength (MUArray _ end _)   = end
117
118withMutablePtrHint :: forall ty prim a . (PrimMonad prim, PrimType ty)
119                   => Bool
120                   -> Bool
121                   -> MUArray ty (PrimState prim)
122                   -> (Ptr ty -> prim a)
123                   -> prim a
124withMutablePtrHint skipCopy skipCopyBack (MUArray start _ back) f =
125    case back of
126        MUArrayAddr fptr -> withFinalPtr fptr (\ptr -> f (ptr `plusPtr` os))
127        MUArrayMBA mb    -> MBLK.withMutablePtrHint skipCopy skipCopyBack mb $ \ptr -> f (ptr `plusPtr` os)
128  where
129    sz           = primSizeInBytes (Proxy :: Proxy ty)
130    !(Offset os) = offsetOfE sz start
131
132-- | Create a pointer on the beginning of the mutable array
133-- and call a function 'f'.
134--
135-- The mutable buffer can be mutated by the 'f' function
136-- and the change will be reflected in the mutable array
137--
138-- If the mutable array is unpinned, a trampoline buffer
139-- is created and the data is only copied when 'f' return.
140withMutablePtr :: (PrimMonad prim, PrimType ty)
141               => MUArray ty (PrimState prim)
142               -> (Ptr ty -> prim a)
143               -> prim a
144withMutablePtr = withMutablePtrHint False False
145
146-- | Copy from a pointer, @count@ elements, into the mutable array
147copyFromPtr :: forall prim ty . (PrimMonad prim, PrimType ty)
148            => Ptr ty -> CountOf ty -> MUArray ty (PrimState prim) -> prim ()
149copyFromPtr src@(Ptr src#) count marr
150    | count > arrSz = primOutOfBound OOB_MemCopy (sizeAsOffset count) arrSz
151    | otherwise     = onMutableBackend copyNative copyPtr marr
152  where
153    arrSz = mutableLength marr
154    ofs = mutableOffset marr
155
156    sz = primSizeInBytes (Proxy :: Proxy ty)
157    !count'@(CountOf bytes@(I# bytes#)) = sizeOfE sz count
158    !off'@(Offset od@(I# od#)) = offsetOfE sz ofs
159
160    copyNative mba = MBLK.unsafeCopyBytesPtr mba off' src count'
161    copyPtr fptr = withFinalPtr fptr $ \dst ->
162        unsafePrimFromIO $ copyBytes (dst `plusPtr` od) src bytes
163
164-- | Copy all the block content to the memory starting at the destination address
165copyToPtr :: forall ty prim . (PrimType ty, PrimMonad prim)
166          => MUArray ty (PrimState prim) -- ^ the source mutable array to copy
167          -> Ptr ty                      -- ^ The destination address where the copy is going to start
168          -> prim ()
169copyToPtr marr dst@(Ptr dst#) = onMutableBackend copyNative copyPtr marr
170  where
171    copyNative (MutableBlock mba) = primitive $ \s1 ->
172        case unsafeFreezeByteArray# mba s1 of
173            (# s2, ba #) -> (# copyByteArrayToAddr# ba os# dst# szBytes# s2, () #)
174    copyPtr fptr = unsafePrimFromIO $ withFinalPtr fptr $ \ptr ->
175        copyBytes dst (ptr `plusPtr` os) szBytes
176
177    !(Offset os@(I# os#)) = offsetInBytes $ mutableOffset marr
178    !(CountOf szBytes@(I# szBytes#)) = sizeInBytes $ mutableLength marr
179
180mutableOffset :: MUArray ty st -> Offset ty
181mutableOffset (MUArray ofs _ _) = ofs
182