module List.Shuffle
(
shuffle,
shuffle_,
shuffleIO,
sample,
sample_,
sampleIO,
)
where
import Control.Monad.IO.Class (MonadIO)
import Control.Monad.ST (runST)
import Control.Monad.ST.Strict (ST)
import Data.Foldable qualified as Foldable
import Data.Primitive.Array qualified as Array
import System.Random (RandomGen)
import System.Random qualified as Random
shuffle :: (RandomGen g) => [a] -> g -> ([a], g)
shuffle :: forall g a. RandomGen g => [a] -> g -> ([a], g)
shuffle [a]
list g
gen0 =
forall a. (forall s. ST s a) -> a
runST do
MutableArray s a
array <- forall a s. [a] -> ST s (MutableArray s a)
listToMutableArray [a]
list
g
gen1 <- forall a g s. RandomGen g => Int -> MutableArray s a -> g -> ST s g
shuffleN (forall s a. MutableArray s a -> Int
Array.sizeofMutableArray MutableArray s a
array forall a. Num a => a -> a -> a
- Int
1) MutableArray s a
array g
gen0
Array a
array1 <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
Array.unsafeFreezeArray MutableArray s a
array
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList Array a
array1, g
gen1)
{-# SPECIALIZE shuffle :: [a] -> Random.StdGen -> ([a], Random.StdGen) #-}
shuffleN :: forall a g s. (RandomGen g) => Int -> Array.MutableArray s a -> g -> ST s g
shuffleN :: forall a g s. RandomGen g => Int -> MutableArray s a -> g -> ST s g
shuffleN Int
n0 MutableArray s a
array =
Int -> g -> ST s g
go Int
0
where
go :: Int -> g -> ST s g
go :: Int -> g -> ST s g
go !Int
i g
gen0
| Int
i forall a. Ord a => a -> a -> Bool
>= Int
n = forall (f :: * -> *) a. Applicative f => a -> f a
pure g
gen0
| Bool
otherwise = do
let (Int
j, g
gen1) = forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
Random.uniformR (Int
i, Int
m) g
gen0
forall s a. Int -> Int -> MutableArray s a -> ST s ()
swapArrayElems Int
i Int
j MutableArray s a
array
Int -> g -> ST s g
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) g
gen1
n :: Int
n = forall a. Ord a => a -> a -> a
min Int
n0 Int
m
m :: Int
m = forall s a. MutableArray s a -> Int
Array.sizeofMutableArray MutableArray s a
array forall a. Num a => a -> a -> a
- Int
1
{-# SPECIALIZE shuffleN :: Int -> Array.MutableArray s a -> Random.StdGen -> ST s Random.StdGen #-}
shuffle_ :: (RandomGen g) => [a] -> g -> [a]
shuffle_ :: forall g a. RandomGen g => [a] -> g -> [a]
shuffle_ [a]
list g
g =
forall a b. (a, b) -> a
fst (forall g a. RandomGen g => [a] -> g -> ([a], g)
shuffle [a]
list g
g)
{-# SPECIALIZE shuffle_ :: [a] -> Random.StdGen -> [a] #-}
shuffleIO :: (MonadIO m) => [a] -> m [a]
shuffleIO :: forall (m :: * -> *) a. MonadIO m => [a] -> m [a]
shuffleIO [a]
list =
forall g a. RandomGen g => [a] -> g -> [a]
shuffle_ [a]
list forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => m StdGen
Random.newStdGen
{-# SPECIALIZE shuffleIO :: [a] -> IO [a] #-}
sample :: (RandomGen g) => Int -> [a] -> g -> ([a], g)
sample :: forall g a. RandomGen g => Int -> [a] -> g -> ([a], g)
sample Int
n [a]
list g
gen0 =
forall a. (forall s. ST s a) -> a
runST do
MutableArray s a
array <- forall a s. [a] -> ST s (MutableArray s a)
listToMutableArray [a]
list
g
gen1 <- forall a g s. RandomGen g => Int -> MutableArray s a -> g -> ST s g
shuffleN Int
n MutableArray s a
array g
gen0
Array a
array1 <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
Array.unsafeFreezeArray MutableArray s a
array
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Int -> [a] -> [a]
take Int
n (forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList Array a
array1), g
gen1)
{-# SPECIALIZE sample :: Int -> [a] -> Random.StdGen -> ([a], Random.StdGen) #-}
sample_ :: (RandomGen g) => Int -> [a] -> g -> [a]
sample_ :: forall g a. RandomGen g => Int -> [a] -> g -> [a]
sample_ Int
n [a]
list g
g =
forall a b. (a, b) -> a
fst (forall g a. RandomGen g => Int -> [a] -> g -> ([a], g)
sample Int
n [a]
list g
g)
{-# SPECIALIZE sample_ :: Int -> [a] -> Random.StdGen -> [a] #-}
sampleIO :: (MonadIO m) => Int -> [a] -> m [a]
sampleIO :: forall (m :: * -> *) a. MonadIO m => Int -> [a] -> m [a]
sampleIO Int
n [a]
list =
forall g a. RandomGen g => Int -> [a] -> g -> [a]
sample_ Int
n [a]
list forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => m StdGen
Random.newStdGen
{-# SPECIALIZE sampleIO :: Int -> [a] -> IO [a] #-}
swapArrayElems :: Int -> Int -> Array.MutableArray s a -> ST s ()
swapArrayElems :: forall s a. Int -> Int -> MutableArray s a -> ST s ()
swapArrayElems Int
i Int
j MutableArray s a
array = do
a
x <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
Array.readArray MutableArray s a
array Int
i
a
y <- forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
Array.readArray MutableArray s a
array Int
j
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
Array.writeArray MutableArray s a
array Int
i a
y
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
Array.writeArray MutableArray s a
array Int
j a
x
{-# INLINE swapArrayElems #-}
listToMutableArray :: [a] -> ST s (Array.MutableArray s a)
listToMutableArray :: forall a s. [a] -> ST s (MutableArray s a)
listToMutableArray [a]
list = do
MutableArray (PrimState (ST s)) a
array <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
Array.newArray (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
list) forall a. HasCallStack => a
undefined
let writeElems :: Int -> [a] -> ST s ()
writeElems !Int
i = \case
[] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
a
x : [a]
xs -> do
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
Array.writeArray MutableArray (PrimState (ST s)) a
array Int
i a
x
Int -> [a] -> ST s ()
writeElems (Int
i forall a. Num a => a -> a -> a
+ Int
1) [a]
xs
Int -> [a] -> ST s ()
writeElems Int
0 [a]
list
forall (f :: * -> *) a. Applicative f => a -> f a
pure MutableArray (PrimState (ST s)) a
array
{-# INLINE listToMutableArray #-}