Fix handling of multiple assignments in inlining

include/verilogAST.hpp

@@ -807,6 +807,13 @@ class If : public BehavioralStatement {
         else_ifs(std::move(else_ifs)),
         else_body(std::move(else_body)){};
 
+  If(std::unique_ptr<Expression> cond,
+     std::vector<std::unique_ptr<BehavioralStatement>> true_body,
+     std::vector<std::unique_ptr<BehavioralStatement>> else_body)
+      : cond(std::move(cond)),
+        true_body(std::move(true_body)),
+        else_body(std::move(else_body)){};
+
   std::string toString();
   ~If(){};
 };

include/verilogAST/assign_inliner.hpp

@@ -67,6 +67,9 @@ class Blacklister : public Transformer {
 
  protected:
   bool blacklist = false;
+  // Allow numeric literals as valid drivers (okay for module instances, not
+  // okay for slice/index)
+  virtual bool allowNumDriver() { return false; };
   void blacklist_invalid_driver(std::unique_ptr<Identifier> node);
 
  public:
@@ -124,14 +127,17 @@ class ModuleInstanceBlacklister : public Blacklister {
   // We can make this configurable, but for now we keep it as the default since
   // some tools do not support general expressions inside module instance
   // statements
+ protected:
+  bool allowNumDriver() override { return true; };
+
  public:
   ModuleInstanceBlacklister(
       std::set<std::string> &wire_blacklist,
       std::map<std::string, std::unique_ptr<Expression>> &assign_map)
       : Blacklister(wire_blacklist, assign_map){};
   using Blacklister::visit;
   virtual std::unique_ptr<ModuleInstantiation> visit(
-      std::unique_ptr<ModuleInstantiation> node);
+      std::unique_ptr<ModuleInstantiation> node) override;
 };
 
 class AssignInliner : public Transformer {

src/assign_inliner.cpp

@@ -4,17 +4,22 @@
 namespace verilogAST {
 
 void Blacklister::blacklist_invalid_driver(std::unique_ptr<Identifier> node) {
+  if (this->wire_blacklist.count(node->value)) {
+    // Already blacklisted
+    return;
+  }
   if (!assign_map.count(node->toString())) {
     // Not in assign map, means it's a module input, don't need to do anything
     // because it won't be inlined
     return;
   }
   auto driver = assign_map[node->toString()]->clone();
   // Can only inline if driven by identifier, index, or slice
-  bool valid_driver = dynamic_cast<Identifier*>(driver.get()) ||
-                      dynamic_cast<Index*>(driver.get()) ||
-                      dynamic_cast<Slice*>(driver.get()) ||
-                      dynamic_cast<NumericLiteral*>(driver.get());
+  bool valid_driver =
+      dynamic_cast<Identifier*>(driver.get()) ||
+      dynamic_cast<Index*>(driver.get()) ||
+      dynamic_cast<Slice*>(driver.get()) ||
+      (this->allowNumDriver() && dynamic_cast<NumericLiteral*>(driver.get()));
   if (!valid_driver) {
     this->wire_blacklist.insert(node->value);
   } else if (auto ptr = dynamic_cast<Identifier*>(driver.get())) {
@@ -136,7 +141,7 @@ bool AssignInliner::can_inline(std::string key) {
     return false;
   }
   auto it = assign_map.find(key);
-  return it != assign_map.end() && (this->assign_count[key] == 1) &&
+  return it != assign_map.end() &&
          (this->read_count[key] == 1 ||
           dynamic_cast<Identifier*>(it->second.get()) ||
           dynamic_cast<NumericLiteral*>(it->second.get()));
@@ -259,6 +264,13 @@ std::unique_ptr<Module> AssignInliner::visit(std::unique_ptr<Module> node) {
                            this->non_input_ports, this->output_ports,
                            this->input_ports);
   node = builder.visit(std::move(node));
+  for (auto entry : assign_count) {
+    if (entry.second > 1) {
+      // Do not inline things assigned more than once, e.g. a reg inside
+      // if/else statements
+      this->wire_blacklist.insert(entry.first);
+    }
+  }
 
   WireReadCounter counter(this->read_count);
   node = counter.visit(std::move(node));

tests/assign_inliner.cpp

@@ -961,6 +961,172 @@ TEST(InlineAssignTests, TestInlineSliceOfIndex) {
   EXPECT_EQ(transformer.visit(std::move(module))->toString(), expected_str);
 }
 
+TEST(InlineAssignTests, TestInlineMultipleAssign) {
+  std::vector<std::unique_ptr<vAST::AbstractPort>> ports;
+  ports.push_back(std::make_unique<vAST::Port>(
+      std::make_unique<vAST::Vector>(vAST::make_id("i0"), vAST::make_num("4"),
+                                     vAST::make_num("0")),
+      vAST::INPUT, vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(
+      std::make_unique<vAST::Vector>(vAST::make_id("i1"), vAST::make_num("4"),
+                                     vAST::make_num("0")),
+      vAST::INPUT, vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("s"), vAST::INPUT,
+                                               vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(
+      std::make_unique<vAST::Vector>(vAST::make_id("o"), vAST::make_num("3"),
+                                     vAST::make_num("0")),
+      vAST::OUTPUT, vAST::WIRE));
+
+  std::vector<std::variant<std::unique_ptr<vAST::StructuralStatement>,
+                           std::unique_ptr<vAST::Declaration>>>
+      body;
+
+  body.push_back(std::make_unique<vAST::Wire>(std::make_unique<vAST::Vector>(
+      vAST::make_id("x"), vAST::make_num("4"), vAST::make_num("0"))));
+
+  body.push_back(std::make_unique<vAST::Wire>(std::make_unique<vAST::Vector>(
+      vAST::make_id("y"), vAST::make_num("4"), vAST::make_num("0"))));
+
+  body.push_back(std::make_unique<vAST::Wire>(std::make_unique<vAST::Vector>(
+      vAST::make_id("z"), vAST::make_num("4"), vAST::make_num("0"))));
+
+  body.push_back(std::make_unique<vAST::ContinuousAssign>(vAST::make_id("x"),
+                                                          vAST::make_id("i0")));
+
+  body.push_back(std::make_unique<vAST::ContinuousAssign>(vAST::make_id("y"),
+                                                          vAST::make_id("i1")));
+
+  std::vector<std::variant<
+      std::unique_ptr<vAST::Identifier>, std::unique_ptr<vAST::PosEdge>,
+      std::unique_ptr<vAST::NegEdge>, std::unique_ptr<vAST::Star>>>
+      sensitivity_list;
+  sensitivity_list.push_back(std::make_unique<vAST::Star>());
+
+  std::vector<std::unique_ptr<vAST::BehavioralStatement>> always_body;
+
+  std::vector<std::unique_ptr<vAST::BehavioralStatement>> true_body;
+  true_body.push_back(std::make_unique<vAST::BlockingAssign>(
+      std::make_unique<vAST::Identifier>("z"),
+      std::make_unique<vAST::Identifier>("x")));
+
+  std::vector<std::unique_ptr<vAST::BehavioralStatement>> else_body;
+  else_body.push_back(std::make_unique<vAST::BlockingAssign>(
+      std::make_unique<vAST::Identifier>("z"),
+      std::make_unique<vAST::Identifier>("y")));
+
+  always_body.push_back(std::make_unique<vAST::If>(
+      vAST::make_id("s"), std::move(true_body), std::move(else_body)));
+
+  body.push_back(std::make_unique<vAST::Always>(std::move(sensitivity_list),
+                                                std::move(always_body)));
+
+  body.push_back(std::make_unique<vAST::ContinuousAssign>(
+      vAST::make_id("o"),
+      std::make_unique<vAST::Slice>(vAST::make_id("z"), vAST::make_num("3"),
+                                    vAST::make_num("0"))));
+
+  std::unique_ptr<vAST::AbstractModule> module = std::make_unique<vAST::Module>(
+      "test_module", std::move(ports), std::move(body));
+
+  std::string raw_str =
+      "module test_module (\n"
+      "    input [4:0] i0,\n"
+      "    input [4:0] i1,\n"
+      "    input s,\n"
+      "    output [3:0] o\n"
+      ");\n"
+      "wire [4:0] x;\n"
+      "wire [4:0] y;\n"
+      "wire [4:0] z;\n"
+      "assign x = i0;\n"
+      "assign y = i1;\n"
+      "always @(*) begin\n"
+      "if (s) begin\n"
+      "    z = x;\n"
+      "end else begin\n"
+      "    z = y;\n"
+      "end\n"
+      "end\n"
+      "\n"
+      "assign o = z[3:0];\n"
+      "endmodule\n";
+
+  EXPECT_EQ(module->toString(), raw_str);
+
+  std::string expected_str =
+      "module test_module (\n"
+      "    input [4:0] i0,\n"
+      "    input [4:0] i1,\n"
+      "    input s,\n"
+      "    output [3:0] o\n"
+      ");\n"
+      "wire [4:0] z;\n"
+      "always @(*) begin\n"
+      "if (s) begin\n"
+      "    z = i0;\n"
+      "end else begin\n"
+      "    z = i1;\n"
+      "end\n"
+      "end\n"
+      "\n"
+      "assign o = z[3:0];\n"
+      "endmodule\n";
+
+  vAST::AssignInliner transformer;
+  EXPECT_EQ(transformer.visit(std::move(module))->toString(), expected_str);
+}
+
+TEST(InlineAssignTests, TestNoInlineNumToIndexSlice) {
+  std::vector<std::unique_ptr<vAST::AbstractPort>> ports;
+  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("o0"),
+                                               vAST::OUTPUT, vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(
+      std::make_unique<vAST::Vector>(vAST::make_id("o1"), vAST::make_num("1"),
+                                     vAST::make_num("0")),
+      vAST::OUTPUT, vAST::WIRE));
+
+  std::vector<std::variant<std::unique_ptr<vAST::StructuralStatement>,
+                           std::unique_ptr<vAST::Declaration>>>
+      body;
+
+  body.push_back(std::make_unique<vAST::Wire>(std::make_unique<vAST::Vector>(
+      vAST::make_id("x"), vAST::make_num("2"), vAST::make_num("0"))));
+
+  body.push_back(std::make_unique<vAST::ContinuousAssign>(vAST::make_id("x"),
+                                                          vAST::make_num("7")));
+
+  body.push_back(std::make_unique<vAST::ContinuousAssign>(
+      vAST::make_id("o0"),
+      std::make_unique<vAST::Index>(vAST::make_id("x"), vAST::make_num("0"))));
+
+  body.push_back(std::make_unique<vAST::ContinuousAssign>(
+      vAST::make_id("o1"),
+      std::make_unique<vAST::Slice>(vAST::make_id("x"), vAST::make_num("1"),
+                                    vAST::make_num("0"))));
+
+  std::unique_ptr<vAST::AbstractModule> module = std::make_unique<vAST::Module>(
+      "test_module", std::move(ports), std::move(body));
+
+  std::string raw_str =
+      "module test_module (\n"
+      "    output o0,\n"
+      "    output [1:0] o1\n"
+      ");\n"
+      "wire [2:0] x;\n"
+      "assign x = 7;\n"
+      "assign o0 = x[0];\n"
+      "assign o1 = x[1:0];\n"
+      "endmodule\n";
+
+  EXPECT_EQ(module->toString(), raw_str);
+
+  std::string expected_str = raw_str;
+
+  vAST::AssignInliner transformer;
+  EXPECT_EQ(transformer.visit(std::move(module))->toString(), expected_str);
+}
+
 }  // namespace
 
 int main(int argc, char **argv) {