module Bench.Vector.TestData.Graph
  ( randomGraph
  ) where

import System.Random.Stateful
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed as U

randomGraph
  :: (StatefulGen g m, MV.PrimMonad m)
  => g
  -> Int
  -> m (Int, U.Vector Int, U.Vector Int)
randomGraph :: forall g (m :: * -> *).
(StatefulGen g m, PrimMonad m) =>
g -> Int -> m (Int, Vector Int, Vector Int)
randomGraph g
g Int
edges = do
  let vertices :: Int
vertices = Int
edges Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
10
  MVector (PrimState m) [Int]
marr <- Int -> [Int] -> m (MVector (PrimState m) [Int])
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
MV.replicate Int
vertices []
  g -> Int -> MVector (PrimState m) [Int] -> Int -> m ()
forall g (m :: * -> *).
(StatefulGen g m, PrimMonad m) =>
g -> Int -> MVector (PrimState m) [Int] -> Int -> m ()
addRandomEdges g
g Int
vertices MVector (PrimState m) [Int]
marr Int
edges
  Vector [Int]
arr <- MVector (PrimState m) [Int] -> m (Vector [Int])
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze MVector (PrimState m) [Int]
marr
  let ([Int]
as, [Int]
bs) = [(Int, Int)] -> ([Int], [Int])
forall a b. [(a, b)] -> ([a], [b])
unzip [ (Int
i, Int
j) | Int
i <- [Int
0 .. Int
vertices Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
j <- Vector [Int]
arr Vector [Int] -> Int -> [Int]
forall a. Vector a -> Int -> a
V.! Int
i ]
  (Int, Vector Int, Vector Int) -> m (Int, Vector Int, Vector Int)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
vertices, [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
U.fromList [Int]
as, [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
U.fromList [Int]
bs)

addRandomEdges
  :: (StatefulGen g m, MV.PrimMonad m)
  => g
  -> Int
  -> MV.MVector (MV.PrimState m) [Int]
  -> Int
  -> m ()
addRandomEdges :: forall g (m :: * -> *).
(StatefulGen g m, PrimMonad m) =>
g -> Int -> MVector (PrimState m) [Int] -> Int -> m ()
addRandomEdges g
g Int
vertices MVector (PrimState m) [Int]
arr = Int -> m ()
forall {m :: * -> *} {t}.
(PrimState m ~ PrimState m, Eq t, Num t, StatefulGen g m,
 PrimMonad m) =>
t -> m ()
fill
  where
    fill :: t -> m ()
fill t
0 = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    fill t
e = do
      Int
m1 <- (Int, Int) -> g -> m Int
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (Int, Int) -> g -> m Int
uniformRM (Int
0, Int
vertices Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) g
g
      Int
m2 <- (Int, Int) -> g -> m Int
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (Int, Int) -> g -> m Int
uniformRM (Int
0, Int
vertices Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) g
g
      let lo :: Int
lo = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m1 Int
m2
          hi :: Int
hi = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
m1 Int
m2
      [Int]
ns <- MVector (PrimState m) [Int] -> Int -> m [Int]
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.read MVector (PrimState m) [Int]
MVector (PrimState m) [Int]
arr Int
lo
      if Int
lo Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
hi Bool -> Bool -> Bool
|| Int
hi Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
ns
        then t -> m ()
fill t
e
        else MVector (PrimState m) [Int] -> Int -> [Int] -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) [Int]
MVector (PrimState m) [Int]
arr Int
lo (Int
hi Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
ns) m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> t -> m ()
fill (t
e t -> t -> t
forall a. Num a => a -> a -> a
- t
1)