on master help sign in
Add instances of Unbox for various types
commit2f9281d0abf4e3dad27616c3ecfcd6c903b3dde2 authorAdithya Kumar <adithya@composewell.com> dateTue 4 Oct 2022, 14:29 (4 years ago) parentce8fa28
1 file changed+101−20
core/src/Streamly/Internal/Data/Unboxed.hs+101−20
core/src/Streamly/Internal/Data/Unboxed.hs+101−20
@@ -22,10 +22,20 @@ module Streamly.Internal.Data.Unboxed
2222#include "ArrayMacros.h"
2323
2424import Data.Complex (Complex((:+)))
25−import GHC.Base (IO(..))
26−import GHC.Int (Int32(..), Int64(..))
27−import GHC.Word (Word32(..), Word64(..), Word8(..))
25+import Data.Functor ((<&>))
26+import Data.Functor.Const (Const(..))
27+import Data.Functor.Identity (Identity(..))
28+import Foreign.Ptr (IntPtr(..), WordPtr(..))
2829import Foreign.Storable (Storable(..))
30+import GHC.Base (IO(..))
31+import GHC.Fingerprint.Type (Fingerprint(..))
32+import GHC.Int (Int16(..), Int32(..), Int64(..), Int8(..))
33+import GHC.Real (Ratio(..))
34+import GHC.Stable (StablePtr(..))
35+import GHC.Word (Word16(..), Word32(..), Word64(..), Word8(..))
36+#if MIN_VERSION_base(4,15,0)
37+import GHC.RTS.Flags (IoSubSystem(..))
38+#endif
2939
3040import GHC.Exts
3141
@@ -166,19 +176,51 @@ type Unboxed a = (Unbox a, Storable a)
166176#define DERIVE_UNBOXED(_type, _constructor, _readArray, _writeArray) \
167177instance Unbox _type where { \
168178; {-# INLINE box #-} \
169−; box (MutableByteArray mbarr) (I# n) = IO $ \s -> \
179+; box (MutableByteArray mbarr) (I# n) = IO $ \s -> \
170180 case _readArray mbarr n s of \
171181 { (# s1, i #) -> (# s1, _constructor i #) } \
172182; {-# INLINE unbox #-} \
173−; unbox (MutableByteArray mbarr) (I# n) (_constructor val) = \
183+; unbox (MutableByteArray mbarr) (I# n) (_constructor val) = \
174184 IO $ \s -> (# _writeArray mbarr n val s, () #) \
175185}
176186
187+#define DERIVE_WRAPPED_UNBOX(_constraint, _type, _constructor) \
188+instance _constraint Unbox _type where \
189+; {-# INLINE box #-} \
190+; box arr i = _constructor <$> box (castContents arr) i \
191+; {-# INLINE unbox #-} \
192+; unbox arr i (_constructor a) = unbox (castContents arr) i a
193+
194+#define DERIVE_BINARY_UNBOX(_constraint, _type, _constructor, _innerType) \
195+instance _constraint Unbox _type where { \
196+; {-# INLINE box #-} \
197+; box arr i = \
198+ let contents = castContents arr :: MutableByteArray _innerType \
199+ in box contents i >>= \
200+ (\p1 -> box contents (i + SIZE_OF(_innerType)) \
201+ <&> _constructor p1) \
202+; {-# INLINE unbox #-} \
203+; unbox arr i (_constructor p1 p2) = \
204+ let contents = castContents arr :: MutableByteArray _innerType \
205+ in unbox contents i p1 >> \
206+ unbox contents (i + SIZE_OF(_innerType)) p2 \
207+}
208+
177209DERIVE_UNBOXED( Char
178210 , C#
179211 , readWord8ArrayAsWideChar#
180212 , writeWord8ArrayAsWideChar#)
181213
214+DERIVE_UNBOXED( Int8
215+ , I8#
216+ , readInt8Array#
217+ , writeInt8Array#)
218+
219+DERIVE_UNBOXED( Int16
220+ , I16#
221+ , readWord8ArrayAsInt16#
222+ , writeWord8ArrayAsInt16#)
223+
182224DERIVE_UNBOXED( Int32
183225 , I32#
184226 , readWord8ArrayAsInt32#
@@ -204,6 +246,11 @@ DERIVE_UNBOXED( Word8
204246 , readWord8Array#
205247 , writeWord8Array#)
206248
249+DERIVE_UNBOXED( Word16
250+ , W16#
251+ , readWord8ArrayAsWord16#
252+ , writeWord8ArrayAsWord16#)
253+
207254DERIVE_UNBOXED( Word32
208255 , W32#
209256 , readWord8ArrayAsWord32#
@@ -219,6 +266,55 @@ DERIVE_UNBOXED( Double
219266 , readWord8ArrayAsDouble#
220267 , writeWord8ArrayAsDouble#)
221268
269+DERIVE_UNBOXED( Float
270+ , F#
271+ , readWord8ArrayAsFloat#
272+ , writeWord8ArrayAsFloat#)
273+
274+DERIVE_UNBOXED( (StablePtr a)
275+ , StablePtr
276+ , readWord8ArrayAsStablePtr#
277+ , writeWord8ArrayAsStablePtr#)
278+
279+DERIVE_UNBOXED( (Ptr a)
280+ , Ptr
281+ , readWord8ArrayAsAddr#
282+ , writeWord8ArrayAsAddr#)
283+
284+DERIVE_UNBOXED( (FunPtr a)
285+ , FunPtr
286+ , readWord8ArrayAsAddr#
287+ , writeWord8ArrayAsAddr#)
288+
289+DERIVE_WRAPPED_UNBOX(,IntPtr,IntPtr)
290+DERIVE_WRAPPED_UNBOX(,WordPtr,WordPtr)
291+DERIVE_WRAPPED_UNBOX(Unbox a =>,(Identity a),Identity)
292+#if MIN_VERSION_base(4,14,0)
293+DERIVE_WRAPPED_UNBOX(Unbox a =>,(Down a),Down)
294+#endif
295+DERIVE_WRAPPED_UNBOX(Unbox a =>,(Const a b),Const)
296+DERIVE_BINARY_UNBOX(forall a. Unboxed a =>,(Complex a),(:+),a)
297+DERIVE_BINARY_UNBOX(forall a. (Integral a, Unboxed a) =>,(Ratio a),(:%),a)
298+DERIVE_BINARY_UNBOX(,Fingerprint,Fingerprint,Word64)
299+
300+instance Unbox () where
301+
302+ {-# INLINE box #-}
303+ box _ _ = return ()
304+
305+ {-# INLINE unbox #-}
306+ unbox _ _ _ = return ()
307+
308+#if MIN_VERSION_base(4,15,0)
309+instance Unbox IoSubSystem where
310+
311+ {-# INLINE box #-}
312+ box arr i = toEnum <$> box (castContents arr) i
313+
314+ {-# INLINE unbox #-}
315+ unbox arr i a = unbox (castContents arr) i (fromEnum a)
316+#endif
317+
222318instance Unbox Bool where
223319
224320 {-# INLINE box #-}
@@ -232,21 +328,6 @@ instance Unbox Bool where
232328 then unbox (castContents arr) i (1 :: Int32)
233329 else unbox (castContents arr) i (0 :: Int32)
234330
235−instance forall a. Unboxed a => Unbox (Complex a) where
236−
237− {-# INLINE box #-}
238− box arr i = do
239− let contents = castContents arr :: MutableByteArray a
240− real <- box contents i
241− img <- box contents (i + SIZE_OF(a))
242− return $ real :+ img
243−
244− {-# INLINE unbox #-}
245− unbox arr i (real :+ img) = do
246− let contents = castContents arr :: MutableByteArray a
247− unbox contents i real
248− unbox contents (i + SIZE_OF(a)) img
249−
250331--------------------------------------------------------------------------------
251332-- Functions
252333--------------------------------------------------------------------------------