Add lens to access main module in SourceInfo

3 files changed, 23 insertions, 7 deletions

diff --git a/src/VeriFuzz/Icarus.hs b/src/VeriFuzz/Icarus.hs
index 0a807dc..b709967 100644
--- a/src/VeriFuzz/Icarus.hs
+++ b/src/VeriFuzz/Icarus.hs
@@ -12,6 +12,7 @@ Icarus verilog module.
 
 module VeriFuzz.Icarus where
 
+import           Control.Lens
 import           Crypto.Hash            (Digest, hash)
 import           Crypto.Hash.Algorithms (SHA256)
 import           Data.Binary            (encode)
@@ -91,7 +92,7 @@ runSimIcarus sim rinfo bss = do
     let modWithTb = VerilogSrc $ Description <$> [newtb, m]
     writefile "main.v" $ genSource modWithTb
     runSimWithFile sim "main.v" bss
-    where m = mainModule rinfo
+    where m = rinfo ^. mainModule
 
 runSimIcarusWithFile :: Icarus -> FilePath -> [ByteString] -> Sh ByteString
 runSimIcarusWithFile sim f _ = do
diff --git a/src/VeriFuzz/Internal/Simulator.hs b/src/VeriFuzz/Internal/Simulator.hs
index 5712709..1454a0f 100644
--- a/src/VeriFuzz/Internal/Simulator.hs
+++ b/src/VeriFuzz/Internal/Simulator.hs
@@ -12,7 +12,7 @@ Class of the simulator and the synthesize tool.
 
 module VeriFuzz.Internal.Simulator where
 
-import           Control.Lens          ((^.), (^..))
+import           Control.Lens
 import           Data.Bits             (shiftL)
 import           Data.ByteString       (ByteString)
 import qualified Data.ByteString       as B
@@ -48,10 +48,19 @@ class (Tool a) => Synthesisor a where
 data SourceInfo = SourceInfo { runMainModule :: {-# UNPACK #-} !Text
                              , runSource     :: !VerilogSrc
                              }
+                  deriving (Eq, Show)
 
-mainModule :: SourceInfo -> ModDecl
-mainModule (SourceInfo main src) = head . filter ismain $ src ^.. getModule
-    where ismain v = v ^. modId . getIdentifier == main
+-- | 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
 
 rootPath :: Sh FilePath
 rootPath = do
diff --git a/src/VeriFuzz/Yosys.hs b/src/VeriFuzz/Yosys.hs
index d33e399..9f605db 100644
--- a/src/VeriFuzz/Yosys.hs
+++ b/src/VeriFuzz/Yosys.hs
@@ -14,6 +14,7 @@ Yosys simulator implementation.
 
 module VeriFuzz.Yosys where
 
+import           Control.Lens
 import           Prelude               hiding (FilePath)
 import           Shelly
 import           Text.Shakespeare.Text (st)
@@ -72,7 +73,7 @@ runEquivYosys
     -> SourceInfo
     -> Sh ()
 runEquivYosys yosys sim1 sim2 srcInfo = do
-    writefile "top.v" . genSource . initMod . makeTop 2 $ mainModule srcInfo
+    writefile "top.v" . genSource . initMod . makeTop 2 $ srcInfo ^. mainModule
     writefile checkFile $ yosysSatConfig sim1 sim2 srcInfo
     runSynth sim1 srcInfo $ fromText [st|syn_#{toText sim1}.v|]
     runMaybeSynth sim2 srcInfo
@@ -94,7 +95,12 @@ runEquiv _ sim1 sim2 srcInfo = do
     root <- rootPath
     dir  <- pwd
     echoP "SymbiYosys: setup"
-    writefile "top.v" . genSource . initMod . makeTopAssert $ mainModule srcInfo
+    writefile "top.v"
+        .  genSource
+        .  initMod
+        .  makeTopAssert
+        $  srcInfo
+        ^. mainModule
     writefile "test.sby" $ sbyConfig root sim1 sim2 srcInfo
     runSynth sim1 srcInfo $ fromText [st|syn_#{toText sim1}.v|]
     runMaybeSynth sim2 srcInfo