1 files changed, 26 insertions, 13 deletions
diff --git a/src/Test/VeriFuzz/Graph/Random.hs b/src/Test/VeriFuzz/Graph/Random.hs
index a0937d0..2b91157 100644
--- a/src/Test/VeriFuzz/Graph/Random.hs
+++ b/src/Test/VeriFuzz/Graph/Random.hs
@@ -1,13 +1,26 @@
-module Test.VeriFuzz.Graph.Random
-  ( randomDAG
-  ) where
-
-import Data.Graph.Inductive
-import Test.QuickCheck
-
-randomDAG :: (Arbitrary a)
-          => GenIO    -- ^ The random number generator to use
-          -> Int      -- ^ The number of nodes
-          -> IO (Gr (LNode a) e) -- ^ The generated graph. It uses Arbitrary to
-                      -- generate random instances of each node
-randomDAG = do
+module Test.VeriFuzz.Graph.Random where
+
+import Data.Graph.Inductive (Graph, LNode, LEdge, mkGraph)
+import Test.QuickCheck (Arbitrary, Gen, arbitrary, generate, infiniteListOf, suchThat, listOf, scale, resize)
+
+import Test.VeriFuzz.Types
+
+arbitraryEdge :: (Arbitrary e) => Int -> Gen (LEdge e)
+arbitraryEdge n = do
+  x <- with $ \a -> a < n && a > 0 && a /= n-1
+  y <- with $ \a -> x < a && a < n && a > 0
+  z <- arbitrary
+  return (x, y, z)
+    where
+      with = suchThat . resize n $ arbitrary
+
+randomDAG :: (Arbitrary l, Arbitrary e, Graph gr)
+          => Int         -- ^ The number of nodes
+          -> IO (gr l e) -- ^ The generated graph. It uses Arbitrary to
+                         -- generate random instances of each node
+randomDAG n = do
+  list <- generate . infiniteListOf $ arbitrary
+  l <- generate . infiniteListOf $ arbitraryEdge n
+  return . mkGraph (nodes list) $ take (2*n) l
+    where
+      nodes l = (zip [0..n] $ take n l)