Improve slice/index inlining logic

We generally can't inline into slice/index nodes because we can't have
expressions such as `(x + y)[1:0]`.  However, we can have nested
index/slice nodes, so we improve the blacklisting logic to not block
inlining if we see a slice or index driver.

Note, we assume that the index/slice driver is well formed, so if we
have a driver of the form <...>[3:0], the contents of the slice (<...>)
are assumed to be valid (so it could contain an ID, another slice, or an
index).

src/assign_inliner.cpp

@@ -7,9 +7,12 @@ std::unique_ptr<Identifier> Blacklister::visit(
     std::unique_ptr<Identifier> node) {
   if (this->blacklist) {
     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) {
+    // Can only inline if driven by identifier, index, or slice
+    bool valid_driver = it != assign_map.end() &&
+                        (dynamic_cast<Identifier*>(it->second.get()) ||
+                         dynamic_cast<Index*>(it->second.get()) ||
+                         dynamic_cast<Slice*>(it->second.get()));
+    if (!valid_driver) {
       this->wire_blacklist.insert(node->value);
     }
   }

tests/assign_inliner.cpp

@@ -868,6 +868,66 @@ TEST(InlineAssignTests, TestNoInlineSliceUnlessID) {
   EXPECT_EQ(transformer.visit(std::move(module))->toString(), expected_str);
 }
 
+TEST(InlineAssignTests, TestInlineSliceOfIndex) {
+  std::vector<std::unique_ptr<vAST::AbstractPort>> ports;
+  std::vector<std::pair<std::unique_ptr<vAST::Expression>,
+                        std::unique_ptr<vAST::Expression>>>
+      outer_dims;
+  outer_dims.push_back({vAST::make_num("7"), vAST::make_num("0")});
+  ports.push_back(std::make_unique<vAST::Port>(
+      std::make_unique<vAST::NDVector>(vAST::make_id("i"), vAST::make_num("4"),
+                                       vAST::make_num("0"),
+                                       std::move(outer_dims)),
+      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::Identifier>("x")));
+
+  body.push_back(std::make_unique<vAST::ContinuousAssign>(
+      vAST::make_id("x"), std::make_unique<vAST::Index>(
+                              std::make_unique<vAST::Identifier>("i"),
+                              std::make_unique<vAST::NumericLiteral>("0"))));
+
+  body.push_back(std::make_unique<vAST::ContinuousAssign>(
+      vAST::make_id("o"),
+      std::make_unique<vAST::Slice>(vAST::make_id("x"), 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] i [7:0],\n"
+      "    output [3:0] o\n"
+      ");\n"
+      "wire x;\n"
+      "assign x = i[0];\n"
+      "assign o = x[3:0];\n"
+      "endmodule\n";
+
+  EXPECT_EQ(module->toString(), raw_str);
+
+  std::string expected_str =
+      "module test_module (\n"
+      "    input [4:0] i [7:0],\n"
+      "    output [3:0] o\n"
+      ");\n"
+      "assign o = i[0][3:0];\n"
+      "endmodule\n";
+
+  vAST::AssignInliner transformer;
+  EXPECT_EQ(transformer.visit(std::move(module))->toString(), expected_str);
+}
+
 }  // namespace
 
 int main(int argc, char **argv) {