git.iko.soy/streamly
on dc58c59 help sign in
streamly/core/src/Streamly/Internal/Data/MutByteArray.hs
249 lines 8.1 kB history raw
1{-# LANGUAGE UnboxedTuples #-}
3-- |
4-- Module : Streamly.Internal.Data.MutByteArray
5-- Copyright : (c) 2023 Composewell Technologies
6-- License : BSD3-3-Clause
7-- Maintainer : streamly@composewell.com
8-- Stability : experimental
9-- Portability : GHC
10--
11module Streamly.Internal.Data.MutByteArray
12 (
13 -- ** MutByteArray
14 MutByteArray(..)
15 , MutableByteArray
16 , getMutableByteArray#
18 -- ** Pinning
19 , PinnedState(..)
20 , isPinned
21 , pin
22 , unpin
24 -- ** Allocation
25 , nil
26 , newBytesAs
27 , newByteArray
28 , pinnedNewByteArray
29 , pinnedNewAlignedBytes
31 -- ** Access
32 , sizeOfMutableByteArray
33 , putSliceUnsafe
34 , asPtrUnsafe
35 ) where
37import Control.Monad.IO.Class (MonadIO(..))
38#ifdef DEBUG
39import Debug.Trace (trace)
40#endif
41import GHC.Base (IO(..))
42import System.IO.Unsafe (unsafePerformIO)
44import GHC.Exts
46--------------------------------------------------------------------------------
47-- The ArrayContents type
48--------------------------------------------------------------------------------
50data PinnedState
51 = Pinned
52 | Unpinned
54-- XXX can use UnliftedNewtypes
56-- | A lifted mutable byte array type wrapping @MutableByteArray# RealWorld@.
57-- This is a low level array used to back high level unboxed arrays and
58-- serialized data.
59data MutByteArray = MutByteArray (MutableByteArray# RealWorld)
61{-# DEPRECATED MutableByteArray "Please use MutByteArray instead" #-}
62type MutableByteArray = MutByteArray
64{-# INLINE getMutableByteArray# #-}
65getMutableByteArray# :: MutByteArray -> MutableByteArray# RealWorld
66getMutableByteArray# (MutByteArray mbarr) = mbarr
68-- | Return the size of the array in bytes.
69{-# INLINE sizeOfMutableByteArray #-}
70sizeOfMutableByteArray :: MutByteArray -> IO Int
71sizeOfMutableByteArray (MutByteArray arr) =
72 IO $ \s ->
73 case getSizeofMutableByteArray# arr s of
74 (# s1, i #) -> (# s1, I# i #)
76{-# INLINE touch #-}
77touch :: MutByteArray -> IO ()
78touch (MutByteArray contents) =
79 IO $ \s -> case touch# contents s of s' -> (# s', () #)
81-- XXX We can provide another API for "unsafe" FFI calls passing an unlifted
82-- pointer to the FFI call. For unsafe calls we do not need to pin the array.
83-- We can pass an unlifted pointer to the FFI routine to avoid GC kicking in
84-- before the pointer is wrapped.
85--
86-- From the GHC manual:
87--
88-- GHC, since version 8.4, guarantees that garbage collection will never occur
89-- during an unsafe call, even in the bytecode interpreter, and further
90-- guarantees that unsafe calls will be performed in the calling thread. Making
91-- it safe to pass heap-allocated objects to unsafe functions.
93-- | Use a @MutByteArray@ as @Ptr a@. This is useful when we want to pass
94-- an array as a pointer to some operating system call or to a "safe" FFI call.
95--
96-- If the array is not pinned it is copied to pinned memory before passing it
97-- to the monadic action.
98--
99-- /Performance Notes:/ Forces a copy if the array is not pinned. It is advised
100-- that the programmer keeps this in mind and creates a pinned array
101-- opportunistically before this operation occurs, to avoid the cost of a copy
102-- if possible.
103--
104-- /Unsafe/ because of direct pointer operations. The user must ensure that
105-- they are writing within the legal bounds of the array.
106--
107-- /Pre-release/
108--
109{-# INLINE asPtrUnsafe #-}
110asPtrUnsafe :: MonadIO m => MutByteArray -> (Ptr a -> m b) -> m b
111asPtrUnsafe arr f = do
112 contents <- liftIO $ pin arr
113 let !ptr = Ptr (byteArrayContents#
114 (unsafeCoerce# (getMutableByteArray# contents)))
115 r <- f ptr
116 liftIO $ touch contents
117 return r
119--------------------------------------------------------------------------------
120-- Creation
121--------------------------------------------------------------------------------
123{-# NOINLINE nil #-}
124nil :: MutByteArray
125nil = unsafePerformIO $ newByteArray 0
127{-# INLINE newByteArray #-}
128newByteArray :: Int -> IO MutByteArray
129newByteArray nbytes | nbytes < 0 =
130 errorWithoutStackTrace "newByteArray: size must be >= 0"
131newByteArray (I# nbytes) = IO $ \s ->
132 case newByteArray# nbytes s of
133 (# s', mbarr# #) ->
134 let c = MutByteArray mbarr#
135 in (# s', c #)
137{-# INLINE pinnedNewByteArray #-}
138pinnedNewByteArray :: Int -> IO MutByteArray
139pinnedNewByteArray nbytes | nbytes < 0 =
140 errorWithoutStackTrace "pinnedNewByteArray: size must be >= 0"
141pinnedNewByteArray (I# nbytes) = IO $ \s ->
142 case newPinnedByteArray# nbytes s of
143 (# s', mbarr# #) ->
144 let c = MutByteArray mbarr#
145 in (# s', c #)
147{-# INLINE pinnedNewAlignedBytes #-}
148pinnedNewAlignedBytes :: Int -> Int -> IO MutByteArray
149pinnedNewAlignedBytes nbytes _align | nbytes < 0 =
150 errorWithoutStackTrace "pinnedNewAlignedBytes: size must be >= 0"
151pinnedNewAlignedBytes (I# nbytes) (I# align) = IO $ \s ->
152 case newAlignedPinnedByteArray# nbytes align s of
153 (# s', mbarr# #) ->
154 let c = MutByteArray mbarr#
155 in (# s', c #)
157{-# INLINE newBytesAs #-}
158newBytesAs :: PinnedState -> Int -> IO MutByteArray
159newBytesAs Unpinned = newByteArray
160newBytesAs Pinned = pinnedNewByteArray
162-------------------------------------------------------------------------------
163-- Copying
164-------------------------------------------------------------------------------
166-- | Put a sub range of a source array into a subrange of a destination array.
167-- This is not safe as it does not check the bounds.
168{-# INLINE putSliceUnsafe #-}
169putSliceUnsafe ::
170 MonadIO m
171 => MutByteArray
172 -> Int
173 -> MutByteArray
174 -> Int
175 -> Int
176 -> m ()
177putSliceUnsafe src srcStartBytes dst dstStartBytes lenBytes = liftIO $ do
178#ifdef DEBUG
179 srcLen <- sizeOfMutableByteArray src
180 dstLen <- sizeOfMutableByteArray src
181 when (srcLen - srcStartBytes < lenBytes) $ error "SRC: Insufficient length."
182 when (dstLen - dstStartBytes < lenBytes) $ error "DST: Insufficient length."
183#endif
184 let !(I# srcStartBytes#) = srcStartBytes
185 !(I# dstStartBytes#) = dstStartBytes
186 !(I# lenBytes#) = lenBytes
187 let arrS# = getMutableByteArray# src
188 arrD# = getMutableByteArray# dst
189 IO $ \s# -> (# copyMutableByteArray#
190 arrS# srcStartBytes# arrD# dstStartBytes# lenBytes# s#
191 , () #)
193-------------------------------------------------------------------------------
194-- Pinning & Unpinning
195-------------------------------------------------------------------------------
197{-# INLINE isPinned #-}
198isPinned :: MutByteArray -> Bool
199isPinned (MutByteArray arr#) =
200 let pinnedInt = I# (isMutableByteArrayPinned# arr#)
201 in pinnedInt /= 0
204{-# INLINE cloneMutableArrayWith# #-}
205cloneMutableArrayWith#
206 :: (Int# -> State# RealWorld -> (# State# RealWorld
207 , MutableByteArray# RealWorld #))
208 -> MutableByteArray# RealWorld
209 -> State# RealWorld
210 -> (# State# RealWorld, MutableByteArray# RealWorld #)
211cloneMutableArrayWith# alloc# arr# s# =
212 case getSizeofMutableByteArray# arr# s# of
213 (# s1#, i# #) ->
214 case alloc# i# s1# of
215 (# s2#, arr1# #) ->
216 case copyMutableByteArray# arr# 0# arr1# 0# i# s2# of
217 s3# -> (# s3#, arr1# #)
219{-# INLINE pin #-}
220pin :: MutByteArray -> IO MutByteArray
221pin arr@(MutByteArray marr#) =
222 if isPinned arr
223 then return arr
224 else
225#ifdef DEBUG
226 do
227 -- XXX dump stack trace
228 trace ("pin: Copying array") (return ())
229#endif
230 IO
231 $ \s# ->
232 case cloneMutableArrayWith# newPinnedByteArray# marr# s# of
233 (# s1#, marr1# #) -> (# s1#, MutByteArray marr1# #)
235{-# INLINE unpin #-}
236unpin :: MutByteArray -> IO MutByteArray
237unpin arr@(MutByteArray marr#) =
238 if not (isPinned arr)
239 then return arr
240 else
241#ifdef DEBUG
242 do
243 -- XXX dump stack trace
244 trace ("unpin: Copying array") (return ())
245#endif
246 IO
247 $ \s# ->
248 case cloneMutableArrayWith# newByteArray# marr# s# of
249 (# s1#, marr1# #) -> (# s1#, MutByteArray marr1# #)