FYShuffle.hs

{-# LANGUAGE RankNTypes, BangPatterns #-}

module FYSHuffle (fyShuffle) where

import Control.Monad
import Control.Monad.Primitive
import Control.Monad.Random
import Control.Monad.ST

import qualified Data.Vector.Generic as V
import Data.Vector.Generic (Vector, Mutable)

import qualified Data.Vector.Generic.Mutable as M
import Data.Vector.Generic.Mutable (MVector)

fyShuffle :: (Vector v a, RandomGen g) => v a -> Rand g (v a)
fyShuffle v = liftRand $ \g -> modify (flip execRandT g . fyShuffleM) v

modify :: Vector v a => (forall s. Mutable v s a -> ST s b) -> v a -> (v a, b)
{-# INLINE modify #-}
modify f !v =
  runST $ do
    mv <- M.new (V.length v)
    V.unsafeCopy mv v
    b <- f mv
    mv' <- V.unsafeFreeze mv
    return (mv', b)

fyShuffleM :: (MVector v a, RandomGen g) => forall s. v s a -> RandT g (ST s) ()
fyShuffleM v = forM_ [0 .. n - 1] $ \i -> getRandomR (i, n - 1) >>= M.unsafeSwap v i
  where
    n = M.length v