From 1889089632d7b079d82cf852538d09de4584d832 Mon Sep 17 00:00:00 2001 From: Alexey Kuleshevich Date: Sun, 1 Dec 2024 13:10:41 -0700 Subject: [PATCH 1/2] Extract all array helper functionality into its own internal module --- random.cabal | 1 + src/System/Random/Array.hs | 206 ++++++++++++++++++++++++++++++++++ src/System/Random/Internal.hs | 167 +-------------------------- stack-old.yaml | 1 - stack-oldish.yaml | 1 + 5 files changed, 210 insertions(+), 166 deletions(-) create mode 100644 src/System/Random/Array.hs diff --git a/random.cabal b/random.cabal index b8246409..1af91647 100644 --- a/random.cabal +++ b/random.cabal @@ -88,6 +88,7 @@ library System.Random.Seed System.Random.Stateful other-modules: + System.Random.Array System.Random.GFinite hs-source-dirs: src diff --git a/src/System/Random/Array.hs b/src/System/Random/Array.hs new file mode 100644 index 00000000..4cdf9e7d --- /dev/null +++ b/src/System/Random/Array.hs @@ -0,0 +1,206 @@ +{-# 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 diff --git a/src/System/Random/Internal.hs b/src/System/Random/Internal.hs index 0a102de7..66d0e81a 100644 --- a/src/System/Random/Internal.hs +++ b/src/System/Random/Internal.hs @@ -12,7 +12,6 @@ {-# LANGUAGE Trustworthy #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE TypeFamilyDependencies #-} -{-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE UnliftedFFITypes #-} {-# OPTIONS_HADDOCK hide, not-home #-} @@ -105,8 +104,8 @@ import Control.Monad.State.Strict (MonadState(..), State, StateT(..), execStateT import Control.Monad.Trans (lift, MonadTrans) import Data.Array.Byte (ByteArray(..), MutableByteArray(..)) import Data.Bits +import Data.ByteString (ByteString) import Data.ByteString.Short.Internal (ShortByteString(SBS)) -import qualified Data.ByteString.Short.Internal as SBS (fromShort) import Data.IORef (IORef, newIORef) import Data.Int import Data.List (sortOn) @@ -120,19 +119,11 @@ import GHC.ST (ST(..)) import GHC.Word import Numeric.Natural (Natural) import System.IO.Unsafe (unsafePerformIO) +import System.Random.Array import System.Random.GFinite (Cardinality(..), GFinite(..), Finite) import qualified System.Random.SplitMix as SM import qualified System.Random.SplitMix32 as SM32 import Data.Kind -#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" -- | This is a binary form of pseudo-random number generator's state. It is designed to be -- safe and easy to use for input/output operations like restoring from file, transmitting @@ -612,157 +603,6 @@ uniformByteString n g = (shortByteStringToByteString $ byteArrayToShortByteString byteArray, g') {-# INLINE uniformByteString #-} --- 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 - -- | Same as 'genShortByteStringIO', but runs in 'ST'. -- -- @since 1.2.0 @@ -1262,9 +1102,6 @@ instance UniformRange Int64 where {-# INLINE uniformRM #-} isInRange = isInRangeOrd -wordSizeInBits :: Int -wordSizeInBits = finiteBitSize (0 :: Word) - instance Uniform Int where uniformM | wordSizeInBits == 64 = diff --git a/stack-old.yaml b/stack-old.yaml index 8b10ec6e..35ca7f5f 100644 --- a/stack-old.yaml +++ b/stack-old.yaml @@ -13,5 +13,4 @@ extra-deps: - inspection-testing-0.4.5.0@sha256:938e7ce2ef42033071a5e60198c6e19ab61c411f5879b85821247a504f131768,8058 - tasty-inspection-testing-0.1@sha256:9c5e76345168fd3a59b43d305eebf8df3c792ce324c66bbdee45b54aa7d2c0ad,1214 - primitive-0.7.4.0@sha256:89b88a3e08493b7727fa4089b0692bfbdf7e1e666ef54635f458644eb8358764,2857 -- vector-0.12.3.1@sha256:fffbd00912d69ed7be9bc7eeb09f4f475e0d243ec43f916a9fd5bbd219ce7f3e,8238 - data-array-byte-0.1.0.1@sha256:ad89e28b2b046175698fbf542af2ce43e5d2af50aae9f48d12566b1bb3de1d3c,1989 diff --git a/stack-oldish.yaml b/stack-oldish.yaml index 06ec8886..89a22647 100644 --- a/stack-oldish.yaml +++ b/stack-oldish.yaml @@ -3,3 +3,4 @@ packages: - . extra-deps: - data-array-byte-0.1.0.1@sha256:2ef1bd3511e82ba56f7f23cd793dd2da84338a1e7c2cbea5b151417afe3baada,1989 +- primitive-0.9.0.0@sha256:de20bf4eff1f972088854c8efda6eaca2d3147aff62232c3707f059152638759,3203 From 9c491eb5904d90bcba70b377e9ef55b425c85d55 Mon Sep 17 00:00:00 2001 From: Alexey Kuleshevich Date: Sun, 1 Dec 2024 12:54:17 -0700 Subject: [PATCH 2/2] Remove unused stack file and redundant import --- stack-oldish.yaml | 6 ------ test/Spec/Range.hs | 1 - 2 files changed, 7 deletions(-) delete mode 100644 stack-oldish.yaml diff --git a/stack-oldish.yaml b/stack-oldish.yaml deleted file mode 100644 index 89a22647..00000000 --- a/stack-oldish.yaml +++ /dev/null @@ -1,6 +0,0 @@ -resolver: lts-18.28 -packages: -- . -extra-deps: -- data-array-byte-0.1.0.1@sha256:2ef1bd3511e82ba56f7f23cd793dd2da84338a1e7c2cbea5b151417afe3baada,1989 -- primitive-0.9.0.0@sha256:de20bf4eff1f972088854c8efda6eaca2d3147aff62232c3707f059152638759,3203 diff --git a/test/Spec/Range.hs b/test/Spec/Range.hs index 6a34dbfb..f10fa9c2 100644 --- a/test/Spec/Range.hs +++ b/test/Spec/Range.hs @@ -7,7 +7,6 @@ module Spec.Range , uniformRangeWithinExcludedD ) where -import System.Random.Internal import System.Random.Stateful import Data.Proxy