Random generation of DAG

3 files changed, 42 insertions, 16 deletions

diff --git a/src/Test/VeriFuzz/CodeGen.hs b/src/Test/VeriFuzz/CodeGen.hs
index 1872b28..3965b16 100644
--- a/src/Test/VeriFuzz/CodeGen.hs
+++ b/src/Test/VeriFuzz/CodeGen.hs
@@ -1 +1 @@
-module VeriFuzz.CodeGen where
+module Test.VeriFuzz.CodeGen where
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)
diff --git a/src/Test/VeriFuzz/Types.hs b/src/Test/VeriFuzz/Types.hs
index 9c0de17..0528972 100644
--- a/src/Test/VeriFuzz/Types.hs
+++ b/src/Test/VeriFuzz/Types.hs
@@ -1,8 +1,21 @@
-module VeriFuzz.Types where
+module Test.VeriFuzz.Types where
+
+import Test.QuickCheck
+import System.Random
 
 data Gate = And
           | Or
           | Xor
           | Nor
           | Nand
-          deriving (Show, Eq, Ord)
+          deriving (Show, Eq, Enum, Bounded)
+
+instance Random Gate where
+  randomR (a, b) g =
+    case randomR (fromEnum a, fromEnum b) g of
+      (x, g') -> (toEnum x, g')
+
+  random g = randomR (minBound, maxBound) g
+
+instance Arbitrary Gate where
+  arbitrary = elements [And, Or, Xor, Nor, Nand]