Support slice and attribute for value of index

include/verilogAST.hpp

@@ -83,34 +83,19 @@ class NumericLiteral : public Expression {
 
   NumericLiteral(std::string value, unsigned int size, bool _signed,
                  Radix radix)
-      : value(value),
-        size(size),
-        _signed(_signed),
-        radix(radix){};
+      : value(value), size(size), _signed(_signed), radix(radix){};
 
   NumericLiteral(std::string value, unsigned int size, bool _signed)
-      : value(value),
-        size(size),
-        _signed(_signed),
-        radix(Radix::DECIMAL){};
+      : value(value), size(size), _signed(_signed), radix(Radix::DECIMAL){};
 
   NumericLiteral(std::string value, unsigned int size)
-      : value(value),
-        size(size),
-        _signed(false),
-        radix(Radix::DECIMAL){};
+      : value(value), size(size), _signed(false), radix(Radix::DECIMAL){};
 
   NumericLiteral(std::string value)
-      : value(value),
-        size(32),
-        _signed(false),
-        radix(Radix::DECIMAL){};
+      : value(value), size(32), _signed(false), radix(Radix::DECIMAL){};
 
   NumericLiteral(std::string value, Radix radix)
-      : value(value),
-        size(32),
-        _signed(false),
-        radix(radix){};
+      : value(value), size(32), _signed(false), radix(radix){};
 
   std::string toString() override;
   auto clone() const { return std::unique_ptr<NumericLiteral>(clone_impl()); }
@@ -209,26 +194,6 @@ class String : public Expression {
   auto clone() const { return std::unique_ptr<String>(clone_impl()); }
 };
 
-class Index : public Expression {
- protected:
-  virtual Index* clone_impl() const override {
-    return new Index(this->id->clone(), this->index->clone());
-  };
-
- public:
-  std::unique_ptr<Identifier> id;
-  std::unique_ptr<Expression> index;
-
-  Index(std::unique_ptr<Identifier> id, std::unique_ptr<Expression> index)
-      : id(std::move(id)), index(std::move(index)){};
-
-  Index(const Index& rhs) : id(rhs.id->clone()), index(rhs.index->clone()){};
-
-  std::string toString() override;
-  ~Index(){};
-  auto clone() const { return std::unique_ptr<Index>(clone_impl()); }
-};
-
 class Slice : public Expression {
  protected:
   virtual Slice* clone_impl() const override {
@@ -256,6 +221,44 @@ class Slice : public Expression {
   auto clone() const { return std::unique_ptr<Slice>(clone_impl()); }
 };
 
+class Index : public Expression {
+  std::variant<std::unique_ptr<Identifier>, std::unique_ptr<Attribute>,
+               std::unique_ptr<Slice>>
+  clone_index_value() const {
+    return std::visit(
+        [](auto&& value) -> std::variant<std::unique_ptr<Identifier>,
+                                         std::unique_ptr<Attribute>,
+                                         std::unique_ptr<Slice>> {
+          return value->clone();
+        },
+        this->value);
+  }
+
+ protected:
+  virtual Index* clone_impl() const override {
+    return new Index(this->clone_index_value(), this->index->clone());
+  };
+
+ public:
+  std::variant<std::unique_ptr<Identifier>, std::unique_ptr<Attribute>,
+               std::unique_ptr<Slice>>
+      value;
+  std::unique_ptr<Expression> index;
+
+  Index(std::variant<std::unique_ptr<Identifier>, std::unique_ptr<Attribute>,
+                     std::unique_ptr<Slice>>
+            value,
+        std::unique_ptr<Expression> index)
+      : value(std::move(value)), index(std::move(index)){};
+
+  Index(const Index& rhs)
+      : value(rhs.clone_index_value()), index(rhs.index->clone()){};
+
+  std::string toString() override;
+  ~Index(){};
+  auto clone() const { return std::unique_ptr<Index>(clone_impl()); }
+};
+
 namespace BinOp {
 enum BinOp {
   LSHIFT,
@@ -578,7 +581,8 @@ class Connections {
     connections.push_back(std::make_pair(name, std::move(expr)));
   }
 
-  // Releases ownership of expression at @name if exists, othwerwise throws error.
+  // Releases ownership of expression at @name if exists, othwerwise throws
+  // error.
   std::unique_ptr<Expression> at(std::string name) {
     auto is_name = [name](auto& element) { return element.first == name; };
     auto it = std::find_if(connections.begin(), connections.end(), is_name);
@@ -610,7 +614,8 @@ class ModuleInstantiation : public StructuralStatement {
   // TODO Need to make sure that the instance parameters are a subset of the
   // module parameters
   ModuleInstantiation(std::string module_name, Parameters parameters,
-                      std::string instance_name, std::unique_ptr<Connections> connections)
+                      std::string instance_name,
+                      std::unique_ptr<Connections> connections)
       : module_name(module_name),
         parameters(std::move(parameters)),
         instance_name(instance_name),

include/verilogAST/assign_inliner.hpp

@@ -59,13 +59,15 @@ class IndexBlacklister : public Transformer {
   //
   // Verilog does not support (y + z)[0]
   std::set<std::string> &wire_blacklist;
+  bool inside_index = false;
 
  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);
+  virtual std::unique_ptr<Identifier> visit(std::unique_ptr<Identifier> node);
 };
 
 class AssignInliner : public Transformer {
@@ -87,9 +89,9 @@ class AssignInliner : public Transformer {
   bool can_inline(std::string key);
 
  public:
-  AssignInliner() : wire_blacklist() {};
-  explicit AssignInliner(std::set<std::string> wire_blacklist) :
-      wire_blacklist(wire_blacklist) {};
+  AssignInliner() : wire_blacklist(){};
+  explicit AssignInliner(std::set<std::string> wire_blacklist)
+      : wire_blacklist(wire_blacklist){};
   using Transformer::visit;
   virtual std::unique_ptr<Expression> visit(std::unique_ptr<Expression> node);
   virtual std::unique_ptr<ContinuousAssign> visit(

src/assign_inliner.cpp

@@ -3,9 +3,18 @@
 
 namespace verilogAST {
 
-std::unique_ptr<Index> IndexBlacklister::visit(
-    std::unique_ptr<Index> node) {
-  this->wire_blacklist.insert(node->id->value);
+std::unique_ptr<Index> IndexBlacklister::visit(std::unique_ptr<Index> node) {
+  bool prev = this->inside_index;
+  this->inside_index = true;
+  node = Transformer::visit(std::move(node));
+  // Restore prev value, since we could be nested inside an index
+  this->inside_index = prev;
+  return node;
+}
+
+std::unique_ptr<Identifier> IndexBlacklister::visit(
+    std::unique_ptr<Identifier> node) {
+  if (this->inside_index) this->wire_blacklist.insert(node->value);
   return node;
 }
 
@@ -62,7 +71,7 @@ std::unique_ptr<ContinuousAssign> AssignMapBuilder::visit(
 
 bool AssignInliner::can_inline(std::string key) {
   if (this->wire_blacklist.count(key)) {
-      return false;
+    return false;
   }
   auto it = assign_map.find(key);
   return it != assign_map.end() && (this->assign_count[key] == 1) &&

src/concat_coalescer.cpp

@@ -1,5 +1,5 @@
-#include <cassert>
 #include "verilogAST/concat_coalescer.hpp"
+#include <cassert>
 
 namespace verilogAST {
 
@@ -8,7 +8,7 @@ namespace {
 struct Run {
   std::string name;
   int first;  // inclusive
-  int last;  // inclusive
+  int last;   // inclusive
 };
 
 class RunOrExpr {
@@ -69,7 +69,11 @@ RunOrExpr makeRunOrExpr(const Expression* arg) {
   if (not index) return RunOrExpr(arg);
   auto as_int = expr_to_int(index->index.get());
   if (not as_int.first) return RunOrExpr(arg);
-  return RunOrExpr(index->id->value, as_int.second, as_int.second);
+  if (not std::holds_alternative<std::unique_ptr<Identifier>>(index->value)) {
+    return RunOrExpr(arg);
+  }
+  auto& id = std::get<std::unique_ptr<Identifier>>(index->value);
+  return RunOrExpr(id->value, as_int.second, as_int.second);
 }
 
 }  // namespace

src/transformer.cpp

@@ -88,7 +88,7 @@ std::unique_ptr<String> Transformer::visit(std::unique_ptr<String> node) {
 }
 
 std::unique_ptr<Index> Transformer::visit(std::unique_ptr<Index> node) {
-  node->id = this->visit(std::move(node->id));
+  node->value = this->visit(std::move(node->value));
   node->index = this->visit(std::move(node->index));
   return node;
 }

src/verilogAST.cpp

@@ -133,7 +133,7 @@ std::string Attribute::toString() {
 std::string String::toString() { return "\"" + value + "\""; }
 
 std::string Index::toString() {
-  return id->toString() + '[' + index->toString() + ']';
+  return variant_to_string(value) + '[' + index->toString() + ']';
 }
 
 std::string Slice::toString() {

tests/basic.cpp

@@ -77,6 +77,16 @@ TEST(BasicTests, TestString) {
 TEST(BasicTests, TestIndex) {
   vAST::Index index(vAST::make_id("x"), vAST::make_num("0"));
   EXPECT_EQ(index.toString(), "x[0]");
+
+  vAST::Index index2(
+      std::make_unique<vAST::Slice>(vAST::make_id("x"), vAST::make_num("3"),
+                                    vAST::make_num("0")),
+      vAST::make_num("0"));
+  EXPECT_EQ(index2.toString(), "x[3:0][0]");
+
+  vAST::Index index3(std::make_unique<vAST::Attribute>(vAST::make_id("x"), "y"),
+                     vAST::make_num("0"));
+  EXPECT_EQ(index3.toString(), "x.y[0]");
 }
 
 TEST(BasicTests, TestSlice) {
@@ -289,14 +299,14 @@ TEST(BasicTests, TestModuleInst) {
                                         make_simple_connections());
 
   EXPECT_EQ(module_inst.toString(),
-          "test_module #(\n"
-          "    .param0(0),\n"
-          "    .param1(1)\n"
-          ") test_module_inst (\n"
-          "    .a(a),\n"
-          "    .b(b[0]),\n"
-          "    .c(c[31:0])\n"
-          ");");
+            "test_module #(\n"
+            "    .param0(0),\n"
+            "    .param1(1)\n"
+            ") test_module_inst (\n"
+            "    .a(a),\n"
+            "    .b(b[0]),\n"
+            "    .c(c[31:0])\n"
+            ");");
 }
 
 TEST(BasicTests, TestModule) {
@@ -580,7 +590,7 @@ TEST(BasicTests, TestIndexCopy) {
   std::unique_ptr<vAST::Index> x1 = std::make_unique<vAST::Index>(*x);
   EXPECT_EQ(x->toString(), "x[y]");
   EXPECT_EQ(x1->toString(), "x[y]");
-  x1->id->value = "z";
+  x1->value = vAST::make_id("z");
   x1->index = std::make_unique<vAST::Identifier>("a");
   EXPECT_EQ(x->toString(), "x[y]");
   EXPECT_EQ(x1->toString(), "z[a]");