{-# LANGUAGE CPP #-}
module Test.Hspec.Core.Shuffle (
  shuffleForest
#ifdef TEST
, shuffle
, mkArray
#endif
) where

import           Prelude ()
import           Test.Hspec.Core.Compat
import           Test.Hspec.Core.Tree

import           System.Random
import           Control.Monad.ST
import           Data.STRef
import           Data.Array.ST

shuffleForest :: STRef st StdGen -> [Tree c a] -> ST st [Tree c a]
shuffleForest :: forall st c a. STRef st StdGen -> [Tree c a] -> ST st [Tree c a]
shuffleForest STRef st StdGen
ref [Tree c a]
xs = (forall st a. STRef st StdGen -> [a] -> ST st [a]
shuffle STRef st StdGen
ref [Tree c a]
xs forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall st c a. STRef st StdGen -> Tree c a -> ST st (Tree c a)
shuffleTree STRef st StdGen
ref))

shuffleTree :: STRef st StdGen -> Tree c a -> ST st (Tree c a)
shuffleTree :: forall st c a. STRef st StdGen -> Tree c a -> ST st (Tree c a)
shuffleTree STRef st StdGen
ref Tree c a
t = case Tree c a
t of
  Node String
d [Tree c a]
xs -> forall c a. String -> [Tree c a] -> Tree c a
Node String
d forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall st c a. STRef st StdGen -> [Tree c a] -> ST st [Tree c a]
shuffleForest STRef st StdGen
ref [Tree c a]
xs
  NodeWithCleanup Maybe (String, Location)
loc c
c [Tree c a]
xs -> forall c a. Maybe (String, Location) -> c -> [Tree c a] -> Tree c a
NodeWithCleanup Maybe (String, Location)
loc c
c forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall st c a. STRef st StdGen -> [Tree c a] -> ST st [Tree c a]
shuffleForest STRef st StdGen
ref [Tree c a]
xs
  Leaf {} -> forall (m :: * -> *) a. Monad m => a -> m a
return Tree c a
t

shuffle :: STRef st StdGen -> [a] -> ST st [a]
shuffle :: forall st a. STRef st StdGen -> [a] -> ST st [a]
shuffle STRef st StdGen
ref [a]
xs = do
  STArray st Int a
arr <- forall a st. [a] -> ST st (STArray st Int a)
mkArray [a]
xs
  bounds :: (Int, Int)
bounds@(Int
_, Int
n) <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
getBounds STArray st Int a
arr
  forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a. Ix a => (a, a) -> [a]
range (Int, Int)
bounds) forall a b. (a -> b) -> a -> b
$ \ Int
i -> do
    Int
j <- forall {b}. Random b => (b, b) -> ST st b
randomIndex (Int
i, Int
n)
    a
vi <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray st Int a
arr Int
i
    a
vj <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray st Int a
arr Int
j
    forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray st Int a
arr Int
j a
vi
    forall (m :: * -> *) a. Monad m => a -> m a
return a
vj
  where
    randomIndex :: (b, b) -> ST st b
randomIndex (b, b)
bounds = do
      (b
a, StdGen
gen) <- forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (b, b)
bounds forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s a. STRef s a -> ST s a
readSTRef STRef st StdGen
ref
      forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef st StdGen
ref StdGen
gen
      forall (m :: * -> *) a. Monad m => a -> m a
return b
a

mkArray :: [a] -> ST st (STArray st Int a)
mkArray :: forall a st. [a] -> ST st (STArray st Int a)
mkArray [a]
xs = forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> [e] -> m (a i e)
newListArray (Int
1, forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) [a]
xs