Move declaration of SourceInfo

3 files changed, 27 insertions, 8 deletions

diff --git a/src/VeriFuzz/Verilog/AST.hs b/src/VeriFuzz/Verilog/AST.hs
index 3a4d2b9..0ef9057 100644
--- a/src/VeriFuzz/Verilog/AST.hs
+++ b/src/VeriFuzz/Verilog/AST.hs
@@ -22,7 +22,8 @@ Defines the types to build a Verilog AST.
 
 module VeriFuzz.Verilog.AST
     ( -- * Top level types
-      Verilog(..)
+      SourceInfo(..)
+    , Verilog(..)
     -- * Primitives
     -- ** Identifier
     , Identifier(..)
@@ -127,6 +128,7 @@ module VeriFuzz.Verilog.AST
     -- * Useful Lenses and Traversals
     , getModule
     , getSourceId
+    , mainModule
     )
 where
 
@@ -452,6 +454,11 @@ traverseModItem _ e = pure e
 newtype Verilog = Verilog { _getVerilog :: [ModDecl] }
                    deriving (Eq, Show, Ord, Data, Semigroup, Monoid)
 
+data SourceInfo = SourceInfo { runMainModule :: {-# UNPACK #-} !Text
+                             , runSource     :: !Verilog
+                             }
+                  deriving (Eq, Show)
+
 $(makeLenses ''Expr)
 $(makeLenses ''ConstExpr)
 $(makeLenses ''Task)
@@ -480,3 +487,15 @@ getModule = _Wrapped . traverse
 getSourceId :: Traversal' Verilog Text
 getSourceId = getModule . modId . _Wrapped
 {-# INLINE getSourceId #-}
+
+-- | May need to change this to Traversal to be safe. For now it will fail when
+-- the main has not been properly set with.
+mainModule :: Lens' SourceInfo ModDecl
+mainModule = lens get_ set_
+  where
+    set_ (SourceInfo top main) v =
+        SourceInfo top (main & getModule %~ update top v)
+    update top v m@(ModDecl (Identifier i) _ _ _ _) | i == top  = v
+                                                    | otherwise = m
+    get_ (SourceInfo top main) = head . filter (f top) $ main ^.. getModule
+    f top (ModDecl (Identifier i) _ _ _ _) = i == top
diff --git a/src/VeriFuzz/Verilog/CodeGen.hs b/src/VeriFuzz/Verilog/CodeGen.hs
index 2030bc7..3b9c7ad 100644
--- a/src/VeriFuzz/Verilog/CodeGen.hs
+++ b/src/VeriFuzz/Verilog/CodeGen.hs
@@ -28,7 +28,6 @@ import qualified Data.Text               as T
 import qualified Data.Text.IO            as T
 import           Numeric                 (showHex)
 import           VeriFuzz.Internal
-import           VeriFuzz.Sim.Internal
 import           VeriFuzz.Verilog.AST
 import           VeriFuzz.Verilog.BitVec
 
diff --git a/src/VeriFuzz/Verilog/Gen.hs b/src/VeriFuzz/Verilog/Gen.hs
index a6ebbd9..46cdc25 100644
--- a/src/VeriFuzz/Verilog/Gen.hs
+++ b/src/VeriFuzz/Verilog/Gen.hs
@@ -436,10 +436,11 @@ moduleDef top = do
 
 -- | Procedural generation method for random Verilog. Uses internal 'Reader' and
 -- 'State' to keep track of the current Verilog code structure.
-procedural :: Config -> Gen Verilog
-procedural config = do
-    (mainMod, st) <- Hog.resize num
-        $ runReaderT (runStateT (moduleDef (Just "top")) context) config
+procedural :: T.Text -> Config -> Gen Verilog
+procedural top config = do
+    (mainMod, st) <- Hog.resize num $ runReaderT
+        (runStateT (moduleDef (Just $ Identifier top)) context)
+        config
     return . Verilog $ mainMod : st ^. modules
   where
     context =
@@ -447,5 +448,5 @@ procedural config = do
     num = fromIntegral $ confProp propSize
     confProp i = config ^. configProperty . i
 
-proceduralIO :: Config -> IO Verilog
-proceduralIO = Hog.sample . procedural
+proceduralIO :: T.Text -> Config -> IO Verilog
+proceduralIO t = Hog.sample . procedural t