haskell / random / 369

Committed 14 Dec 2024 04:22PM UTC coverage: 69.751%. Remained the same

Build # 369

Build Type

push

github

Committed by

web-flow

Commit Message

Merge bf95d1190 into ff595fa46

Run Details

588 of 843 relevant lines covered (69.75%)

1.31 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

79.31

/src/System/Random/Array.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE UnboxedTuples #-}
-- |
-- Module      :  System.Random.Array
-- Copyright   :  (c) Alexey Kuleshevich 2024
-- License     :  BSD-style (see the file LICENSE in the 'random' repository)
-- Maintainer  :  libraries@haskell.org
--
module System.Random.Array
  ( -- * Helper array functionality
    ioToST
  , wordSizeInBits
    -- ** MutableByteArray
  , newMutableByteArray
  , newPinnedMutableByteArray
  , freezeMutableByteArray
  , writeWord8
  , writeWord64LE
  , writeByteSliceWord64LE
  , indexWord8
  , indexWord64LE
  , indexByteSliceWord64LE
  , sizeOfByteArray
  , shortByteStringToByteArray
  , byteArrayToShortByteString
  , getSizeOfMutableByteArray
  , shortByteStringToByteString
  ) where

import Control.Monad (when)
import Control.Monad.ST
import Data.Array.Byte (ByteArray(..), MutableByteArray(..))
import Data.Bits
import Data.ByteString.Short.Internal (ShortByteString(SBS))
import qualified Data.ByteString.Short.Internal as SBS (fromShort)
import Data.Word
import GHC.Exts
import GHC.IO (IO(..))
import GHC.ST (ST(..))
import GHC.Word
#if __GLASGOW_HASKELL__ >= 802
import Data.ByteString.Internal (ByteString(PS))
import GHC.ForeignPtr
#else
import Data.ByteString (ByteString)
#endif

-- Needed for WORDS_BIGENDIAN
#include "MachDeps.h"

wordSizeInBits :: Int
wordSizeInBits = finiteBitSize (0 :: Word)

-- Architecture independent helpers:

sizeOfByteArray :: ByteArray -> Int
sizeOfByteArray (ByteArray ba#) = I# (sizeofByteArray# ba#)

st_ :: (State# s -> State# s) -> ST s ()
st_ m# = ST $ \s# -> (# m# s#, () #)
{-# INLINE st_ #-}

ioToST :: IO a -> ST RealWorld a
ioToST (IO m#) = ST m#
{-# INLINE ioToST #-}

newMutableByteArray :: Int -> ST s (MutableByteArray s)
newMutableByteArray (I# n#) =
  ST $ \s# ->
    case newByteArray# n# s# of
      (# s'#, mba# #) -> (# s'#, MutableByteArray mba# #)
{-# INLINE newMutableByteArray #-}

newPinnedMutableByteArray :: Int -> ST s (MutableByteArray s)
newPinnedMutableByteArray (I# n#) =
  ST $ \s# ->
    case newPinnedByteArray# n# s# of
      (# s'#, mba# #) -> (# s'#, MutableByteArray mba# #)
{-# INLINE newPinnedMutableByteArray #-}

freezeMutableByteArray :: MutableByteArray s -> ST s ByteArray
freezeMutableByteArray (MutableByteArray mba#) =
  ST $ \s# ->
    case unsafeFreezeByteArray# mba# s# of
      (# s'#, ba# #) -> (# s'#, ByteArray ba# #)

writeWord8 :: MutableByteArray s -> Int -> Word8 -> ST s ()
writeWord8 (MutableByteArray mba#) (I# i#) (W8# w#) = st_ (writeWord8Array# mba# i# w#)
{-# INLINE writeWord8 #-}

writeByteSliceWord64LE :: MutableByteArray s -> Int -> Int -> Word64 -> ST s ()
writeByteSliceWord64LE mba fromByteIx toByteIx = go fromByteIx
  where
    go !i !z =
      when (i < toByteIx) $ do
        writeWord8 mba i (fromIntegral z :: Word8)
        go (i + 1) (z `shiftR` 8)
{-# INLINE writeByteSliceWord64LE #-}

indexWord8 ::
     ByteArray
  -> Int -- ^ Offset into immutable byte array in number of bytes
  -> Word8
indexWord8 (ByteArray ba#) (I# i#) =
  W8# (indexWord8Array# ba# i#)
{-# INLINE indexWord8 #-}

indexWord64LE ::
     ByteArray
  -> Int -- ^ Offset into immutable byte array in number of bytes
  -> Word64
#if defined WORDS_BIGENDIAN || !(__GLASGOW_HASKELL__ >= 806)
indexWord64LE ba i = indexByteSliceWord64LE ba i (i + 8)
#else
indexWord64LE (ByteArray ba#) (I# i#)
  | wordSizeInBits == 64 = W64# (indexWord8ArrayAsWord64# ba# i#)
  | otherwise =
    let !w32l = W32# (indexWord8ArrayAsWord32# ba# i#)
        !w32u = W32# (indexWord8ArrayAsWord32# ba# (i# +# 4#))
    in (fromIntegral w32u `shiftL` 32) .|. fromIntegral w32l
#endif
{-# INLINE indexWord64LE #-}

indexByteSliceWord64LE ::
     ByteArray
  -> Int -- ^ Starting offset in number of bytes
  -> Int -- ^ Ending offset in number of bytes
  -> Word64
indexByteSliceWord64LE ba fromByteIx toByteIx = goWord8 fromByteIx 0
  where
    r = (toByteIx - fromByteIx) `rem` 8
    nPadBits = if r == 0 then 0 else 8 * (8 - r)
    goWord8 i !w64
      | i < toByteIx = goWord8 (i + 1) (shiftL w64 8 .|. fromIntegral (indexWord8 ba i))
      | otherwise = byteSwap64 (shiftL w64 nPadBits)
{-# INLINE indexByteSliceWord64LE #-}

-- On big endian machines we need to write one byte at a time for consistency with little
-- endian machines. Also for GHC versions prior to 8.6 we don't have primops that can
-- write with byte offset, eg. writeWord8ArrayAsWord64# and writeWord8ArrayAsWord32#, so we
-- also must fallback to writing one byte a time. Such fallback results in about 3 times
-- slow down, which is not the end of the world.
writeWord64LE ::
     MutableByteArray s
  -> Int -- ^ Offset into mutable byte array in number of bytes
  -> Word64 -- ^ 8 bytes that will be written into the supplied array
  -> ST s ()
#if defined WORDS_BIGENDIAN || !(__GLASGOW_HASKELL__ >= 806)
writeWord64LE mba i w64 =
  writeByteSliceWord64LE mba i (i + 8) w64
#else
writeWord64LE (MutableByteArray mba#) (I# i#) w64@(W64# w64#)
  | wordSizeInBits == 64 = st_ (writeWord8ArrayAsWord64# mba# i# w64#)
  | otherwise = do
    let !(W32# w32l#) = fromIntegral w64
        !(W32# w32u#) = fromIntegral (w64 `shiftR` 32)
    st_ (writeWord8ArrayAsWord32# mba# i# w32l#)
    st_ (writeWord8ArrayAsWord32# mba# (i# +# 4#) w32u#)
#endif
{-# INLINE writeWord64LE #-}

getSizeOfMutableByteArray :: MutableByteArray s -> ST s Int
getSizeOfMutableByteArray (MutableByteArray mba#) =
#if __GLASGOW_HASKELL__ >=802
  ST $ \s ->
    case getSizeofMutableByteArray# mba# s of
      (# s', n# #) -> (# s', I# n# #)
#else
  pure $! I# (sizeofMutableByteArray# mba#)
#endif
{-# INLINE getSizeOfMutableByteArray #-}

shortByteStringToByteArray :: ShortByteString -> ByteArray
shortByteStringToByteArray (SBS ba#) = ByteArray ba#
{-# INLINE shortByteStringToByteArray #-}

byteArrayToShortByteString :: ByteArray -> ShortByteString
byteArrayToShortByteString (ByteArray ba#) = SBS ba#
{-# INLINE byteArrayToShortByteString #-}

-- | Convert a ShortByteString to ByteString by casting, whenever memory is pinned,
-- otherwise make a copy into a new pinned ByteString
shortByteStringToByteString :: ShortByteString -> ByteString
shortByteStringToByteString ba =
#if __GLASGOW_HASKELL__ < 802
  SBS.fromShort ba
#else
  let !(SBS ba#) = ba in
  if isTrue# (isByteArrayPinned# ba#)
    then pinnedByteArrayToByteString ba#
    else SBS.fromShort ba
{-# INLINE shortByteStringToByteString #-}

pinnedByteArrayToByteString :: ByteArray# -> ByteString
pinnedByteArrayToByteString ba# =
  PS (pinnedByteArrayToForeignPtr ba#) 0 (I# (sizeofByteArray# ba#))
{-# INLINE pinnedByteArrayToByteString #-}

pinnedByteArrayToForeignPtr :: ByteArray# -> ForeignPtr a
pinnedByteArrayToForeignPtr ba# =
  ForeignPtr (byteArrayContents# ba#) (PlainPtr (unsafeCoerce# ba#))
{-# INLINE pinnedByteArrayToForeignPtr #-}
#endif

1	{-# LANGUAGE BangPatterns #-}
2	{-# LANGUAGE CPP #-}
3	{-# LANGUAGE MagicHash #-}
4	{-# LANGUAGE Trustworthy #-}
5	{-# LANGUAGE UnboxedTuples #-}
6	-- \|
7	-- Module : System.Random.Array
8	-- Copyright : (c) Alexey Kuleshevich 2024
9	-- License : BSD-style (see the file LICENSE in the 'random' repository)
10	-- Maintainer : libraries@haskell.org
11	--
12	module System.Random.Array
13	( -- * Helper array functionality
14	ioToST
15	, wordSizeInBits
16	-- ** MutableByteArray
17	, newMutableByteArray
18	, newPinnedMutableByteArray
19	, freezeMutableByteArray
20	, writeWord8
21	, writeWord64LE
22	, writeByteSliceWord64LE
23	, indexWord8
24	, indexWord64LE
25	, indexByteSliceWord64LE
26	, sizeOfByteArray
27	, shortByteStringToByteArray
28	, byteArrayToShortByteString
29	, getSizeOfMutableByteArray
30	, shortByteStringToByteString
31	) where
32
33	import Control.Monad (when)
34	import Control.Monad.ST
35	import Data.Array.Byte (ByteArray(..), MutableByteArray(..))
36	import Data.Bits
37	import Data.ByteString.Short.Internal (ShortByteString(SBS))
38	import qualified Data.ByteString.Short.Internal as SBS (fromShort)
39	import Data.Word
40	import GHC.Exts
41	import GHC.IO (IO(..))
42	import GHC.ST (ST(..))
43	import GHC.Word
44	#if __GLASGOW_HASKELL__ >= 802
45	import Data.ByteString.Internal (ByteString(PS))
46	import GHC.ForeignPtr
47	#else
48	import Data.ByteString (ByteString)
49	#endif
50
51	-- Needed for WORDS_BIGENDIAN
52	#include "MachDeps.h"
53
54	wordSizeInBits :: Int
55	wordSizeInBits = finiteBitSize (0 :: Word)	1✔
56
57	-- Architecture independent helpers:
58
59	sizeOfByteArray :: ByteArray -> Int
60	sizeOfByteArray (ByteArray ba#) = I# (sizeofByteArray# ba#)	2✔
61
62	st_ :: (State# s -> State# s) -> ST s ()
63	st_ m# = ST $ \s# -> (# m# s#, () #)	1✔
64	{-# INLINE st_ #-}
65
66	ioToST :: IO a -> ST RealWorld a
67	ioToST (IO m#) = ST m#	×
68	{-# INLINE ioToST #-}
69
70	newMutableByteArray :: Int -> ST s (MutableByteArray s)
71	newMutableByteArray (I# n#) =	2✔
72	ST $ \s# ->	2✔
73	case newByteArray# n# s# of	2✔
74	(# s'#, mba# #) -> (# s'#, MutableByteArray mba# #)	2✔
75	{-# INLINE newMutableByteArray #-}
76
77	newPinnedMutableByteArray :: Int -> ST s (MutableByteArray s)
78	newPinnedMutableByteArray (I# n#) =	2✔
79	ST $ \s# ->	2✔
80	case newPinnedByteArray# n# s# of	2✔
81	(# s'#, mba# #) -> (# s'#, MutableByteArray mba# #)	2✔
82	{-# INLINE newPinnedMutableByteArray #-}
83
84	freezeMutableByteArray :: MutableByteArray s -> ST s ByteArray
85	freezeMutableByteArray (MutableByteArray mba#) =	2✔
86	ST $ \s# ->	2✔
87	case unsafeFreezeByteArray# mba# s# of	2✔
88	(# s'#, ba# #) -> (# s'#, ByteArray ba# #)	2✔
89
90	writeWord8 :: MutableByteArray s -> Int -> Word8 -> ST s ()
91	writeWord8 (MutableByteArray mba#) (I# i#) (W8# w#) = st_ (writeWord8Array# mba# i# w#)	2✔
92	{-# INLINE writeWord8 #-}
93
94	writeByteSliceWord64LE :: MutableByteArray s -> Int -> Int -> Word64 -> ST s ()
95	writeByteSliceWord64LE mba fromByteIx toByteIx = go fromByteIx	2✔
96	where
97	go !i !z =	2✔
98	when (i < toByteIx) $ do	2✔
99	writeWord8 mba i (fromIntegral z :: Word8)	2✔
100	go (i + 1) (z `shiftR` 8)	2✔
101	{-# INLINE writeByteSliceWord64LE #-}
102
103	indexWord8 ::
104	ByteArray
105	-> Int -- ^ Offset into immutable byte array in number of bytes
106	-> Word8
107	indexWord8 (ByteArray ba#) (I# i#) =	2✔
108	W8# (indexWord8Array# ba# i#)	2✔
109	{-# INLINE indexWord8 #-}
110
111	indexWord64LE ::
112	ByteArray
113	-> Int -- ^ Offset into immutable byte array in number of bytes
114	-> Word64
115	#if defined WORDS_BIGENDIAN \|\| !(__GLASGOW_HASKELL__ >= 806)
116	indexWord64LE ba i = indexByteSliceWord64LE ba i (i + 8)
117	#else
118	indexWord64LE (ByteArray ba#) (I# i#)	2✔
119	\| wordSizeInBits == 64 = W64# (indexWord8ArrayAsWord64# ba# i#)	1✔
120	\| otherwise =	×
121	let !w32l = W32# (indexWord8ArrayAsWord32# ba# i#)	×
122	!w32u = W32# (indexWord8ArrayAsWord32# ba# (i# +# 4#))	×
123	in (fromIntegral w32u `shiftL` 32) .\|. fromIntegral w32l	×
124	#endif
125	{-# INLINE indexWord64LE #-}
126
127	indexByteSliceWord64LE ::
128	ByteArray
129	-> Int -- ^ Starting offset in number of bytes
130	-> Int -- ^ Ending offset in number of bytes
131	-> Word64
132	indexByteSliceWord64LE ba fromByteIx toByteIx = goWord8 fromByteIx 0	2✔
133	where
134	r = (toByteIx - fromByteIx) `rem` 8	2✔
135	nPadBits = if r == 0 then 0 else 8 * (8 - r)	1✔
136	goWord8 i !w64	2✔
137	\| i < toByteIx = goWord8 (i + 1) (shiftL w64 8 .\|. fromIntegral (indexWord8 ba i))	2✔
138	\| otherwise = byteSwap64 (shiftL w64 nPadBits)	1✔
139	{-# INLINE indexByteSliceWord64LE #-}
140
141	-- On big endian machines we need to write one byte at a time for consistency with little
142	-- endian machines. Also for GHC versions prior to 8.6 we don't have primops that can
143	-- write with byte offset, eg. writeWord8ArrayAsWord64# and writeWord8ArrayAsWord32#, so we
144	-- also must fallback to writing one byte a time. Such fallback results in about 3 times
145	-- slow down, which is not the end of the world.
146	writeWord64LE ::
147	MutableByteArray s
148	-> Int -- ^ Offset into mutable byte array in number of bytes
149	-> Word64 -- ^ 8 bytes that will be written into the supplied array
150	-> ST s ()
151	#if defined WORDS_BIGENDIAN \|\| !(__GLASGOW_HASKELL__ >= 806)
152	writeWord64LE mba i w64 =
153	writeByteSliceWord64LE mba i (i + 8) w64
154	#else
155	writeWord64LE (MutableByteArray mba#) (I# i#) w64@(W64# w64#)	2✔
156	\| wordSizeInBits == 64 = st_ (writeWord8ArrayAsWord64# mba# i# w64#)	1✔
157	\| otherwise = do	×
158	let !(W32# w32l#) = fromIntegral w64	×
159	!(W32# w32u#) = fromIntegral (w64 `shiftR` 32)	×
160	st_ (writeWord8ArrayAsWord32# mba# i# w32l#)	×
161	st_ (writeWord8ArrayAsWord32# mba# (i# +# 4#) w32u#)	×
162	#endif
163	{-# INLINE writeWord64LE #-}
164
165	getSizeOfMutableByteArray :: MutableByteArray s -> ST s Int
166	getSizeOfMutableByteArray (MutableByteArray mba#) =	2✔
167	#if __GLASGOW_HASKELL__ >=802
168	ST $ \s ->	2✔
169	case getSizeofMutableByteArray# mba# s of	2✔
170	(# s', n# #) -> (# s', I# n# #)	2✔
171	#else
172	pure $! I# (sizeofMutableByteArray# mba#)
173	#endif
174	{-# INLINE getSizeOfMutableByteArray #-}
175
176	shortByteStringToByteArray :: ShortByteString -> ByteArray
177	shortByteStringToByteArray (SBS ba#) = ByteArray ba#	×
178	{-# INLINE shortByteStringToByteArray #-}
179
180	byteArrayToShortByteString :: ByteArray -> ShortByteString
181	byteArrayToShortByteString (ByteArray ba#) = SBS ba#	2✔
182	{-# INLINE byteArrayToShortByteString #-}
183
184	-- \| Convert a ShortByteString to ByteString by casting, whenever memory is pinned,
185	-- otherwise make a copy into a new pinned ByteString
186	shortByteStringToByteString :: ShortByteString -> ByteString
187	shortByteStringToByteString ba =	2✔
188	#if __GLASGOW_HASKELL__ < 802
189	SBS.fromShort ba
190	#else
191	let !(SBS ba#) = ba in	2✔
192	if isTrue# (isByteArrayPinned# ba#)	1✔
193	then pinnedByteArrayToByteString ba#	2✔
194	else SBS.fromShort ba	×
195	{-# INLINE shortByteStringToByteString #-}
196
197	pinnedByteArrayToByteString :: ByteArray# -> ByteString
198	pinnedByteArrayToByteString ba# =	2✔
199	PS (pinnedByteArrayToForeignPtr ba#) 0 (I# (sizeofByteArray# ba#))	2✔
200	{-# INLINE pinnedByteArrayToByteString #-}
201
202	pinnedByteArrayToForeignPtr :: ByteArray# -> ForeignPtr a
203	pinnedByteArrayToForeignPtr ba# =	2✔
204	ForeignPtr (byteArrayContents# ba#) (PlainPtr (unsafeCoerce# ba#))	1✔
205	{-# INLINE pinnedByteArrayToForeignPtr #-}
206	#endif

haskell / random / 369

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous