[Explicit State Access] Support multiple reads for state element, verify labels are propagated, and decouple next nodes from state reads when transforming state elements.

xls/ir/BUILD

@@ -1893,6 +1893,7 @@ cc_test(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:span",
         "@googletest//:gtest",
     ],
 )

xls/ir/ir_matcher.cc

@@ -654,8 +654,19 @@ bool NextMatcher::MatchAndExplain(
   if (!NodeMatcher::MatchAndExplain(node, listener)) {
     return false;
   }
-  if (label_.has_value() &&
-      !label_->MatchAndExplain(node->As<xls::Next>()->label(), listener)) {
+  const xls::Next* next = node->As<xls::Next>();
+  if (state_element_.has_value()) {
+    if (next->has_state_read()) {
+      *listener << " expected decoupled Next node with only state_element, but "
+                   "it has state_read";
+      return false;
+    }
+    if (!state_element_->MatchAndExplain(next->state_element(), listener)) {
+      *listener << " has incorrect state_element";
+      return false;
+    }
+  }
+  if (label_.has_value() && !label_->MatchAndExplain(next->label(), listener)) {
     *listener << " has incorrect label";
     return false;
   }
@@ -664,6 +675,12 @@ bool NextMatcher::MatchAndExplain(
 
 void NextMatcher::DescribeTo(::std::ostream* os) const {
   std::vector<std::string> additional_fields;
+  if (state_element_.has_value()) {
+    std::stringstream ss;
+    ss << "state_element=";
+    state_element_->DescribeTo(&ss);
+    additional_fields.push_back(ss.str());
+  }
   if (label_.has_value()) {
     std::stringstream ss;
     ss << "label=";

xls/ir/ir_matcher.h

@@ -1329,9 +1329,18 @@ inline ::testing::Matcher<const ::xls::Node*> StateRead() {
 //   EXPECT_THAT(x, m::Next());
 //   EXPECT_THAT(x, m::Next(m::StateRead("foo"), m::Literal(1)));
 //   EXPECT_THAT(x, m::Next(m::StateRead("foo"), m::Literal(1),
-//   m::Literal(1)))
-//   EXPECT_THAT(x, m::Next(m::StateRead("foo"), m::Literal(1),
-//   m::Literal(1), "some_label"))
+//                          m::Literal(1)));
+//   EXPECT_THAT(x, m::NextWithLabel(m::StateRead("foo"), m::Literal(1),
+//                                   "some_label"));
+//   EXPECT_THAT(x, m::NextWithLabel(m::StateRead("foo"), m::Literal(1),
+//                                   m::Literal(1), "some_label"));
+//
+// NextStateElement matcher. Supported forms:
+//
+//   EXPECT_THAT(x, m::NextStateElement(state_element, m::Literal(1),
+//                                      m::Literal(1)));
+//   EXPECT_THAT(x, m::NextStateElementWithLabel(state_element, m::Literal(1),
+//                                               m::Literal(1), "some_label"));
 class NextMatcher : public NodeMatcher {
  public:
   explicit NextMatcher(
@@ -1340,11 +1349,21 @@ class NextMatcher : public NodeMatcher {
           label = std::nullopt)
       : NodeMatcher(Op::kNext, operands), label_(std::move(label)) {}
 
+  explicit NextMatcher(
+      ::testing::Matcher<const ::xls::StateElement*> state_element,
+      absl::Span<const ::testing::Matcher<const Node*>> operands = {},
+      std::optional<::testing::Matcher<const std::optional<std::string>&>>
+          label = std::nullopt)
+      : NodeMatcher(Op::kNext, operands),
+        state_element_(std::move(state_element)),
+        label_(std::move(label)) {}
+
   bool MatchAndExplain(const Node* node,
                        ::testing::MatchResultListener* listener) const override;
   void DescribeTo(::std::ostream* os) const override;
 
  private:
+  std::optional<::testing::Matcher<const ::xls::StateElement*>> state_element_;
   std::optional<::testing::Matcher<const std::optional<std::string>&>> label_;
 };
 
@@ -1374,6 +1393,20 @@ inline ::testing::Matcher<const ::xls::Node*> NextWithLabel(
   return NextMatcher({state_read, value, predicate}, label);
 }
 
+inline ::testing::Matcher<const ::xls::Node*> NextStateElement(
+    ::testing::Matcher<const ::xls::StateElement*> state_element,
+    ::testing::Matcher<const Node*> value,
+    ::testing::Matcher<const Node*> predicate) {
+  return NextMatcher(std::move(state_element), {value, predicate});
+}
+inline ::testing::Matcher<const ::xls::Node*> NextStateElementWithLabel(
+    ::testing::Matcher<const ::xls::StateElement*> state_element,
+    ::testing::Matcher<const Node*> value,
+    ::testing::Matcher<const Node*> predicate,
+    ::testing::Matcher<const std::optional<std::string>&> label) {
+  return NextMatcher(std::move(state_element), {value, predicate}, label);
+}
+
 // RegisterRead matcher. Matches register name only. Supported forms:
 //
 //   EXPECT_THAT(x, m::RegisterRead());

xls/ir/proc.cc

@@ -15,6 +15,7 @@
 #include "xls/ir/proc.h"
 
 #include <algorithm>
+#include <cstddef>
 #include <cstdint>
 #include <iterator>
 #include <memory>
@@ -940,84 +941,149 @@ absl::Status Proc::ConvertToNewStyle() {
 absl::StatusOr<StateElement*> Proc::TransformStateElement(
     StateElement* old_state_element, const Value& init_value,
     Proc::StateElementTransformer& transform) {
-  StateRead* old_state_read = GetStateReadByStateElement(old_state_element);
+  absl::Span<StateRead* const> old_state_reads =
+      GetStateReadsByStateElement(old_state_element);
   std::string orig_name(old_state_element->name());
-  std::string orig_read_name(old_state_read->GetNameView());
-  XLS_ASSIGN_OR_RETURN(std::optional<Node*> read_predicate,
-                       transform.TransformReadPredicate(this, old_state_read));
+
+  std::vector<std::string> orig_read_names;
+  orig_read_names.reserve(old_state_reads.size());
+  for (StateRead* old_state_read : old_state_reads) {
+    orig_read_names.push_back(std::string(old_state_read->GetNameView()));
+  }
+
+  // Create new state element (unread initially).
   XLS_ASSIGN_OR_RETURN(
-      StateRead * new_state_read,
-      AppendStateElement(absl::StrFormat("TEMP_NAME__%s__", orig_name),
-                         init_value, read_predicate,
-                         /*next_state=*/std::nullopt));
-  new_state_read->SetLoc(old_state_read->loc());
-  new_state_read->set_label(old_state_read->label());
-  if (old_state_read->state_element()->non_synthesizable()) {
-    new_state_read->state_element()->SetNonSynthesizable();
+      StateElement * new_state_element,
+      InsertUnreadStateElement(GetStateElementCount(),
+                               absl::StrFormat("TEMP_NAME__%s__", orig_name),
+                               init_value));
+  if (old_state_element->non_synthesizable()) {
+    new_state_element->SetNonSynthesizable();
   }
-  StateElement* new_state_element = new_state_read->state_element();
   std::string temp_name = new_state_element->name();
 
-  XLS_ASSIGN_OR_RETURN(
-      Node * new_state_value,
-      transform.TransformStateRead(this, new_state_read, old_state_read));
-  std::vector<std::pair<Node*, Node*>> to_replace{
-      {old_state_read, new_state_value}};
+  std::vector<std::pair<Node*, Node*>> to_replace;
+  std::vector<StateRead*> new_state_reads;
+  new_state_reads.reserve(old_state_reads.size());
+  // We track the first state read because we need to check the return type of
+  // the state read after transformation. The type is identical across all
+  // state reads.
+  StateRead* first_new_read = nullptr;
+
+  // Transform and create new reads for each old read.
+  for (StateRead* old_state_read : old_state_reads) {
+    XLS_ASSIGN_OR_RETURN(
+        std::optional<Node*> read_predicate,
+        transform.TransformReadPredicate(this, old_state_read));
+    XLS_ASSIGN_OR_RETURN(
+        StateRead * new_state_read,
+        AddStateRead(new_state_element, read_predicate, old_state_read->label(),
+                     old_state_read->loc()));
+    new_state_reads.push_back(new_state_read);
+    if (first_new_read == nullptr) {
+      first_new_read = new_state_read;
+    }
+
+    XLS_ASSIGN_OR_RETURN(
+        Node * new_state_value,
+        transform.TransformStateRead(this, new_state_read, old_state_read));
+    to_replace.push_back({old_state_read, new_state_value});
+  }
+
   struct NextTransformation {
     Next* old_next;
     Node* new_value;
     std::optional<Node*> new_predicate;
+    StateRead* new_state_read;
   };
   std::vector<NextTransformation> transforms;
   for (Next* nxt : next_values(old_state_element)) {
     NextTransformation& new_next = transforms.emplace_back();
     new_next.old_next = nxt;
-    XLS_ASSIGN_OR_RETURN(new_next.new_value, transform.TransformNextValue(
-                                                 this, new_state_read, nxt));
-    XLS_RET_CHECK(new_next.new_value->GetType() == new_state_read->GetType())
+
+    // Find the corresponding new state read.
+    new_next.new_state_read = first_new_read;
+    if (nxt->has_state_read()) {
+      StateRead* old_read = nxt->state_read()->As<StateRead>();
+      for (size_t i = 0; i < old_state_reads.size(); ++i) {
+        if (old_state_reads[i] == old_read) {
+          new_next.new_state_read = new_state_reads[i];
+          break;
+        }
+      }
+    }
+
+    XLS_ASSIGN_OR_RETURN(
+        new_next.new_value,
+        transform.TransformNextValue(this, new_next.new_state_read, nxt));
+    XLS_RET_CHECK(new_next.new_value->GetType() ==
+                  new_next.new_state_read->GetType())
         << "New value is not compatible type. Expected: "
-        << new_state_read->GetType() << " got " << new_next.new_value;
+        << new_next.new_state_read->GetType() << " got " << new_next.new_value;
     XLS_ASSIGN_OR_RETURN(
         new_next.new_predicate,
-        transform.TransformNextPredicate(this, new_state_read, nxt));
+        transform.TransformNextPredicate(this, new_next.new_state_read, nxt));
   }
 
-  // We've transformed all the graph elements. Start replacing them.
-
-  // Switch old_state_read's name to a temporary to-remove name
+  // Rename old element & reads to a temporary to-remove name.
   std::string to_remove_name = UniquifyStateName(
       absl::StrFormat("TO_REMOVE_TRANSFORMED_STATE__%s__", orig_name));
   auto orig_storage = state_elements_.extract(orig_name);
   orig_storage.key() = to_remove_name;
   old_state_element->SetName(to_remove_name);
-  old_state_read->SetName(to_remove_name);
+  for (StateRead* old_state_read : old_state_reads) {
+    old_state_read->SetName(to_remove_name);
+  }
   CHECK(state_elements_.insert(std::move(orig_storage)).inserted);
 
   // Take over the old state element & read names.
   auto new_storage = state_elements_.extract(temp_name);
   new_storage.key() = orig_name;
   new_state_element->SetName(orig_name);
-  new_state_read->SetNameDirectly(orig_read_name);
+  for (size_t i = 0; i < new_state_reads.size(); ++i) {
+    new_state_reads[i]->SetNameDirectly(orig_read_names[i]);
+  }
   CHECK(state_elements_.insert(std::move(new_storage)).inserted);
 
   // Identity-ify the old next nodes and create new ones.
   for (const NextTransformation& nt : transforms) {
-    // Make the next
-    XLS_ASSIGN_OR_RETURN(
-        Next * nxt,
-        MakeNodeWithName<Next>(nt.old_next->loc(), new_state_read, nt.new_value,
-                               nt.new_predicate, nt.old_next->label(),
-                               nt.old_next->GetName()));
+    Next* nxt;
+    if (nt.old_next->has_state_read()) {
+      // Coupled: use the matched new_state_read
+      XLS_ASSIGN_OR_RETURN(
+          nxt,
+          MakeNodeWithName<Next>(nt.old_next->loc(), nt.new_state_read,
+                                 nt.new_value, nt.new_predicate,
+                                 nt.old_next->label(), nt.old_next->GetName()));
+    } else {
+      // Decoupled: use new_state_element directly
+      XLS_ASSIGN_OR_RETURN(
+          nxt,
+          MakeNodeWithName<Next>(nt.old_next->loc(), new_state_element,
+                                 nt.new_value, nt.new_predicate,
+                                 nt.old_next->label(), nt.old_next->GetName()));
+    }
     to_replace.push_back({nt.old_next, nxt});
     // Identity-ify the old next.
-    XLS_RETURN_IF_ERROR(nt.old_next->ReplaceOperandNumber(
-        Next::kValueOperand, nt.old_next->state_read()));
+    if (nt.old_next->has_state_read()) {
+      XLS_RETURN_IF_ERROR(nt.old_next->ReplaceOperandNumber(
+          Next::kValueOperand, nt.old_next->state_read()));
+    } else {
+      XLS_ASSIGN_OR_RETURN(
+          Node * dummy,
+          MakeNode<Literal>(nt.old_next->loc(),
+                            ZeroOfType(old_state_element->type())));
+      XLS_RET_CHECK(nt.old_next->ReplaceOperand(nt.old_next->value(), dummy));
+    }
   }
+
+  // Replace uses.
   for (const auto& [old_n, new_n] : to_replace) {
     XLS_RETURN_IF_ERROR(old_n->ReplaceUsesWith(
         new_n,
         [&](Node* n) {
-          if (n->Is<Next>() && n->As<Next>()->state_read() == old_n) {
+          if (n->Is<Next>() && n->As<Next>()->has_state_read() &&
+              n->As<Next>()->state_read() == old_n) {
             return false;
           }
           return true;