haskell / random / 434

Committed 05 Jul 2025 08:47PM UTC coverage: 69.035% (+0.3%) from 68.696%

Build # 434

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #189 from haskell/lehins/fourmolize

fourmolize

Run Details

141 of 209 new or added lines in 6 files covered. (67.46%)

98 existing lines in 6 files now uncovered.

651 of 943 relevant lines covered (69.03%)

1.3 hits per line

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

69.83

/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,

  -- ** MutableArray
  Array (..),
  MutableArray (..),
  newMutableArray,
  freezeMutableArray,
  writeArray,
  shuffleListM,
  shuffleListST,
) where

import Control.Monad (when)
import Control.Monad.ST
import Control.Monad.Trans (MonadTrans, lift)
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)

----------------
-- Byte Array --
----------------

-- 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 ->
  -- | Offset into immutable byte array in number of bytes
  Int ->
  Word8
indexWord8 (ByteArray ba#) (I# i#) =
  W8# (indexWord8Array# ba# i#)
{-# INLINE indexWord8 #-}

indexWord64LE ::
  ByteArray ->
  -- | Offset into immutable byte array in number of bytes
  Int ->
  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 ->
  -- | Starting offset in number of bytes
  Int ->
  -- | Ending offset in number of bytes
  Int ->
  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 ->
  -- | Offset into mutable byte array in number of bytes
  Int ->
  -- | 8 bytes that will be written into the supplied array
  Word64 ->
  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
#if __GLASGOW_HASKELL__ >=802
getSizeOfMutableByteArray (MutableByteArray mba#) =
  ST $ \s ->
    case getSizeofMutableByteArray# mba# s of
      (# s', n# #) -> (# s', I# n# #)
#else
getSizeOfMutableByteArray (MutableByteArray mba#) =
  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
#if __GLASGOW_HASKELL__ < 802
shortByteStringToByteString ba = SBS.fromShort ba
#else
shortByteStringToByteString ba =
  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

-----------------
-- Boxed Array --
-----------------

data Array a = Array (Array# a)

data MutableArray s a = MutableArray (MutableArray# s a)

newMutableArray :: Int -> a -> ST s (MutableArray s a)
newMutableArray (I# n#) a =
  ST $ \s# ->
    case newArray# n# a s# of
      (# s'#, ma# #) -> (# s'#, MutableArray ma# #)
{-# INLINE newMutableArray #-}

freezeMutableArray :: MutableArray s a -> ST s (Array a)
freezeMutableArray (MutableArray ma#) =
  ST $ \s# ->
    case unsafeFreezeArray# ma# s# of
      (# s'#, a# #) -> (# s'#, Array a# #)
{-# INLINE freezeMutableArray #-}

sizeOfMutableArray :: MutableArray s a -> Int
sizeOfMutableArray (MutableArray ma#) = I# (sizeofMutableArray# ma#)
{-# INLINE sizeOfMutableArray #-}

readArray :: MutableArray s a -> Int -> ST s a
readArray (MutableArray ma#) (I# i#) = ST (readArray# ma# i#)
{-# INLINE readArray #-}

writeArray :: MutableArray s a -> Int -> a -> ST s ()
writeArray (MutableArray ma#) (I# i#) a = st_ (writeArray# ma# i# a)
{-# INLINE writeArray #-}

swapArray :: MutableArray s a -> Int -> Int -> ST s ()
swapArray ma i j = do
  x <- readArray ma i
  y <- readArray ma j
  writeArray ma j x
  writeArray ma i y
{-# INLINE swapArray #-}

-- | Write contents of the list into the mutable array. Make sure that array is big
-- enough or segfault will happen.
fillMutableArrayFromList :: MutableArray s a -> [a] -> ST s ()
fillMutableArrayFromList ma = go 0
  where
    go _ [] = pure ()
    go i (x : xs) = writeArray ma i x >> go (i + 1) xs
{-# INLINE fillMutableArrayFromList #-}

readListFromMutableArray :: MutableArray s a -> ST s [a]
readListFromMutableArray ma = go (len - 1) []
  where
    len = sizeOfMutableArray ma
    go i !acc
      | i >= 0 = do
          x <- readArray ma i
          go (i - 1) (x : acc)
      | otherwise = pure acc
{-# INLINE readListFromMutableArray #-}

-- | Generate a list of indices that will be used for swapping elements in uniform shuffling:
--
-- @
-- [ (0, n - 1)
-- , (0, n - 2)
-- , (0, n - 3)
-- , ...
-- , (0, 3)
-- , (0, 2)
-- , (0, 1)
-- ]
-- @
genSwapIndices ::
  Monad m =>
  -- | Action that generates a Word in the supplied range.
  (Word -> m Word) ->
  -- | Number of index swaps to generate.
  Word ->
  m [Int]
genSwapIndices genWordR n = go 1 []
  where
    go i !acc
      | i >= n = pure acc
      | otherwise = do
          x <- genWordR i
          let !xi = fromIntegral x
          go (i + 1) (xi : acc)
{-# INLINE genSwapIndices #-}

-- | Implementation of mutable version of Fisher-Yates shuffle. Unfortunately, we cannot generally
-- interleave pseudo-random number generation and mutation of `ST` monad, therefore we have to
-- pre-generate all of the index swaps with `genSwapIndices` and store them in a list before we can
-- perform the actual swaps.
shuffleListM :: Monad m => (Word -> m Word) -> [a] -> m [a]
shuffleListM genWordR ls
  | len <= 1 = pure ls
  | otherwise = do
      swapIxs <- genSwapIndices genWordR (fromIntegral len)
      pure $ runST $ do
        ma <- newMutableArray len $ error "Impossible: shuffleListM"
        fillMutableArrayFromList ma ls

        -- Shuffle elements of the mutable array according to the uniformly generated index swap list
        let goSwap _ [] = pure ()
            goSwap i (j : js) = swapArray ma i j >> goSwap (i - 1) js
        goSwap (len - 1) swapIxs

        readListFromMutableArray ma
  where
    len = length ls
{-# INLINE shuffleListM #-}

-- | This is a ~x2-x3 more efficient version of `shuffleListM`. It is more efficient because it does
-- not need to pregenerate a list of indices and instead generates them on demand. Because of this the
-- result that will be produced will differ for the same generator, since the order in which index
-- swaps are generated is reversed.
--
-- Unfortunately, most stateful generator monads can't handle `MonadTrans`, so this version is only
-- used for implementing the pure shuffle.
shuffleListST :: (Monad (t (ST s)), MonadTrans t) => (Word -> t (ST s) Word) -> [a] -> t (ST s) [a]
shuffleListST genWordR ls
  | len <= 1 = pure ls
  | otherwise = do
      ma <- lift $ newMutableArray len $ error "Impossible: shuffleListST"
      lift $ fillMutableArrayFromList ma ls

      -- Shuffle elements of the mutable array according to the uniformly generated index swap
      let goSwap i =
            when (i > 0) $ do
              j <- genWordR $ (fromIntegral :: Int -> Word) i
              lift $ swapArray ma i ((fromIntegral :: Word -> Int) j)
              goSwap (i - 1)
      goSwap (len - 1)

      lift $ readListFromMutableArray ma
  where
    len = length ls
{-# INLINE shuffleListST #-}

1	{-# LANGUAGE BangPatterns #-}
2	{-# LANGUAGE CPP #-}
3	{-# LANGUAGE MagicHash #-}
4	{-# LANGUAGE Trustworthy #-}
5	{-# LANGUAGE UnboxedTuples #-}
6
7	-- \|
8	-- Module : System.Random.Array
9	-- Copyright : (c) Alexey Kuleshevich 2024
10	-- License : BSD-style (see the file LICENSE in the 'random' repository)
11	-- Maintainer : libraries@haskell.org
12	module System.Random.Array (
13	-- * Helper array functionality
14	ioToST,
15	wordSizeInBits,
16
17	-- ** MutableByteArray
18	newMutableByteArray,
19	newPinnedMutableByteArray,
20	freezeMutableByteArray,
21	writeWord8,
22	writeWord64LE,
23	writeByteSliceWord64LE,
24	indexWord8,
25	indexWord64LE,
26	indexByteSliceWord64LE,
27	sizeOfByteArray,
28	shortByteStringToByteArray,
29	byteArrayToShortByteString,
30	getSizeOfMutableByteArray,
31	shortByteStringToByteString,
32
33	-- ** MutableArray
34	Array (..),
35	MutableArray (..),
36	newMutableArray,
37	freezeMutableArray,
38	writeArray,
39	shuffleListM,
40	shuffleListST,
41	) where
42
43	import Control.Monad (when)
44	import Control.Monad.ST
45	import Control.Monad.Trans (MonadTrans, lift)
46	import Data.Array.Byte (ByteArray (..), MutableByteArray (..))
47	import Data.Bits
48	import Data.ByteString.Short.Internal (ShortByteString (SBS))
49	import qualified Data.ByteString.Short.Internal as SBS (fromShort)
50	import Data.Word
51	import GHC.Exts
52	import GHC.IO (IO (..))
53	import GHC.ST (ST (..))
54	import GHC.Word
55	#if __GLASGOW_HASKELL__ >= 802
56	import Data.ByteString.Internal (ByteString(PS))
57	import GHC.ForeignPtr
58	#else
59	import Data.ByteString (ByteString)
60	#endif
61
62	-- Needed for WORDS_BIGENDIAN
63	#include "MachDeps.h"
64
65	wordSizeInBits :: Int
66	wordSizeInBits = finiteBitSize (0 :: Word)	1✔
67
68	----------------
69	-- Byte Array --
70	----------------
71
72	-- Architecture independent helpers:
73
74	sizeOfByteArray :: ByteArray -> Int
75	sizeOfByteArray (ByteArray ba#) = I# (sizeofByteArray# ba#)	2✔
76
77	st_ :: (State# s -> State# s) -> ST s ()
78	st_ m# = ST $ \s# -> (# m# s#, () #)	1✔
79	{-# INLINE st_ #-}
80
81	ioToST :: IO a -> ST RealWorld a
UNCOV 82	ioToST (IO m#) = ST m#	×
83	{-# INLINE ioToST #-}
84
85	newMutableByteArray :: Int -> ST s (MutableByteArray s)
86	newMutableByteArray (I# n#) =	2✔
87	ST $ \s# ->	2✔
88	case newByteArray# n# s# of	2✔
89	(# s'#, mba# #) -> (# s'#, MutableByteArray mba# #)	2✔
90	{-# INLINE newMutableByteArray #-}
91
92	newPinnedMutableByteArray :: Int -> ST s (MutableByteArray s)
93	newPinnedMutableByteArray (I# n#) =	2✔
94	ST $ \s# ->	2✔
95	case newPinnedByteArray# n# s# of	2✔
96	(# s'#, mba# #) -> (# s'#, MutableByteArray mba# #)	2✔
97	{-# INLINE newPinnedMutableByteArray #-}
98
99	freezeMutableByteArray :: MutableByteArray s -> ST s ByteArray
100	freezeMutableByteArray (MutableByteArray mba#) =	2✔
101	ST $ \s# ->	2✔
102	case unsafeFreezeByteArray# mba# s# of	2✔
103	(# s'#, ba# #) -> (# s'#, ByteArray ba# #)	2✔
104
105	writeWord8 :: MutableByteArray s -> Int -> Word8 -> ST s ()
106	writeWord8 (MutableByteArray mba#) (I# i#) (W8# w#) = st_ (writeWord8Array# mba# i# w#)	2✔
107	{-# INLINE writeWord8 #-}
108
109	writeByteSliceWord64LE :: MutableByteArray s -> Int -> Int -> Word64 -> ST s ()
110	writeByteSliceWord64LE mba fromByteIx toByteIx = go fromByteIx	2✔
111	where
112	go !i !z =	2✔
113	when (i < toByteIx) $ do	2✔
114	writeWord8 mba i (fromIntegral z :: Word8)	2✔
115	go (i + 1) (z `shiftR` 8)	2✔
116	{-# INLINE writeByteSliceWord64LE #-}
117
118	indexWord8 ::
119	ByteArray ->
120	-- \| Offset into immutable byte array in number of bytes
121	Int ->
122	Word8
123	indexWord8 (ByteArray ba#) (I# i#) =	2✔
124	W8# (indexWord8Array# ba# i#)	2✔
125	{-# INLINE indexWord8 #-}
126
127	indexWord64LE ::
128	ByteArray ->
129	-- \| Offset into immutable byte array in number of bytes
130	Int ->
131	Word64
132	#if defined WORDS_BIGENDIAN \|\| !(__GLASGOW_HASKELL__ >= 806)
133	indexWord64LE ba i = indexByteSliceWord64LE ba i (i + 8)
134	#else
135	indexWord64LE (ByteArray ba#) (I# i#)	2✔
136	\| wordSizeInBits == 64 = W64# (indexWord8ArrayAsWord64# ba# i#)	1✔
UNCOV 137	\| otherwise =	×
138	let !w32l = W32# (indexWord8ArrayAsWord32# ba# i#)	×
139	!w32u = W32# (indexWord8ArrayAsWord32# ba# (i# +# 4#))	×
140	in (fromIntegral w32u `shiftL` 32) .\|. fromIntegral w32l	×
141	#endif
142	{-# INLINE indexWord64LE #-}
143
144	indexByteSliceWord64LE ::
145	ByteArray ->
146	-- \| Starting offset in number of bytes
147	Int ->
148	-- \| Ending offset in number of bytes
149	Int ->
150	Word64
151	indexByteSliceWord64LE ba fromByteIx toByteIx = goWord8 fromByteIx 0	2✔
152	where
153	r = (toByteIx - fromByteIx) `rem` 8	2✔
154	nPadBits = if r == 0 then 0 else 8 * (8 - r)	1✔
155	goWord8 i !w64	2✔
156	\| i < toByteIx = goWord8 (i + 1) (shiftL w64 8 .\|. fromIntegral (indexWord8 ba i))	2✔
157	\| otherwise = byteSwap64 (shiftL w64 nPadBits)	1✔
158	{-# INLINE indexByteSliceWord64LE #-}
159
160	-- On big endian machines we need to write one byte at a time for consistency with little
161	-- endian machines. Also for GHC versions prior to 8.6 we don't have primops that can
162	-- write with byte offset, eg. writeWord8ArrayAsWord64# and writeWord8ArrayAsWord32#, so we
163	-- also must fallback to writing one byte a time. Such fallback results in about 3 times
164	-- slow down, which is not the end of the world.
165	writeWord64LE ::
166	MutableByteArray s ->
167	-- \| Offset into mutable byte array in number of bytes
168	Int ->
169	-- \| 8 bytes that will be written into the supplied array
170	Word64 ->
171	ST s ()
172	#if defined WORDS_BIGENDIAN \|\| !(__GLASGOW_HASKELL__ >= 806)
173	writeWord64LE mba i w64 =
174	writeByteSliceWord64LE mba i (i + 8) w64
175	#else
176	writeWord64LE (MutableByteArray mba#) (I# i#) w64@(W64# w64#)	2✔
177	\| wordSizeInBits == 64 = st_ (writeWord8ArrayAsWord64# mba# i# w64#)	1✔
UNCOV 178	\| otherwise = do	×
179	let !(W32# w32l#) = fromIntegral w64	×
180	!(W32# w32u#) = fromIntegral (w64 `shiftR` 32)	×
181	st_ (writeWord8ArrayAsWord32# mba# i# w32l#)	×
182	st_ (writeWord8ArrayAsWord32# mba# (i# +# 4#) w32u#)	×
183	#endif
184	{-# INLINE writeWord64LE #-}
185
186	getSizeOfMutableByteArray :: MutableByteArray s -> ST s Int
187	#if __GLASGOW_HASKELL__ >=802
188	getSizeOfMutableByteArray (MutableByteArray mba#) =	2✔
189	ST $ \s ->	2✔
190	case getSizeofMutableByteArray# mba# s of	2✔
191	(# s', n# #) -> (# s', I# n# #)	2✔
192	#else
193	getSizeOfMutableByteArray (MutableByteArray mba#) =
194	pure $! I# (sizeofMutableByteArray# mba#)
195	#endif
196	{-# INLINE getSizeOfMutableByteArray #-}
197
198	shortByteStringToByteArray :: ShortByteString -> ByteArray
UNCOV 199	shortByteStringToByteArray (SBS ba#) = ByteArray ba#	×
200	{-# INLINE shortByteStringToByteArray #-}
201
202	byteArrayToShortByteString :: ByteArray -> ShortByteString
203	byteArrayToShortByteString (ByteArray ba#) = SBS ba#	2✔
204	{-# INLINE byteArrayToShortByteString #-}
205
206	-- \| Convert a ShortByteString to ByteString by casting, whenever memory is pinned,
207	-- otherwise make a copy into a new pinned ByteString
208	shortByteStringToByteString :: ShortByteString -> ByteString
209	#if __GLASGOW_HASKELL__ < 802
210	shortByteStringToByteString ba = SBS.fromShort ba
211	#else
212	shortByteStringToByteString ba =	2✔
213	let !(SBS ba#) = ba in	2✔
214	if isTrue# (isByteArrayPinned# ba#)	1✔
215	then pinnedByteArrayToByteString ba#	2✔
UNCOV 216	else SBS.fromShort ba	×
217	{-# INLINE shortByteStringToByteString #-}
218
219	pinnedByteArrayToByteString :: ByteArray# -> ByteString
220	pinnedByteArrayToByteString ba# =	2✔
221	PS (pinnedByteArrayToForeignPtr ba#) 0 (I# (sizeofByteArray# ba#))	2✔
222	{-# INLINE pinnedByteArrayToByteString #-}
223
224	pinnedByteArrayToForeignPtr :: ByteArray# -> ForeignPtr a
225	pinnedByteArrayToForeignPtr ba# =	2✔
226	ForeignPtr (byteArrayContents# ba#) (PlainPtr (unsafeCoerce# ba#))	1✔
227	{-# INLINE pinnedByteArrayToForeignPtr #-}
228	#endif
229
230	-----------------
231	-- Boxed Array --
232	-----------------
233
234	data Array a = Array (Array# a)
235
236	data MutableArray s a = MutableArray (MutableArray# s a)
237
238	newMutableArray :: Int -> a -> ST s (MutableArray s a)
239	newMutableArray (I# n#) a =	2✔
240	ST $ \s# ->	2✔
241	case newArray# n# a s# of	1✔
242	(# s'#, ma# #) -> (# s'#, MutableArray ma# #)	2✔
243	{-# INLINE newMutableArray #-}
244
245	freezeMutableArray :: MutableArray s a -> ST s (Array a)
UNCOV 246	freezeMutableArray (MutableArray ma#) =	×
247	ST $ \s# ->	×
248	case unsafeFreezeArray# ma# s# of	×
249	(# s'#, a# #) -> (# s'#, Array a# #)	×
250	{-# INLINE freezeMutableArray #-}
251
252	sizeOfMutableArray :: MutableArray s a -> Int
253	sizeOfMutableArray (MutableArray ma#) = I# (sizeofMutableArray# ma#)	2✔
254	{-# INLINE sizeOfMutableArray #-}
255
256	readArray :: MutableArray s a -> Int -> ST s a
257	readArray (MutableArray ma#) (I# i#) = ST (readArray# ma# i#)	2✔
258	{-# INLINE readArray #-}
259
260	writeArray :: MutableArray s a -> Int -> a -> ST s ()
261	writeArray (MutableArray ma#) (I# i#) a = st_ (writeArray# ma# i# a)	2✔
262	{-# INLINE writeArray #-}
263
264	swapArray :: MutableArray s a -> Int -> Int -> ST s ()
265	swapArray ma i j = do	2✔
266	x <- readArray ma i	2✔
267	y <- readArray ma j	2✔
268	writeArray ma j x	2✔
269	writeArray ma i y	2✔
270	{-# INLINE swapArray #-}
271
272	-- \| Write contents of the list into the mutable array. Make sure that array is big
273	-- enough or segfault will happen.
274	fillMutableArrayFromList :: MutableArray s a -> [a] -> ST s ()
275	fillMutableArrayFromList ma = go 0	2✔
276	where
277	go _ [] = pure ()	1✔
278	go i (x : xs) = writeArray ma i x >> go (i + 1) xs	2✔
279	{-# INLINE fillMutableArrayFromList #-}
280
281	readListFromMutableArray :: MutableArray s a -> ST s [a]
282	readListFromMutableArray ma = go (len - 1) []	2✔
283	where
284	len = sizeOfMutableArray ma	2✔
285	go i !acc	2✔
286	\| i >= 0 = do	2✔
287	x <- readArray ma i	2✔
288	go (i - 1) (x : acc)	2✔
289	\| otherwise = pure acc	1✔
290	{-# INLINE readListFromMutableArray #-}
291
292	-- \| Generate a list of indices that will be used for swapping elements in uniform shuffling:
293	--
294	-- @
295	-- [ (0, n - 1)
296	-- , (0, n - 2)
297	-- , (0, n - 3)
298	-- , ...
299	-- , (0, 3)
300	-- , (0, 2)
301	-- , (0, 1)
302	-- ]
303	-- @
304	genSwapIndices ::
305	Monad m =>
306	-- \| Action that generates a Word in the supplied range.
307	(Word -> m Word) ->
308	-- \| Number of index swaps to generate.
309	Word ->
310	m [Int]
311	genSwapIndices genWordR n = go 1 []	×
312	where
UNCOV 313	go i !acc	×
314	\| i >= n = pure acc	×
315	\| otherwise = do	×
316	x <- genWordR i	×
317	let !xi = fromIntegral x	×
318	go (i + 1) (xi : acc)	×
319	{-# INLINE genSwapIndices #-}
320
321	-- \| Implementation of mutable version of Fisher-Yates shuffle. Unfortunately, we cannot generally
322	-- interleave pseudo-random number generation and mutation of `ST` monad, therefore we have to
323	-- pre-generate all of the index swaps with `genSwapIndices` and store them in a list before we can
324	-- perform the actual swaps.
325	shuffleListM :: Monad m => (Word -> m Word) -> [a] -> m [a]
UNCOV 326	shuffleListM genWordR ls	×
327	\| len <= 1 = pure ls	×
328	\| otherwise = do	×
NEW 329	swapIxs <- genSwapIndices genWordR (fromIntegral len)	×
NEW 330	pure $ runST $ do	×
NEW 331	ma <- newMutableArray len $ error "Impossible: shuffleListM"	×
NEW 332	fillMutableArrayFromList ma ls	×
333
334	-- Shuffle elements of the mutable array according to the uniformly generated index swap list
NEW 335	let goSwap _ [] = pure ()	×
NEW 336	goSwap i (j : js) = swapArray ma i j >> goSwap (i - 1) js	×
NEW 337	goSwap (len - 1) swapIxs	×
338
NEW 339	readListFromMutableArray ma	×
340	where
UNCOV 341	len = length ls	×
342	{-# INLINE shuffleListM #-}
343
344	-- \| This is a ~x2-x3 more efficient version of `shuffleListM`. It is more efficient because it does
345	-- not need to pregenerate a list of indices and instead generates them on demand. Because of this the
346	-- result that will be produced will differ for the same generator, since the order in which index
347	-- swaps are generated is reversed.
348	--
349	-- Unfortunately, most stateful generator monads can't handle `MonadTrans`, so this version is only
350	-- used for implementing the pure shuffle.
351	shuffleListST :: (Monad (t (ST s)), MonadTrans t) => (Word -> t (ST s) Word) -> [a] -> t (ST s) [a]
352	shuffleListST genWordR ls	2✔
353	\| len <= 1 = pure ls	2✔
354	\| otherwise = do	1✔
355	ma <- lift $ newMutableArray len $ error "Impossible: shuffleListST"	1✔
356	lift $ fillMutableArrayFromList ma ls	2✔
357
358	-- Shuffle elements of the mutable array according to the uniformly generated index swap
359	let goSwap i =	2✔
360	when (i > 0) $ do	2✔
361	j <- genWordR $ (fromIntegral :: Int -> Word) i	2✔
362	lift $ swapArray ma i ((fromIntegral :: Word -> Int) j)	2✔
363	goSwap (i - 1)	2✔
364	goSwap (len - 1)	2✔
365
366	lift $ readListFromMutableArray ma	2✔
367	where
368	len = length ls	2✔
369	{-# INLINE shuffleListST #-}

haskell / random / 434

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