From 96860eedf60cf465eefb68dcbadcb938c4ceaec8 Mon Sep 17 00:00:00 2001
From: Lenny Truong <lenny@stanford.edu>
Date: Tue, 17 Mar 2020 20:21:37 -0700
Subject: [PATCH] Add logic to blacklist inlining of index references

---
 include/verilogAST/assign_inliner.hpp | 17 ++++++++++
 src/assign_inliner.cpp                |  9 +++++
 tests/assign_inliner.cpp              | 49 +++++++++++++++++++++++++++
 3 files changed, 75 insertions(+)

diff --git a/include/verilogAST/assign_inliner.hpp b/include/verilogAST/assign_inliner.hpp
index df569d4..0611c3a 100644
--- a/include/verilogAST/assign_inliner.hpp
+++ b/include/verilogAST/assign_inliner.hpp
@@ -51,6 +51,23 @@ class WireReadCounter : public Transformer {
   virtual std::unique_ptr<Declaration> visit(std::unique_ptr<Declaration> node);
 };
 
+class IndexBlacklister : public Transformer {
+  // 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;
+
+ public:
+  IndexBlacklister(std::set<std::string> &wire_blacklist)
+      : wire_blacklist(wire_blacklist){};
+
+  using Transformer::visit;
+  virtual std::unique_ptr<Index> visit(std::unique_ptr<Index> node);
+};
+
 class AssignInliner : public Transformer {
   std::map<std::string, int> read_count;
   std::map<std::string, int> assign_count;
diff --git a/src/assign_inliner.cpp b/src/assign_inliner.cpp
index d567e7a..1243396 100644
--- a/src/assign_inliner.cpp
+++ b/src/assign_inliner.cpp
@@ -3,6 +3,12 @@
 
 namespace verilogAST {
 
+std::unique_ptr<Index> IndexBlacklister::visit(
+    std::unique_ptr<Index> node) {
+  this->wire_blacklist.insert(node->id->value);
+  return node;
+}
+
 std::unique_ptr<Identifier> WireReadCounter::visit(
     std::unique_ptr<Identifier> node) {
   this->read_count[node->toString()]++;
@@ -160,6 +166,9 @@ std::unique_ptr<Module> AssignInliner::visit(std::unique_ptr<Module> node) {
   WireReadCounter counter(this->read_count);
   node = counter.visit(std::move(node));
 
+  IndexBlacklister index_blacklist(this->wire_blacklist);
+  node = index_blacklist.visit(std::move(node));
+
   std::vector<std::unique_ptr<AbstractPort>> new_ports;
   for (auto&& item : node->ports) {
     new_ports.push_back(this->visit(std::move(item)));
diff --git a/tests/assign_inliner.cpp b/tests/assign_inliner.cpp
index b10d2fb..5d32d4b 100644
--- a/tests/assign_inliner.cpp
+++ b/tests/assign_inliner.cpp
@@ -623,6 +623,55 @@ TEST(InlineAssignTests, TestInstConn) {
   EXPECT_EQ(transformer.visit(std::move(module))->toString(), expected_str);
 }
 
+TEST(InlineAssignTests, TestNoInlineIndex) {
+  std::vector<std::unique_ptr<vAST::AbstractPort>> ports;
+  ports.push_back(std::make_unique<vAST::Port>(
+      std::make_unique<vAST::Vector>(vAST::make_id("i1"),
+                                     vAST::make_num("1"), vAST::make_num("0")),
+      vAST::INPUT, vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(
+      std::make_unique<vAST::Vector>(vAST::make_id("i2"),
+                                     vAST::make_num("1"), vAST::make_num("0")),
+      vAST::INPUT, vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("o"), 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::Identifier>("x")));
+
+  body.push_back(std::make_unique<vAST::ContinuousAssign>(
+      vAST::make_id("x"),
+      std::make_unique<vAST::BinaryOp>(vAST::make_id("i1"), vAST::BinOp::ADD,
+                                       vAST::make_id("i2"))));
+
+  body.push_back(std::make_unique<vAST::ContinuousAssign>(
+      vAST::make_id("o"),
+      std::make_unique<vAST::Index>(vAST::make_id("x"), 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 [1:0] i1,\n"
+      "    input [1:0] i2,\n"
+      "    output o\n"
+      ");\n"
+      "wire x;\n"
+      "assign x = i1 + i2;\n"
+      "assign o = x[0];\n"
+      "endmodule\n";
+
+  EXPECT_EQ(module->toString(), raw_str);
+
+  vAST::AssignInliner transformer;
+  EXPECT_EQ(transformer.visit(std::move(module))->toString(), raw_str);
+}
+
 }  // namespace
 
 int main(int argc, char **argv) {