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
  marr <- Int -> [Int] -> m (MVector (PrimState m) [Int])
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
MV.replicate Int
vertices []
  addRandomEdges g vertices marr edges
  arr <- V.unsafeFreeze marr
  let (as, bs) = unzip [ (i, j) | i <- [0 .. vertices - 1], j <- arr V.! i ]
  return (vertices, U.fromList as, U.fromList 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
      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
      m2 <- uniformRM (0, vertices - 1) g
      let lo = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m1 Int
m2
          hi = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
m1 Int
m2
      ns <- MV.read arr lo
      if lo == hi || hi `elem` ns
        then fill e
        else MV.write arr lo (hi : ns) >> fill (e - 1)