diff --git a/src/assign_inliner.cpp b/src/assign_inliner.cpp index 879508a..a0ebcd7 100644 --- a/src/assign_inliner.cpp +++ b/src/assign_inliner.cpp @@ -7,9 +7,12 @@ std::unique_ptr Blacklister::visit( std::unique_ptr node) { if (this->blacklist) { auto it = assign_map.find(node->toString()); - bool assigned_to_id = - it != assign_map.end() && dynamic_cast(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(it->second.get()) || + dynamic_cast(it->second.get()) || + dynamic_cast(it->second.get())); + if (!valid_driver) { this->wire_blacklist.insert(node->value); } } diff --git a/tests/assign_inliner.cpp b/tests/assign_inliner.cpp index e6177b0..65fb345 100644 --- a/tests/assign_inliner.cpp +++ b/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> ports; + std::vector, + std::unique_ptr>> + outer_dims; + outer_dims.push_back({vAST::make_num("7"), vAST::make_num("0")}); + ports.push_back(std::make_unique( + std::make_unique(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( + std::make_unique(vAST::make_id("o"), vAST::make_num("3"), + vAST::make_num("0")), + vAST::OUTPUT, vAST::WIRE)); + + std::vector, + std::unique_ptr>> + body; + + body.push_back( + std::make_unique(std::make_unique("x"))); + + body.push_back(std::make_unique( + vAST::make_id("x"), std::make_unique( + std::make_unique("i"), + std::make_unique("0")))); + + body.push_back(std::make_unique( + vAST::make_id("o"), + std::make_unique(vAST::make_id("x"), vAST::make_num("3"), + vAST::make_num("0")))); + + std::unique_ptr module = std::make_unique( + "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) {