Type fixes

1 files changed, 8 insertions, 7 deletions

diff --git a/src/Test/VeriFuzz/Graph/ASTGen.hs b/src/Test/VeriFuzz/Graph/ASTGen.hs
index d3e6ea5..b5fb944 100644
--- a/src/Test/VeriFuzz/Graph/ASTGen.hs
+++ b/src/Test/VeriFuzz/Graph/ASTGen.hs
@@ -35,11 +35,12 @@ fromGate Xor = BinXor
 
 genPortsAST :: Circuit -> [Port]
 genPortsAST c =
-  (Port Input . frNode <$> inp) ++ (Port Output . frNode <$> out)
+  (port Input . frNode <$> inp) ++ (port Output . frNode <$> out)
   where
     inp = inputs graph
     out = outputs graph
     graph = getCircuit c
+    port x = Port (Just x) Nothing
 
 -- | Generates the nested expression AST, so that it can then generate the
 -- assignment expressions.
@@ -54,25 +55,25 @@ genAssignExpr g (n:ns) = OpExpr wire op <$> genAssignExpr g ns
 -- | Generate the continuous assignment AST for a particular node. If it does
 -- not have any nodes that link to it then return 'Nothing', as that means that
 -- the assignment will just be empty.
-genContAssignAST :: Circuit -> LNode Gate -> Maybe ContAssign
-genContAssignAST c (n, g) = ContAssign name <$> genAssignExpr g nodes
+genContAssignAST :: Circuit -> LNode Gate -> Maybe ModItem
+genContAssignAST c (n, g) = ModCA . ContAssign name <$> genAssignExpr g nodes
   where
     gr = getCircuit c
     nodes = G.pre gr n
     name = frNode n
 
-genAssignAST :: Circuit -> [ContAssign]
+genAssignAST :: Circuit -> [ModItem]
 genAssignAST c = catMaybes $ genContAssignAST c <$> nodes
   where
     gr = getCircuit c
     nodes = G.labNodes gr
 
-genModuleDeclAST :: Circuit -> ModuleDecl
-genModuleDeclAST c = ModuleDecl id ports items
+genModuleDeclAST :: Circuit -> ModDecl
+genModuleDeclAST c = ModDecl id ports items
   where
     id = Identifier "gen_module"
     ports = genPortsAST c
-    items = Assign <$> genAssignAST c
+    items = genAssignAST c
 
 generateAST :: Circuit -> SourceText
 generateAST c = SourceText [Description $ genModuleDeclAST c]