moved some expression traversal logic from PackedArray to Traverse

src/Convert/PackedArray.hs

@@ -29,7 +29,7 @@ convert = traverseDescriptions convertDescription
 convertDescription :: Description -> Description
 convertDescription description =
     hoistPortDecls $
-    traverseModuleItems (flattenModuleItem info . convertModuleItem dimMap') description
+    traverseModuleItems (flattenModuleItem info . rewriteModuleItem dimMap') description
     where
         info = execState
             (collectModuleItemsM collectDecl description)
@@ -152,53 +152,6 @@ simplify other = other
 prefix :: Identifier -> Identifier
 prefix ident = "_sv2v_" ++ ident
 
-
--- TODO FIXME XXX: There is a huge opportunity here to simplify the code after
--- this point in the module. Each of these mappings have a bit of their own
--- quirks. They cover all LHSs, expressions, and statements, at every level.
-
-
-rewriteRange :: DimMap -> Range -> Range
-rewriteRange dimMap (a, b) = (r a, r b)
-    where r = rewriteExpr dimMap
-
-rewriteIdentifier :: DimMap -> Identifier -> Identifier
-rewriteIdentifier dimMap x =
-    if Map.member x dimMap
-        then prefix x
-        else x
-
-rewriteExpr :: DimMap -> Expr -> Expr
-rewriteExpr dimMap = rewriteExpr'
-    where
-        ri :: Identifier -> Identifier
-        ri = rewriteIdentifier dimMap
-        re = rewriteExpr'
-        rewriteExpr' :: Expr -> Expr
-        rewriteExpr' (String     s) = String    s
-        rewriteExpr' (Number     s) = Number    s
-        rewriteExpr' (ConstBool  b) = ConstBool b
-        rewriteExpr' (Ident      i  ) = Ident      (ri i)
-        rewriteExpr' (IdentRange i (r @ (s, e))) =
-            case Map.lookup i dimMap of
-                Nothing -> IdentRange (ri i) (rewriteRange dimMap r)
-                Just (t, _) ->
-                    IdentRange i (simplify s', simplify e')
-                    where
-                        (a, b) = head $ snd $ typeDims t
-                        size = BinOp Add (BinOp Sub a b) (Number "1")
-                        s' = BinOp Sub (BinOp Mul size (BinOp Add s (Number "1"))) (Number "1")
-                        e' = BinOp Mul size e
-        rewriteExpr' (IdentBit   i e) = IdentBit   (ri i) (re e)
-        rewriteExpr' (Repeat     e l) = Repeat (re e) (map re l)
-        rewriteExpr' (Concat     l  ) = Concat (map re l)
-        rewriteExpr' (Call       f l) = Call f (map re l)
-        rewriteExpr' (UniOp      o e) = UniOp o (re e)
-        rewriteExpr' (BinOp      o e1 e2) = BinOp o (re e1) (re e2)
-        rewriteExpr' (Mux        e1 e2 e3) = Mux (re e1) (re e2) (re e3)
-        rewriteExpr' (Bit        e n) = Bit (re e) n
-        rewriteExpr' (Cast       t e) = Cast t (re e)
-
 -- combines (flattens) the bottom two ranges in the given list of ranges
 flattenRanges :: [Range] -> [Range]
 flattenRanges rs =
@@ -214,40 +167,37 @@ flattenRanges rs =
         r' = (simplify upper, e1)
         rs' = (tail $ tail rs) ++ [r']
 
-rewriteLHS :: DimMap -> LHS -> LHS
-rewriteLHS dimMap (LHS      x  ) = LHS      (rewriteIdentifier dimMap x)
-rewriteLHS dimMap (LHSBit   x e) = LHSBit   (rewriteIdentifier dimMap x) (rewriteExpr  dimMap e)
-rewriteLHS dimMap (LHSRange x r) = LHSRange (rewriteIdentifier dimMap x) (rewriteRange dimMap r)
-rewriteLHS dimMap (LHSConcat ls) = LHSConcat $ map (rewriteLHS dimMap) ls
-
-rewriteStmt :: DimMap -> Stmt -> Stmt
-rewriteStmt dimMap orig = rs orig
+rewriteModuleItem :: DimMap -> ModuleItem -> ModuleItem
+rewriteModuleItem dimMap =
+    traverseStmts rewriteStmt .
+    traverseExprs rewriteExpr
     where
-        rs :: Stmt -> Stmt
-        rs (Block decls stmts) = Block decls (map rs stmts)
-        rs (Case kw e cases def) = Case kw e' cases' def'
-            where
-                re :: Expr -> Expr
-                re = rewriteExpr dimMap
-                rc :: Case -> Case
-                rc (exprs, stmt) = (map re exprs, rs stmt)
-                e' = re e
-                cases' = map rc cases
-                def' = fmap rs def
-        rs (AsgnBlk lhs expr) = convertAssignment AsgnBlk lhs expr
-        rs (Asgn    lhs expr) = convertAssignment Asgn    lhs expr
-        rs (For (x1, e1) cc (x2, e2) stmt) = For (x1, e1') cc' (x2, e2') (rs stmt)
-            where
-                e1' = rewriteExpr dimMap e1
-                e2' = rewriteExpr dimMap e2
-                cc' = rewriteExpr dimMap cc
-        rs (If cc s1 s2) = If (rewriteExpr dimMap cc) (rs s1) (rs s2)
-        rs (Timing sense stmt) = Timing sense (rs stmt)
-        rs (Null) = Null
+        rewriteIdent :: Identifier -> Identifier
+        rewriteIdent x = if Map.member x dimMap then prefix x else x
+
+        rewriteExpr :: Expr -> Expr
+        rewriteExpr (Ident      i)   = Ident    (rewriteIdent i)
+        rewriteExpr (IdentBit   i e) = IdentBit (rewriteIdent i) e
+        rewriteExpr (IdentRange i (r @ (s, e))) =
+            case Map.lookup i dimMap of
+                Nothing -> IdentRange (rewriteIdent i) r
+                Just (t, _) ->
+                    IdentRange i (simplify s', simplify e')
+                    where
+                        (a, b) = head $ snd $ typeDims t
+                        size = BinOp Add (BinOp Sub a b) (Number "1")
+                        s' = BinOp Sub (BinOp Mul size (BinOp Add s (Number "1"))) (Number "1")
+                        e' = BinOp Mul size e
+        rewriteExpr other = other
+
+        rewriteStmt :: Stmt -> Stmt
+        rewriteStmt (AsgnBlk lhs expr) = convertAssignment AsgnBlk lhs expr
+        rewriteStmt (Asgn    lhs expr) = convertAssignment Asgn    lhs expr
+        rewriteStmt other = other
         convertAssignment :: (LHS -> Expr -> Stmt) -> LHS -> Expr -> Stmt
         convertAssignment constructor (lhs @ (LHS ident)) (expr @ (Repeat _ exprs)) =
             case Map.lookup ident dimMap of
-                Nothing -> constructor (rewriteLHS dimMap lhs) (rewriteExpr dimMap expr)
+                Nothing -> constructor lhs expr
                 Just (_, (a, b)) ->
                     For inir chkr incr assign
                     where
@@ -259,27 +209,4 @@ rewriteStmt dimMap orig = rs orig
                         chkr = BinOp Le (Ident index) a
                         incr = (index, BinOp Add (Ident index) (Number "1"))
         convertAssignment constructor lhs expr =
-            constructor (rewriteLHS dimMap lhs) (rewriteExpr dimMap expr)
-
-convertModuleItem :: DimMap -> ModuleItem -> ModuleItem
-convertModuleItem dimMap (MIDecl (Variable d t x a me)) =
-    MIDecl $ Variable d t  x a' me'
-    where
-        a' = map (rewriteRange dimMap) a
-        me' = fmap (rewriteExpr dimMap) me
-convertModuleItem dimMap (Assign lhs expr) =
-    Assign (rewriteLHS dimMap lhs) (rewriteExpr dimMap expr)
-convertModuleItem dimMap (AlwaysC kw stmt) =
-    AlwaysC kw (rewriteStmt dimMap stmt)
-convertModuleItem dimMap (Function ret f decls stmt) =
-    Function ret f decls (rewriteStmt dimMap stmt)
-convertModuleItem dimMap (Instance m params x ml) =
-    Instance m params x $ fmap (map convertPortBinding) ml
-    where
-        convertPortBinding :: PortBinding -> PortBinding
-        convertPortBinding (p, Nothing) = (p, Nothing)
-        convertPortBinding (p, Just  e) = (p, Just $ rewriteExpr dimMap e)
-convertModuleItem _ (Comment  x) = Comment  x
-convertModuleItem _ (Genvar   x) = Genvar   x
-convertModuleItem _ (MIDecl   x) = MIDecl   x
-convertModuleItem _ (Generate x) = Generate x
+            constructor lhs expr

src/Convert/Traverse.hs

@@ -21,8 +21,12 @@ module Convert.Traverse
 , traverseStmtLHSsM
 , traverseStmtLHSs
 , collectStmtLHSsM
+, traverseExprsM
+, traverseExprs
+, collectExprsM
 ) where
 
+import Data.Maybe (fromJust)
 import Control.Monad.State
 import Language.SystemVerilog.AST
 
@@ -135,3 +139,123 @@ traverseStmtLHSs :: Mapper LHS -> Mapper Stmt
 traverseStmtLHSs = unmonad traverseStmtLHSsM
 collectStmtLHSsM :: Monad m => CollectorM m LHS -> CollectorM m Stmt
 collectStmtLHSsM = collectify traverseStmtLHSsM
+
+traverseNestedExprsM :: Monad m => MapperM m Expr -> MapperM m Expr
+traverseNestedExprsM mapper = exprMapper
+    where
+        exprMapper e = mapper e >>= em
+        em (String     s) = return $ String    s
+        em (Number     s) = return $ Number    s
+        em (ConstBool  b) = return $ ConstBool b
+        em (Ident      i) = return $ Ident     i
+        em (IdentRange i (e1, e2)) = do
+            e1' <- exprMapper e1
+            e2' <- exprMapper e2
+            return $ IdentRange i (e1', e2')
+        em (IdentBit   i e) =
+            exprMapper e >>= return . IdentBit i
+        em (Repeat     e l) = do
+            e' <- exprMapper e
+            l' <- mapM exprMapper l
+            return $ Repeat e' l'
+        em (Concat     l) =
+            mapM exprMapper l >>= return . Concat
+        em (Call       f l) =
+            mapM exprMapper l >>= return . Call f
+        em (UniOp      o e) =
+            exprMapper e >>= return . UniOp o
+        em (BinOp      o e1 e2) = do
+            e1' <- exprMapper e1
+            e2' <- exprMapper e2
+            return $ BinOp o e1' e2'
+        em (Mux        e1 e2 e3) = do
+            e1' <- exprMapper e1
+            e2' <- exprMapper e2
+            e3' <- exprMapper e3
+            return $ Mux e1' e2' e3'
+        em (Bit        e n) =
+            exprMapper e >>= \e' -> return $ Bit e' n
+        em (Cast       t e) =
+            exprMapper e >>= return . Cast t
+
+
+traverseExprsM :: Monad m => MapperM m Expr -> MapperM m ModuleItem
+traverseExprsM mapper = moduleItemMapper
+    where
+
+    rangeMapper (a, b) = do
+        a' <- exprMapper a
+        b' <- exprMapper b
+        return (a', b')
+
+    maybeExprMapper Nothing = return Nothing
+    maybeExprMapper (Just e) =
+        exprMapper e >>= return . Just
+
+    declMapper (Parameter  t x e) =
+        exprMapper e >>= return . Parameter  t x
+    declMapper (Localparam t x e) =
+        exprMapper e >>= return . Localparam t x
+    declMapper (Variable d t x a me) = do
+        a' <- mapM rangeMapper a
+        me' <- maybeExprMapper me
+        return $ Variable d t x a' me'
+
+    exprMapper = traverseNestedExprsM mapper
+
+    caseMapper (exprs, stmt) = do
+        exprs' <- mapM exprMapper exprs
+        return (exprs', stmt)
+    stmtMapper = traverseNestedStmtsM flatStmtMapper
+    flatStmtMapper (Block header stmts) = do
+        if header == Nothing
+            then return $ Block Nothing stmts
+            else do
+                let Just (name, decls) = header
+                decls' <- mapM declMapper decls
+                return $ Block (Just (name, decls')) stmts
+    flatStmtMapper (Case kw e cases def) = do
+        e' <- exprMapper e
+        cases' <- mapM caseMapper cases
+        return $ Case kw e' cases' def
+    flatStmtMapper (AsgnBlk lhs expr) =
+        exprMapper expr >>= return . AsgnBlk lhs
+    flatStmtMapper (Asgn    lhs expr) =
+        exprMapper expr >>= return . Asgn    lhs
+    flatStmtMapper (For (x1, e1) cc (x2, e2) stmt) = do
+        e1' <- exprMapper e1
+        e2' <- exprMapper e2
+        cc' <- exprMapper cc
+        return $ For (x1, e1') cc' (x2, e2') stmt
+    flatStmtMapper (If cc s1 s2) =
+        exprMapper cc >>= \cc' -> return $ If cc' s1 s2
+    flatStmtMapper (Timing sense stmt) = return $ Timing sense stmt
+    flatStmtMapper (Null) = return Null
+
+    portBindingMapper (p, me) =
+        maybeExprMapper me >>= \me' -> return (p, me')
+
+    moduleItemMapper (MIDecl decl) =
+        declMapper decl >>= return . MIDecl
+    moduleItemMapper (Assign lhs expr) =
+        exprMapper expr >>= return . Assign lhs
+    moduleItemMapper (AlwaysC kw stmt) =
+        stmtMapper stmt >>= return . AlwaysC kw
+    moduleItemMapper (Function ret f decls stmt) = do
+        decls' <- mapM declMapper decls
+        stmt' <- stmtMapper stmt
+        return $ Function ret f decls' stmt'
+    moduleItemMapper (Instance m params x ml) = do
+        if ml == Nothing
+            then return $ Instance m params x Nothing
+            else do
+                l <- mapM portBindingMapper (fromJust ml)
+                return $ Instance m params x (Just l)
+    moduleItemMapper (Comment  x) = return $ Comment  x
+    moduleItemMapper (Genvar   x) = return $ Genvar   x
+    moduleItemMapper (Generate x) = return $ Generate x
+
+traverseExprs :: Mapper Expr -> Mapper ModuleItem
+traverseExprs = unmonad traverseExprsM
+collectExprsM :: Monad m => CollectorM m Expr -> CollectorM m ModuleItem
+collectExprsM = collectify traverseExprsM