Add nesting to the SourceText

2 files changed, 19 insertions, 2 deletions

diff --git a/app/Main.hs b/app/Main.hs
index bb2697a..cbe570b 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE OverloadedStrings #-}
+
 module Main where
 
 import qualified Data.Graph.Inductive              as G
@@ -18,4 +20,4 @@ main = do
 --  _ <- runGraphviz (graphToDot quickParams $ emap (const "") gr) Png "output.png",
 --  T.putStrLn $ generate gr
   --g <- QC.generate (QC.arbitrary :: QC.Gen SourceText)
-  render . genSourceText . generateAST $ Circuit gr
+  render . genSourceText . nestSource (Identifier "w96") . generateAST $ Circuit gr
diff --git a/src/Test/VeriFuzz/Mutate.hs b/src/Test/VeriFuzz/Mutate.hs
index dd391fb..3e7acae 100644
--- a/src/Test/VeriFuzz/Mutate.hs
+++ b/src/Test/VeriFuzz/Mutate.hs
@@ -11,6 +11,8 @@ Functions to mutate the Verilog AST from "Test.VeriFuzz.VerilogAST" to generate
 more random patterns, such as nesting wires instead of creating new ones.
 -}
 
+{-# LANGUAGE OverloadedStrings #-}
+
 module Test.VeriFuzz.Mutate where
 
 import           Control.Lens
@@ -32,9 +34,22 @@ findAssign id items =
       | ca ^. contAssignNetLVal == id = Just $ ca ^. contAssignExpr
       | otherwise = Nothing
 
+idTrans :: Identifier -> Expression -> Expression -> Expression
+idTrans i expr (PrimExpr (PrimId id))
+  | id == i = expr
+  | otherwise = (PrimExpr (PrimId id))
+idTrans _ _ e = e
+
 -- | Nest expressions for a specific 'Identifier'. If the 'Identifier' is not found,
 -- the AST is not changed.
 nestId :: Identifier -> ModuleDecl -> ModuleDecl
 nestId id mod
-  | not $ inPort id mod = mod
+  | not $ inPort id mod = mod & get %~ trans
   | otherwise = mod
+  where
+    get = moduleItems . traverse . _Assign . contAssignExpr
+    trans = transformOf traverseExpr . idTrans id . PrimExpr . PrimId $ Identifier "RANDOM"
+
+nestSource :: Identifier -> SourceText -> SourceText
+nestSource id src =
+  src & getSourceText . traverse . getDescription %~ nestId id