{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}
module Test.Vector.Shuffle
( -- * Simple
shuffle
, shuffleNonEmpty
-- * Advanced
, mkSeed
, shuffleWith
, shuffleNonEmptyWith
) where
import Prelude
import Control.Monad
( forM_ )
import Control.Monad.Trans.Class
( lift )
import Control.Monad.Trans.State.Strict
( evalStateT, state )
import Crypto.Hash
( hash )
import Crypto.Hash.Algorithms
( MD5 )
import Data.List.NonEmpty
( NonEmpty (..) )
import Data.Maybe
( fromMaybe )
import Data.Text
( Text )
import Data.Vector.Mutable
( IOVector )
import Data.Word
( Word8 )
import System.Random
( RandomGen, StdGen, mkStdGen, newStdGen, randomR )
import qualified Data.ByteArray as BA
import qualified Data.ByteString as BS
import qualified Data.List.NonEmpty as NE
import qualified Data.Text.Encoding as T
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
-- | Generate a random generator seed from a text string
mkSeed :: Text -> StdGen
mkSeed = mkStdGen . toInt . quickHash . T.encodeUtf16LE
where
quickHash = BA.convert . hash @_ @MD5
toInt = snd . BS.foldl' exponentiation (0,0)
where
exponentiation :: (Int, Int) -> Word8 -> (Int, Int)
exponentiation (e, n) i = (e+1, n + fromIntegral i*2^e)
-- | Shuffles a list of elements.
--
-- >>> shuffle (outputs coinSel)
-- [...]
shuffle :: [a] -> IO [a]
shuffle xs = newStdGen >>= flip shuffleWith xs
-- | Like 'shuffle', but from a given seed. 'shuffle' will use a randomly
-- generate seed using 'newStdGen' from @System.Random@.
--
-- __Properties:__
--
-- - @shuffleWith g es == shuffleWith g es@
-- - @∃Δ> 1. g ≠g', length es > Δ⇒ shuffleWith g es ≠shuffleWith g' es@
shuffleWith :: RandomGen g => g -> [a] -> IO [a]
shuffleWith seed = modifyInPlace $ \v -> flip evalStateT seed $ do
let (lo, hi) = (0, MV.length v - 1)
forM_ [lo .. hi] $ \i -> do
j <- fromInteger <$> state (randomR (fromIntegral lo, fromIntegral hi))
lift $ swapElems v i j
where
swapElems :: IOVector a -> Int -> Int -> IO ()
swapElems v i j = do
x <- MV.read v i
y <- MV.read v j
MV.write v i y
MV.write v j x
modifyInPlace :: forall a. (IOVector a -> IO ()) -> [a] -> IO [a]
modifyInPlace f xs = do
v' <- V.thaw $ V.fromList xs
f v'
V.toList <$> V.freeze v'
-- | Shuffles a /non-empty/ list of elements.
--
-- See 'shuffle'.
--
shuffleNonEmpty :: NonEmpty a -> IO (NonEmpty a)
shuffleNonEmpty xs = newStdGen >>= flip shuffleNonEmptyWith xs
-- | Shuffles a /non-empty/ list of elements with the given random generator.
--
-- See 'shuffleWith'.
--
shuffleNonEmptyWith :: RandomGen g => g -> NonEmpty a -> IO (NonEmpty a)
shuffleNonEmptyWith g =
fmap (fromMaybe raiseError . NE.nonEmpty) . shuffleWith g . NE.toList
where
raiseError = error "shuffleNonEmptyWith encountered an empty list."