rando-0.0.0.2: src/System/Random/Shuffle/FisherYates.hs

module System.Random.Shuffle.FisherYates (
     shuffle
   ) where

import Control.Monad
import Control.Monad.ST
import Data.STRef
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import System.Random.TF
import System.Random.TF.Instances

-- The Fisher-Yates shuffle, asymptotically optimal for both space and time:
shuffle :: [x] -> IO [x]
shuffle list = do

   initGen <- newTFGen

   pure $ V.toList $ runST $ do
      genVar <- newSTRef initGen

      v <- V.thaw $ V.fromList list
      let vLen = VM.length v

      forM_ [vLen - 1, vLen - 2 .. 1] $ \n -> do
         indexToSwap <- randSTFromZero genVar n
         -- The "unsafe" means it doesn't have bounds checks, which is fine here:
         VM.unsafeSwap v indexToSwap n

      -- This "unsafe" is also safe here:
      V.unsafeFreeze v



-- | Picks from 0 to maxval, inclusive
randSTFromZero :: STRef s TFGen -> Int -> ST s Int
randSTFromZero genVar maxVal = do
   g0 <- readSTRef genVar
   let (x, g1) = randomR (0, maxVal) g0
   writeSTRef genVar g1
   pure x