Implement module item reduction properly

3 files changed, 18 insertions, 6 deletions

diff --git a/src/VeriFuzz/Reduce.hs b/src/VeriFuzz/Reduce.hs
index 2865898..0de97ae 100644
--- a/src/VeriFuzz/Reduce.hs
+++ b/src/VeriFuzz/Reduce.hs
@@ -27,7 +27,7 @@ where
 
 import           Control.Lens
 import           Data.List                (nub)
-import           Data.Maybe               (catMaybes)
+import           Data.Maybe               (mapMaybe)
 import           Data.Text                (Text)
 import           VeriFuzz.Verilog.AST
 import           VeriFuzz.Verilog.CodeGen
@@ -208,7 +208,7 @@ halveAssigns = combine mainModule halveModAssign
 -- | Checks if a module item is needed in the module declaration.
 relevantModItem :: ModDecl -> ModItem -> Bool
 relevantModItem (ModDecl _ out _ _ _) (ModCA (ContAssign i _)) = i `elem` fmap _portName out
-relevantModItem _ (Decl _ _ _)                          = True
+relevantModItem _ Decl{}                          = True
 relevantModItem _ _                                            = False
 
 isAssign :: Statement -> Bool
@@ -220,7 +220,7 @@ lValName :: LVal -> [Identifier]
 lValName (RegId i) = [i]
 lValName (RegExpr i _) = [i]
 lValName (RegSize i _) = [i]
-lValName (RegConcat e) = catMaybes . fmap getId . concat $ universe <$> e
+lValName (RegConcat e) = mapMaybe getId . concat $ universe <$> e
     where
         getId (Id i) = Just i
         getId _      = Nothing
@@ -235,15 +235,20 @@ findActiveWires :: ModDecl -> [Identifier]
 findActiveWires m@(ModDecl _ i o _ p) = nub $ assignWires <> assignStat <> fmap portToId i <> fmap portToId o <> fmap paramToId p
     where
         assignWires = m ^.. modItems . traverse . modContAssign . contAssignNetLVal
-        assignStat = concat . fmap lValName $ (allStat ^.. traverse . stmntBA . assignReg)
+        assignStat = concatMap lValName $ (allStat ^.. traverse . stmntBA . assignReg)
                      <> (allStat ^.. traverse . stmntNBA . assignReg)
         allStat = filter isAssign . concat $ fmap universe stat
         stat = (m ^.. modItems . traverse . _Initial) <> (m ^.. modItems . traverse . _Always)
 
+cleanSourceInfo :: SourceInfo -> SourceInfo
+cleanSourceInfo src = clean active src
+    where
+        active = findActiveWires (src ^. mainModule)
+
 -- | Reducer for module items. It does a binary search on all the module items,
 -- except assignments to outputs and input-output declarations.
 halveModItems :: Replace SourceInfo
-halveModItems srcInfo = fmap addRelevant src
+halveModItems srcInfo = cleanSourceInfo . addRelevant <$> src
     where
         repl = halve . filter (not . relevantModItem main)
         relevant = filter (relevantModItem main) $ main ^. modItems
diff --git a/src/VeriFuzz/Verilog/CodeGen.hs b/src/VeriFuzz/Verilog/CodeGen.hs
index efacd3c..71ba162 100644
--- a/src/VeriFuzz/Verilog/CodeGen.hs
+++ b/src/VeriFuzz/Verilog/CodeGen.hs
@@ -11,7 +11,8 @@ This module generates the code from the Verilog AST defined in
 "VeriFuzz.Verilog.AST".
 -}
 
-{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE FlexibleInstances  #-}
 
 module VeriFuzz.Verilog.CodeGen
     ( -- * Code Generation
@@ -21,6 +22,7 @@ module VeriFuzz.Verilog.CodeGen
     )
 where
 
+import           Data.Data                 (Data)
 import           Data.List.NonEmpty        (NonEmpty (..), toList)
 import           Data.Text                 (Text)
 import qualified Data.Text                 as T
@@ -318,6 +320,7 @@ instance Source SourceInfo where
     genSource (SourceInfo _ src) = genSource src
 
 newtype GenVerilog a = GenVerilog { unGenVerilog :: a }
+                     deriving (Eq, Ord, Data)
 
 instance (Source a) => Show (GenVerilog a) where
     show = T.unpack . genSource . unGenVerilog
diff --git a/src/VeriFuzz/Verilog/Mutate.hs b/src/VeriFuzz/Verilog/Mutate.hs
index 3f0ae83..66f3c37 100644
--- a/src/VeriFuzz/Verilog/Mutate.hs
+++ b/src/VeriFuzz/Verilog/Mutate.hs
@@ -49,6 +49,7 @@ import           VeriFuzz.Circuit.Internal
 import           VeriFuzz.Internal
 import           VeriFuzz.Verilog.AST
 import           VeriFuzz.Verilog.BitVec
+import           VeriFuzz.Verilog.CodeGen
 import           VeriFuzz.Verilog.Internal
 
 class Mutate a where
@@ -146,6 +147,9 @@ instance Mutate a => Mutate [a] where
 instance Mutate a => Mutate (Maybe a) where
     mutExpr f a = mutExpr f <$> a
 
+instance Mutate a => Mutate (GenVerilog a) where
+    mutExpr f (GenVerilog a) = GenVerilog $ mutExpr f a
+
 -- | Return if the 'Identifier' is in a 'ModDecl'.
 inPort :: Identifier -> ModDecl -> Bool
 inPort i m = inInput