diff --git a/xls/ir/nodes.cc b/xls/ir/nodes.cc index acb4447d71..ec21bfc2d2 100755 --- a/xls/ir/nodes.cc +++ b/xls/ir/nodes.cc @@ -985,9 +985,16 @@ bool StateRead::IsDefinitelyEqualTo(const Node* other) const { std::vector StateRead::GetNextValues() const { std::vector next_values; - for (Node* user : users()) { - if (user->Is() && user->As()->state_read() == this) { - next_values.push_back(user->As()); + if (function_base()->AsProcOrDie()->uses_decoupled_next()) { + const auto& next_values_se = + function_base()->AsProcOrDie()->next_values(state_element_); + next_values.insert(next_values.end(), next_values_se.begin(), + next_values_se.end()); + } else { + for (Node* user : users()) { + if (user->Is() && user->As()->state_read() == this) { + next_values.push_back(user->As()); + } } } return next_values; diff --git a/xls/ir/proc_test.cc b/xls/ir/proc_test.cc index a8df186c98..45f709e871 100644 --- a/xls/ir/proc_test.cc +++ b/xls/ir/proc_test.cc @@ -288,6 +288,41 @@ TEST_F(ProcTest, RemoveStateThatStillHasUse) { HasSubstr("state read st has uses"))); } +TEST_F(ProcTest, GetNextStateReadDecoupled) { + auto p = CreatePackage(); + TokenlessProcBuilder pb(TestName(), "tkn", p.get()); + XLS_ASSERT_OK_AND_ASSIGN(StateElement * state_elem, + pb.UnreadStateElement("x", Value(UBits(42, 32)), + /*non_synthesizable=*/false)); + BValue read = pb.StateRead(state_elem); + BValue next_val = pb.Add(read, pb.Literal(Value(UBits(1, 32)))); + XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build({})); + + StateRead* state_read = proc->GetStateReadByStateElement(state_elem); + EXPECT_EQ(state_read->GetNextValues().size(), 0); + + XLS_ASSERT_OK(proc->MakeNodeWithName(SourceInfo(), state_elem, + next_val.node(), + /*predicate=*/std::nullopt, + /*label=*/std::nullopt, "x_next") + .status()); + EXPECT_EQ(state_read->GetNextValues().size(), 1); + + XLS_ASSERT_OK_AND_ASSIGN( + Literal * next_value, + proc->MakeNode(SourceInfo(), Value(UBits(43, 32)))); + XLS_ASSERT_OK( + proc->MakeNodeWithName(SourceInfo(), state_elem, next_value, + /*predicate=*/std::nullopt, + /*label=*/std::nullopt, "second_next") + .status()); + EXPECT_EQ(state_read->GetNextValues().size(), 2); + EXPECT_THAT(state_read->GetNextValues(), + UnorderedElementsAre( + m::NextWithStateElement(state_elem, m::Add()), + m::NextWithStateElement(state_elem, m::Literal(43)))); +} + TEST_F(ProcTest, InsertStateElementDecoupled) { auto p = CreatePackage(); ProcBuilder pb("p", p.get());