git.iko.soy/streamly
on da36a01 help sign in
streamly/core/src/Streamly/Internal/Data/Unboxed.hs
358 lines 13 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
2{-# LANGUAGE UnboxedTuples #-}
4module Streamly.Internal.Data.Unboxed
5 ( Unboxed
6 , Unbox(..)
7 , alignment
8 , peekWith
9 , pokeWith
10 , sizeOf
11 , MutableByteArray(..)
12 , castContents
13 , touch
14 , getMutableByteArray#
15 , pin
16 , unpin
17 , newUnpinnedBytes
18 , newAlignedPinnedBytes
19 , nil
20 ) where
22#include "MachDeps.h"
24#include "ArrayMacros.h"
26import Data.Complex (Complex((:+)))
27import Data.Functor ((<&>))
28import Data.Functor.Const (Const(..))
29import Data.Functor.Identity (Identity(..))
30import Foreign.Ptr (IntPtr(..), WordPtr(..))
31import Foreign.Storable (Storable(..))
32import GHC.Base (IO(..))
33import GHC.Fingerprint.Type (Fingerprint(..))
34import GHC.Int (Int16(..), Int32(..), Int64(..), Int8(..))
35import GHC.Real (Ratio(..))
36import GHC.Stable (StablePtr(..))
37import GHC.Word (Word16(..), Word32(..), Word64(..), Word8(..))
38#if MIN_VERSION_base(4,15,0)
39import GHC.RTS.Flags (IoSubSystem(..))
40#endif
41import System.IO.Unsafe (unsafePerformIO)
43import GHC.Exts
45--------------------------------------------------------------------------------
46-- The ArrayContents type
47--------------------------------------------------------------------------------
49-- XXX can use UnliftedNewtypes
50data MutableByteArray a = MutableByteArray !(MutableByteArray# RealWorld)
52{-# INLINE getMutableByteArray# #-}
53getMutableByteArray# :: MutableByteArray a -> MutableByteArray# RealWorld
54getMutableByteArray# (MutableByteArray mbarr) = mbarr
56{-# INLINE castContents #-}
57castContents :: MutableByteArray a -> MutableByteArray b
58castContents (MutableByteArray mbarr) = MutableByteArray mbarr
60{-# INLINE touch #-}
61touch :: MutableByteArray a -> IO ()
62touch (MutableByteArray contents) =
63 IO $ \s -> case touch# contents s of s' -> (# s', () #)
65--------------------------------------------------------------------------------
66-- Creation
67--------------------------------------------------------------------------------
69{-# NOINLINE nil #-}
70nil :: MutableByteArray a
71nil = unsafePerformIO $ newUnpinnedBytes 0
73{-# INLINE newUnpinnedBytes #-}
74newUnpinnedBytes :: Int -> IO (MutableByteArray a)
75newUnpinnedBytes nbytes | nbytes < 0 =
76 errorWithoutStackTrace "newUnpinnedBytes: size must be >= 0"
77newUnpinnedBytes (I# nbytes) = IO $ \s ->
78 case newByteArray# nbytes s of
79 (# s', mbarr# #) ->
80 let c = MutableByteArray mbarr#
81 in (# s', c #)
83{-# INLINE newAlignedPinnedBytes #-}
84newAlignedPinnedBytes :: Int -> Int -> IO (MutableByteArray a)
85newAlignedPinnedBytes nbytes _align | nbytes < 0 =
86 errorWithoutStackTrace "newAlignedPinnedBytes: size must be >= 0"
87newAlignedPinnedBytes (I# nbytes) (I# align) = IO $ \s ->
88 case newAlignedPinnedByteArray# nbytes align s of
89 (# s', mbarr# #) ->
90 let c = MutableByteArray mbarr#
91 in (# s', c #)
93-------------------------------------------------------------------------------
94-- Pinning & Unpinning
95-------------------------------------------------------------------------------
97{-# INLINE isPinned #-}
98isPinned :: MutableByteArray a -> Bool
99isPinned (MutableByteArray arr#) =
100 let pinnedInt = I# (isMutableByteArrayPinned# arr#)
101 in pinnedInt == 1
104{-# INLINE cloneMutableArrayWith# #-}
105cloneMutableArrayWith#
106 :: (Int# -> State# RealWorld -> (# State# RealWorld
107 , MutableByteArray# RealWorld #))
108 -> MutableByteArray# RealWorld
109 -> State# RealWorld
110 -> (# State# RealWorld, MutableByteArray# RealWorld #)
111cloneMutableArrayWith# alloc# arr# s# =
112 case getSizeofMutableByteArray# arr# s# of
113 (# s1#, i# #) ->
114 case alloc# i# s1# of
115 (# s2#, arr1# #) ->
116 case copyMutableByteArray# arr# 0# arr1# 0# i# s2# of
117 s3# -> (# s3#, arr1# #)
119{-# INLINE pin #-}
120pin :: MutableByteArray a -> IO (MutableByteArray a)
121pin arr@(MutableByteArray marr#) =
122 if isPinned arr
123 then return arr
124 else IO
125 $ \s# ->
126 case cloneMutableArrayWith# newPinnedByteArray# marr# s# of
127 (# s1#, marr1# #) -> (# s1#, MutableByteArray marr1# #)
129{-# INLINE unpin #-}
130unpin :: MutableByteArray a -> IO (MutableByteArray a)
131unpin arr@(MutableByteArray marr#) =
132 if not (isPinned arr)
133 then return arr
134 else IO
135 $ \s# ->
136 case cloneMutableArrayWith# newByteArray# marr# s# of
137 (# s1#, marr1# #) -> (# s1#, MutableByteArray marr1# #)
139--------------------------------------------------------------------------------
140-- The Unboxed type class
141--------------------------------------------------------------------------------
143-- In theory we could convert a type to and from a byte stream and use that
144-- to implement boxing, unboxing. But that would be inefficient. This type
145-- class allows each primitive type to have its own specific efficient
146-- implementation to read and write the type to memory.
147--
148-- This is a typeclass that uses MutableByteArray but could use ForeignPtr or
149-- any other representation of memory. We could make this a multiparameter type
150-- class if necessary.
151--
152-- If the type class would have to support reading and writing to a Ptr as well,
153-- basically what Storable does. We will also need to touch the anchoring ptr at
154-- the right points which is prone to errors. However, it should be simple to
155-- implement unmanaged/read-only memory arrays by allowing a Ptr type in
156-- ArrayContents, though it would require all instances to support reading from
157-- a Ptr.
158--
159-- There is a reason for using byte offset instead of element index in read and
160-- write operations in the type class. If we use element index slicing of the
161-- array becomes rigid. We can only slice the array at addresses that are
162-- aligned with the type, therefore, we cannot slice at misaligned location and
163-- then cast the slice into another type which does not ncessarily align with
164-- the original type.
165--
166-- As a side note, there seem to be no performance advantage of alignment
167-- anymore, see
168-- https://lemire.me/blog/2012/05/31/data-alignment-for-speed-myth-or-reality/
169--
170-- We have to keep the instances compatible with Storable as we depend on it for
171-- size and alignment. We can possibly write tests for checking compatibility
172-- with Storable.
174-- | An 'Unbox' type supplies operations for reading and writing the type
175-- from and to a byte array in memory. The read operation constructs the boxed
176-- type from an unboxed byte representation in memory. The write operation
177-- writes the boxed type to an unboxed byte representation in memory.
178--
179-- This type class dependes on Storable for size and alignment of the type. It
180-- is similar to Storable in the sense that it has corresponding operations for
181-- peekByteOff and pokeByteOff. While Storable allows writing to a Ptr, this
182-- type class allows writing to a MutableByteArray#.
183class Storable a => Unbox a where
184 -- | Read an element of type "a" from a MutableByteArray given the byte
185 -- index.
186 box :: MutableByteArray a -> Int -> IO a
187 -- | Write an element of type "a" to a MutableByteArray given the byte
188 -- index.
189 unbox :: MutableByteArray a -> Int -> a -> IO ()
191type Unboxed a = (Unbox a, Storable a)
193#define DERIVE_UNBOXED(_type, _constructor, _readArray, _writeArray) \
194instance Unbox _type where { \
195; {-# INLINE box #-} \
196; box (MutableByteArray mbarr) (I# n) = IO $ \s -> \
197 case _readArray mbarr n s of \
198 { (# s1, i #) -> (# s1, _constructor i #) } \
199; {-# INLINE unbox #-} \
200; unbox (MutableByteArray mbarr) (I# n) (_constructor val) = \
201 IO $ \s -> (# _writeArray mbarr n val s, () #) \
204#define DERIVE_WRAPPED_UNBOX(_constraint, _type, _constructor) \
205instance _constraint Unbox _type where \
206; {-# INLINE box #-} \
207; box arr i = _constructor <$> box (castContents arr) i \
208; {-# INLINE unbox #-} \
209; unbox arr i (_constructor a) = unbox (castContents arr) i a
211#define DERIVE_BINARY_UNBOX(_constraint, _type, _constructor, _innerType) \
212instance _constraint Unbox _type where { \
213; {-# INLINE box #-} \
214; box arr i = \
215 let contents = castContents arr :: MutableByteArray _innerType \
216 in box contents i >>= \
217 (\p1 -> box contents (i + SIZE_OF(_innerType)) \
218 <&> _constructor p1) \
219; {-# INLINE unbox #-} \
220; unbox arr i (_constructor p1 p2) = \
221 let contents = castContents arr :: MutableByteArray _innerType \
222 in unbox contents i p1 >> \
223 unbox contents (i + SIZE_OF(_innerType)) p2 \
226DERIVE_UNBOXED( Char
227 , C#
228 , readWord8ArrayAsWideChar#
229 , writeWord8ArrayAsWideChar#)
231DERIVE_UNBOXED( Int8
232 , I8#
233 , readInt8Array#
234 , writeInt8Array#)
236DERIVE_UNBOXED( Int16
237 , I16#
238 , readWord8ArrayAsInt16#
239 , writeWord8ArrayAsInt16#)
241DERIVE_UNBOXED( Int32
242 , I32#
243 , readWord8ArrayAsInt32#
244 , writeWord8ArrayAsInt32#)
246DERIVE_UNBOXED( Int
247 , I#
248 , readWord8ArrayAsInt#
249 , writeWord8ArrayAsInt#)
251DERIVE_UNBOXED( Int64
252 , I64#
253 , readWord8ArrayAsInt64#
254 , writeWord8ArrayAsInt64#)
256DERIVE_UNBOXED( Word
257 , W#
258 , readWord8ArrayAsWord#
259 , writeWord8ArrayAsWord#)
261DERIVE_UNBOXED( Word8
262 , W8#
263 , readWord8Array#
264 , writeWord8Array#)
266DERIVE_UNBOXED( Word16
267 , W16#
268 , readWord8ArrayAsWord16#
269 , writeWord8ArrayAsWord16#)
271DERIVE_UNBOXED( Word32
272 , W32#
273 , readWord8ArrayAsWord32#
274 , writeWord8ArrayAsWord32#)
276DERIVE_UNBOXED( Word64
277 , W64#
278 , readWord8ArrayAsWord64#
279 , writeWord8ArrayAsWord64#)
281DERIVE_UNBOXED( Double
282 , D#
283 , readWord8ArrayAsDouble#
284 , writeWord8ArrayAsDouble#)
286DERIVE_UNBOXED( Float
287 , F#
288 , readWord8ArrayAsFloat#
289 , writeWord8ArrayAsFloat#)
291DERIVE_UNBOXED( (StablePtr a)
292 , StablePtr
293 , readWord8ArrayAsStablePtr#
294 , writeWord8ArrayAsStablePtr#)
296DERIVE_UNBOXED( (Ptr a)
297 , Ptr
298 , readWord8ArrayAsAddr#
299 , writeWord8ArrayAsAddr#)
301DERIVE_UNBOXED( (FunPtr a)
302 , FunPtr
303 , readWord8ArrayAsAddr#
304 , writeWord8ArrayAsAddr#)
306DERIVE_WRAPPED_UNBOX(,IntPtr,IntPtr)
307DERIVE_WRAPPED_UNBOX(,WordPtr,WordPtr)
308DERIVE_WRAPPED_UNBOX(Unbox a =>,(Identity a),Identity)
309#if MIN_VERSION_base(4,14,0)
310DERIVE_WRAPPED_UNBOX(Unbox a =>,(Down a),Down)
311#endif
312DERIVE_WRAPPED_UNBOX(Unbox a =>,(Const a b),Const)
313DERIVE_BINARY_UNBOX(forall a. Unboxed a =>,(Complex a),(:+),a)
314DERIVE_BINARY_UNBOX(forall a. (Integral a, Unboxed a) =>,(Ratio a),(:%),a)
315DERIVE_BINARY_UNBOX(,Fingerprint,Fingerprint,Word64)
317instance Unbox () where
319 {-# INLINE box #-}
320 box _ _ = return ()
322 {-# INLINE unbox #-}
323 unbox _ _ _ = return ()
325#if MIN_VERSION_base(4,15,0)
326instance Unbox IoSubSystem where
328 {-# INLINE box #-}
329 box arr i = toEnum <$> box (castContents arr) i
331 {-# INLINE unbox #-}
332 unbox arr i a = unbox (castContents arr) i (fromEnum a)
333#endif
335instance Unbox Bool where
337 {-# INLINE box #-}
338 box arr i = do
339 res <- box (castContents arr) i
340 return $ res /= (0 :: Int32)
342 {-# INLINE unbox #-}
343 unbox arr i a =
344 if a
345 then unbox (castContents arr) i (1 :: Int32)
346 else unbox (castContents arr) i (0 :: Int32)
348--------------------------------------------------------------------------------
349-- Functions
350--------------------------------------------------------------------------------
352{-# INLINE peekWith #-}
353peekWith :: Unboxed a => MutableByteArray a -> Int -> IO a
354peekWith = box
356{-# INLINE pokeWith #-}
357pokeWith :: Unboxed a => MutableByteArray a -> Int -> a -> IO ()
358pokeWith = unbox