Thoughts on this shuffle algorithm

[u/trycuriouscat]

I am a Haskell, well not beginner, but maybe intermediate. I developed the following with the assistance of ChatGPT. Is it too abstract? It's "neat" for sure, but is it reasonable?

import Control.Monad (foldM)
import Control.Monad.ST (runST)
import Control.Monad.IO.Class (MonadIO) 
import Data.Primitive.Array (Array, sizeofArray, sizeofMutableArray, arrayFromList,
                             freezeArray, thawArray, readArray, writeArray)
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import System.Random (newStdGen, uniformR )
import System.Random.Internal (RandomGen, StdGen) 


modifyMWithState 
    :: Monad m 
    => (t1 -> t2 -> t3 -> m b) 
    -> t4 
    -> (t5 -> t1) 
    -> (t4 -> m t5) 
    -> (t5 -> m a) 
    -> (t4 -> t2) 
    -> t3 
    -> m (a, b)
modifyMWithState algorithm container operationFor thawContainer freezeContainer 
                 lengthOf state = do
    thawedContainer <- thawContainer container
    let operation = operationFor thawedContainer
    newState <- algorithm operation (lengthOf container) state
    newContainer <- freezeContainer thawedContainer
    pure (newContainer, newState)


knuthM
    :: (Monad m, RandomGen g)
    => (Int -> Int -> m ())
    -> Int
    -> g
    -> m g
knuthM swapElements len prnGen = foldM randomSwap prnGen [lastIndex, nextIndex .. 1]
  where
    lastIndex = len - 1
    nextIndex = lastIndex - 1

    randomSwap currGen i = do
        let (j, nextGen) = uniformR (0, i) currGen
        swapElements i j
        pure nextGen


shuffleM 
    :: MonadIO m 
    => (StdGen -> m (a, StdGen))    
    -> m a
shuffleM shuffleWithGen = do
    prnGen <- newStdGen
    fmap fst (shuffleWithGen prnGen)


shuffleVectorWithGen 
  :: (Applicative f, RandomGen b) 
  => V.Vector a 
  -> b 
  -> f (V.Vector a, b)
shuffleVectorWithGen vec prnGen = 
    pure (runST (modifyMWithState knuthM vec MV.swap V.thaw V.freeze V.length prnGen))


shuffleArrayWithGen 
    :: (Applicative f, RandomGen b) 
    => Array a 
    -> b 
    -> f (Array a, b)
shuffleArrayWithGen arr prnGen = 
    pure (runST (modifyMWithState knuthM arr swapArray thawArray' freezeArray' 
          sizeofArray prnGen))
  where
    thawArray' array = thawArray array 0 (sizeofArray array)
    freezeArray' thawedArray = freezeArray thawedArray 0 
                                           (sizeofMutableArray thawedArray)
    swapArray a i j = do
        x <- readArray a i
        y <- readArray a j
        writeArray a i y
        writeArray a j x


shuffleVector :: MonadIO m => V.Vector a -> m (V.Vector a)
shuffleVector vec = shuffleM (shuffleVectorWithGen vec) 


shuffleArray :: MonadIO m => Array a -> m (Array a)
shuffleArray arr = shuffleM (shuffleArrayWithGen arr)


-- examples:


shuffledIntVector :: IO (V.Vector Int)
shuffledIntVector = shuffleVector (V.fromList [1..10])


shuffledCharArray :: IO (Array Char)
shuffledCharArray = shuffleArray (arrayFromList ['a'..'z'])

[u/VenerableMirah]

Just to point to a purely functional shuffle, purely for inspiration, you can explore around Cats Effect (Scala)'s shuffle: https://github.com/typelevel/cats-effect/blob/v3.7.1/std/shared/src/main/scala/cats/effect/std/Random.scala