Refactor blacklist code to share common logic

include/verilogAST/assign_inliner.hpp

@@ -61,48 +61,51 @@ class WireReadCounter : public Transformer {
   virtual std::unique_ptr<Declaration> visit(std::unique_ptr<Declaration> node);
 };
 
-class SliceBlacklister : public Transformer {
+class Blacklister : public Transformer {
+  std::set<std::string> &wire_blacklist;
+  std::map<std::string, std::unique_ptr<Expression>> &assign_map;
+
+ protected:
+  bool blacklist = false;
+
+ public:
+  Blacklister(std::set<std::string> &wire_blacklist,
+              std::map<std::string, std::unique_ptr<Expression>> &assign_map)
+      : wire_blacklist(wire_blacklist), assign_map(assign_map){};
+  using Transformer::visit;
+  virtual std::unique_ptr<Identifier> visit(std::unique_ptr<Identifier> node);
+};
+
+class SliceBlacklister : public Blacklister {
   // Prevent inling wires into slice nodes, e.g.
   // wire [7:0] x;
   // assign x = y + z;
   // assign w = x[4:0];
   //
   // Verilog does not support (y + z)[4:0]
-  std::set<std::string> &wire_blacklist;
-  std::map<std::string, std::unique_ptr<Expression>> &assign_map;
-  bool inside_slice = false;
-
  public:
   SliceBlacklister(
       std::set<std::string> &wire_blacklist,
       std::map<std::string, std::unique_ptr<Expression>> &assign_map)
-      : wire_blacklist(wire_blacklist), assign_map(assign_map){};
-
-  using Transformer::visit;
+      : Blacklister(wire_blacklist, assign_map){};
+  using Blacklister::visit;
   virtual std::unique_ptr<Slice> visit(std::unique_ptr<Slice> node);
-  virtual std::unique_ptr<Identifier> visit(std::unique_ptr<Identifier> node);
 };
 
-class IndexBlacklister : public Transformer {
+class IndexBlacklister : public Blacklister {
   // Prevent inling wires into index nodes, e.g.
   // wire x;
   // assign x = y + z;
   // assign w = x[0];
   //
   // Verilog does not support (y + z)[0]
-  std::set<std::string> &wire_blacklist;
-  std::map<std::string, std::unique_ptr<Expression>> &assign_map;
-  bool inside_index = false;
-
  public:
   IndexBlacklister(
       std::set<std::string> &wire_blacklist,
       std::map<std::string, std::unique_ptr<Expression>> &assign_map)
-      : wire_blacklist(wire_blacklist), assign_map(assign_map){};
-
-  using Transformer::visit;
+      : Blacklister(wire_blacklist, assign_map){};
+  using Blacklister::visit;
   virtual std::unique_ptr<Index> visit(std::unique_ptr<Index> node);
-  virtual std::unique_ptr<Identifier> visit(std::unique_ptr<Identifier> node);
 };
 
 class AssignInliner : public Transformer {

src/assign_inliner.cpp

@@ -3,18 +3,9 @@
 
 namespace verilogAST {
 
-std::unique_ptr<Slice> SliceBlacklister::visit(std::unique_ptr<Slice> node) {
-  bool prev = this->inside_slice;
-  this->inside_slice = true;
-  node = Transformer::visit(std::move(node));
-  // Restore prev value, since we could be nested inside an slice
-  this->inside_slice = prev;
-  return node;
-}
-
-std::unique_ptr<Identifier> SliceBlacklister::visit(
+std::unique_ptr<Identifier> Blacklister::visit(
     std::unique_ptr<Identifier> node) {
-  if (this->inside_slice) {
+  if (this->blacklist) {
     auto it = assign_map.find(node->toString());
     bool assigned_to_id =
         it != assign_map.end() && dynamic_cast<Identifier*>(it->second.get());
@@ -25,25 +16,21 @@ std::unique_ptr<Identifier> SliceBlacklister::visit(
   return node;
 }
 
-std::unique_ptr<Index> IndexBlacklister::visit(std::unique_ptr<Index> node) {
-  bool prev = this->inside_index;
-  this->inside_index = true;
+std::unique_ptr<Slice> SliceBlacklister::visit(std::unique_ptr<Slice> node) {
+  bool prev = this->blacklist;
+  this->blacklist = true;
   node = Transformer::visit(std::move(node));
-  // Restore prev value, since we could be nested inside an index
-  this->inside_index = prev;
+  // Restore prev value, since we could be nested inside an slice
+  this->blacklist = prev;
   return node;
 }
 
-std::unique_ptr<Identifier> IndexBlacklister::visit(
-    std::unique_ptr<Identifier> node) {
-  if (this->inside_index) {
-    auto it = assign_map.find(node->toString());
-    bool assigned_to_id =
-        it != assign_map.end() && dynamic_cast<Identifier*>(it->second.get());
-    if (!assigned_to_id) {
-      this->wire_blacklist.insert(node->value);
-    }
-  };
+std::unique_ptr<Index> IndexBlacklister::visit(std::unique_ptr<Index> node) {
+  bool prev = this->blacklist;
+  this->blacklist = true;
+  node = Transformer::visit(std::move(node));
+  // Restore prev value, since we could be nested inside an index
+  this->blacklist = prev;
   return node;
 }