[Explicit State Access] Allow multiple StateRead nodes to map to a single state element.

Previously State Elements and State Reads had a 1:1 mapping. Explicit state access now permits multiple reads from the same state element across different code paths, such as within If/Else blocks.

PiperOrigin-RevId: 895959816

Author: NL02

xls/ir/proc.cc

@@ -214,13 +214,15 @@ absl::Status Proc::RemoveStateElement(int64_t index) {
   StateElement* old_state_element = GetStateElement(index);
   auto old_state_read_it = state_reads_.find(old_state_element);
   XLS_RET_CHECK(old_state_read_it != state_reads_.end());
-  if (!old_state_read_it->second->users().empty()) {
-    return absl::InvalidArgumentError(absl::StrFormat(
-        "Cannot remove state element %d of proc %s, existing "
-        "state read %s has uses",
-        index, name(), old_state_read_it->second->GetNameView()));
+  for (StateRead* read : old_state_read_it->second) {
+    if (!read->users().empty()) {
+      return absl::InvalidArgumentError(
+          absl::StrFormat("Cannot remove state element %d of proc %s, existing "
+                          "state read %s has uses",
+                          index, name(), read->GetNameView()));
+    }
+    XLS_RETURN_IF_ERROR(RemoveNode(read));
   }
-  XLS_RETURN_IF_ERROR(RemoveNode(old_state_read_it->second));
   // TODO(allight): This should ideally not need to be done manually.
   state_reads_.erase(old_state_read_it);
 
@@ -232,11 +234,14 @@ absl::Status Proc::RemoveStateElement(int64_t index) {
 absl::Status Proc::RemoveAllStateElements() {
   // TODO(allight): This relies on side tables being valid. For now just let it
   // go.
-  for (const auto& [elem, read] : state_reads_) {
-    if (read != nullptr) {
-      XLS_RETURN_IF_ERROR(RemoveNode(read))
-          << "Cannot remove " << elem->ToString() << " of proc " << name()
-          << " because read '" << read->ToString() << "' could not be removed.";
+  for (const auto& [elem, reads] : state_reads_) {
+    for (StateRead* read : reads) {
+      if (read != nullptr) {
+        XLS_RETURN_IF_ERROR(RemoveNode(read))
+            << "Cannot remove " << elem->ToString() << " of proc " << name()
+            << " because read '" << read->ToString()
+            << "' could not be removed.";
+      }
     }
     XLS_RETURN_IF_ERROR(state_name_uniquer_.ReleaseIdentifier(elem->name()))
         << "Cannot release name of " << elem->ToString();
@@ -278,7 +283,7 @@ absl::StatusOr<StateRead*> Proc::InsertStateElement(
                        MakeNodeWithName<StateRead>(
                            loc, state_element, read_predicate,
                            /*label=*/std::nullopt, state_element->name()));
-  state_reads_[state_element] = state_read;
+  state_reads_[state_element].push_back(state_read);
 
   if (next_state.has_value()) {
     if (!ValueConformsToType(init_value, next_state.value()->GetType())) {
@@ -351,14 +356,13 @@ absl::StatusOr<Proc*> Proc::Clone(
     return mapping.at(orig);
   };
   for (StateElement* state_element : StateElements()) {
-    StateRead* state_read = state_reads_.at(state_element);
-    XLS_ASSIGN_OR_RETURN(
-        StateRead * cloned_state_read,
-        cloned_proc->AppendStateElement(
-            remap_name(state_name_remapping, state_element->name()),
-            state_element->initial_value(), state_read->predicate(),
-            /*next_state=*/std::nullopt));
-    original_to_clone[state_read] = cloned_state_read;
+    XLS_RETURN_IF_ERROR(
+        cloned_proc
+            ->InsertUnreadStateElement(
+                cloned_proc->GetStateElementCount(),
+                remap_name(state_name_remapping, state_element->name()),
+                state_element->initial_value())
+            .status());
   }
   if (is_new_style_proc()) {
     absl::flat_hash_map<ChannelInterface*, ChannelInterface*> channel_map;
@@ -445,7 +449,23 @@ absl::StatusOr<Proc*> Proc::Clone(
 
     switch (node->op()) {
       case Op::kStateRead: {
-        continue;
+        StateRead* src = node->As<StateRead>();
+        StateElement* src_elem = src->state_element();
+        XLS_ASSIGN_OR_RETURN(int64_t idx, GetStateElementIndex(src_elem));
+        StateElement* cloned_elem = cloned_proc->GetStateElement(idx);
+
+        std::optional<Node*> cloned_predicate;
+        if (src->predicate().has_value()) {
+          cloned_predicate = original_to_clone.at(src->predicate().value());
+        }
+
+        XLS_ASSIGN_OR_RETURN(StateRead * cloned_state_read,
+                             cloned_proc->MakeNodeWithName<StateRead>(
+                                 src->loc(), cloned_elem, cloned_predicate,
+                                 /*label=*/std::nullopt, cloned_elem->name()));
+        cloned_proc->state_reads_[cloned_elem].push_back(cloned_state_read);
+        original_to_clone[node] = cloned_state_read;
+        break;
       }
       case Op::kReceive: {
         Receive* src = node->As<Receive>();
@@ -1000,10 +1020,8 @@ absl::Status Proc::InternalRebuildSideTables() {
   state_reads_.clear();
   for (Node* n : nodes()) {
     if (n->Is<StateRead>()) {
-      XLS_RET_CHECK(!state_reads_.contains(n->As<StateRead>()->state_element()))
-          << "Duplicate state element read: "
-          << n->As<StateRead>()->state_element();
-      state_reads_[n->As<StateRead>()->state_element()] = n->As<StateRead>();
+      state_reads_[n->As<StateRead>()->state_element()].push_back(
+          n->As<StateRead>());
     } else if (n->Is<Next>()) {
       next_values_.push_back(n->As<Next>());
       next_values_by_state_element_[n->As<Next>()->state_element()].insert(

xls/ir/proc.h

@@ -95,10 +95,23 @@ class Proc : public FunctionBase {
     return state_elements_.contains(name);
   }
 
+  // Remove legacy getters after all downstream passes migrate logic.
   StateRead* GetStateRead(int64_t index) const {
-    return state_reads_.at(GetStateElement(index));
+    return GetStateReads(index).front();
   }
+
   StateRead* GetStateReadByStateElement(StateElement* state_element) const {
+    return GetStateReadsByStateElement(state_element).front();
+  }
+
+  // Get state reads for a state element at the given index.
+  absl::Span<StateRead* const> GetStateReads(int64_t index) const {
+    return state_reads_.at(GetStateElement(index));
+  }
+
+  // Get state reads for a state element.
+  absl::Span<StateRead* const> GetStateReadsByStateElement(
+      StateElement* state_element) const {
     return state_reads_.at(state_element);
   }
 
@@ -403,8 +416,8 @@ class Proc : public FunctionBase {
   absl::flat_hash_map<std::string, std::unique_ptr<StateElement>>
       state_elements_;
 
-  // Map of the unique StateRead node for each state element.
-  absl::flat_hash_map<StateElement*, StateRead*> state_reads_;
+  // Map of StateRead nodes for each state element.
+  absl::flat_hash_map<StateElement*, std::vector<StateRead*>> state_reads_;
 
   // Vector of state element pointers. Kept in sync with the state_elements_
   // map. Enables easy, stable iteration over state elements. With this vector,

xls/ir/proc_test.cc

@@ -205,6 +205,43 @@ TEST_F(ProcTest, StatelessProc) {
   EXPECT_EQ(proc->DumpIr(), "proc p() {\n}\n");
 }
 
+TEST_F(ProcTest, MultipleStateReads) {
+  auto p = CreatePackage();
+  ProcBuilder pb("p", p.get());
+  BValue tkn = pb.StateElement("tkn", Value::Token());
+  BValue state = pb.StateElement("x", Value(UBits(42, 32)));
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc, pb.Build({tkn, state}));
+
+  StateElement* state_elem = proc->GetStateElement(1);
+
+  EXPECT_EQ(proc->GetStateReads(1).size(), 1);
+  StateRead* read1 = proc->GetStateReads(1).front();
+
+  // Second Read
+  XLS_ASSERT_OK_AND_ASSIGN(
+      StateRead * read2,
+      proc->MakeNodeWithName<StateRead>(SourceInfo(), state_elem,
+                                        /*predicate=*/std::nullopt,
+                                        /*label=*/std::nullopt, "x_read2"));
+  XLS_ASSERT_OK(proc->RebuildSideTables());
+
+  EXPECT_EQ(proc->GetStateReads(1).size(), 2);
+  EXPECT_THAT(proc->GetStateReads(1), ElementsAre(read1, read2));
+
+  EXPECT_EQ(proc->GetStateReadsByStateElement(state_elem).size(), 2);
+  EXPECT_THAT(proc->GetStateReadsByStateElement(state_elem),
+              ElementsAre(read1, read2));
+
+  // Remove the second read.
+  std::string read2_name = read2->GetName();
+  XLS_ASSERT_OK(proc->RemoveNode(read2));
+  XLS_ASSERT_OK(proc->RebuildSideTables());
+
+  // Now we should have 1 read again.
+  EXPECT_EQ(proc->GetStateReads(1).size(), 1);
+  EXPECT_EQ(proc->GetStateReads(1).front(), read1);
+}
+
 TEST_F(ProcTest, RemoveStateThatStillHasUse) {
   // Don't call CreatePackage which creates a VerifiedPackage because we
   // intentionally create a malformed proc.
@@ -254,10 +291,10 @@ TEST_F(ProcTest, Clone) {
   EXPECT_EQ(clone->DumpIr(),
             R"(proc cloned(tkn: token, state: bits[32], init={token, 42}) {
   tkn: token = state_read(state_element=tkn, id=12)
-  literal.14: bits[32] = literal(value=1, id=14)
-  state: bits[32] = state_read(state_element=state, id=13)
+  literal.13: bits[32] = literal(value=1, id=13)
+  state: bits[32] = state_read(state_element=state, id=14)
   receive_3: (token, bits[32]) = receive(tkn, channel=cloned_chan, id=15)
-  add.16: bits[32] = add(literal.14, state, id=16)
+  add.16: bits[32] = add(literal.13, state, id=16)
   tuple_index.17: bits[32] = tuple_index(receive_3, index=1, id=17)
   tuple_index.18: token = tuple_index(receive_3, index=0, id=18)
   add.19: bits[32] = add(add.16, tuple_index.17, id=19)
@@ -304,10 +341,10 @@ proc cloned<input_chan: bits[32] in, chan: bits[32] out>(tkn: token, state: bits
   chan_interface input_chan(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
   chan_interface chan(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
   tkn: token = state_read(state_element=tkn, id=1)
-  literal.3: bits[32] = literal(value=1, id=3)
-  state: bits[32] = state_read(state_element=state, id=2)
+  literal.2: bits[32] = literal(value=1, id=2)
+  state: bits[32] = state_read(state_element=state, id=3)
   receive_3: (token, bits[32]) = receive(tkn, channel=input_chan, id=4)
-  add.5: bits[32] = add(literal.3, state, id=5)
+  add.5: bits[32] = add(literal.2, state, id=5)
   tuple_index.6: bits[32] = tuple_index(receive_3, index=1, id=6)
   tuple_index.7: token = tuple_index(receive_3, index=0, id=7)
   add.8: bits[32] = add(add.5, tuple_index.6, id=8)
@@ -355,15 +392,15 @@ TEST_F(ProcTest, CloneNewStyle) {
   chan baz(bits[32], id=0, kind=streaming, ops=send_receive, flow_control=ready_valid, strictness=proven_mutually_exclusive)
   chan_interface baz(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=none, flop_kind=none)
   chan_interface baz(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=none, flop_kind=none)
-  tkn: token = literal(value=token, id=14)
-  receive_3: (token, bits[32]) = receive(tkn, channel=foo, id=15)
-  tuple_index.16: token = tuple_index(receive_3, index=0, id=16)
-  receive_6: (token, bits[32]) = receive(tuple_index.16, channel=baz, id=17)
-  tuple_index.18: token = tuple_index(receive_6, index=0, id=18)
-  state: bits[32] = state_read(state_element=state, id=13)
+  tkn: token = literal(value=token, id=13)
+  receive_3: (token, bits[32]) = receive(tkn, channel=foo, id=14)
+  tuple_index.15: token = tuple_index(receive_3, index=0, id=15)
+  receive_6: (token, bits[32]) = receive(tuple_index.15, channel=baz, id=16)
+  tuple_index.17: token = tuple_index(receive_6, index=0, id=17)
+  state: bits[32] = state_read(state_element=state, id=18)
   tuple_index.19: bits[32] = tuple_index(receive_3, index=1, id=19)
   tuple_index.20: bits[32] = tuple_index(receive_6, index=1, id=20)
-  send_9: token = send(tuple_index.18, state, channel=bar, id=21)
+  send_9: token = send(tuple_index.17, state, channel=bar, id=21)
   add.22: bits[32] = add(tuple_index.19, tuple_index.20, id=22)
   send_10: token = send(send_9, state, channel=baz, id=23)
   next_value.24: () = next_value(param=state, value=add.22, id=24)
@@ -556,7 +593,8 @@ TEST_F(ScheduledProcTest, StageAddAndClear) {
   proc->ClearStages();
   EXPECT_TRUE(proc->stages().empty());
   // Re-stage the state element to satisfy the verifier.
-  XLS_ASSERT_OK(proc->AddNodeToStage(0, proc->GetStateRead(0)).status());
+  XLS_ASSERT_OK(
+      proc->AddNodeToStage(0, proc->GetStateReads(0).front()).status());
 }
 
 TEST_F(ScheduledProcTest, AddEmptyStages) {
@@ -596,7 +634,8 @@ TEST_F(ScheduledProcTest, GetStageIndex) {
   EXPECT_THAT(proc->GetStageIndex(x), IsOkAndHolds(1));
   EXPECT_THAT(proc->GetStageIndex(y), IsOkAndHolds(2));
   EXPECT_THAT(proc->GetStageIndex(add), StatusIs(absl::StatusCode::kNotFound));
-  EXPECT_THAT(proc->GetStageIndex(proc->GetStateRead(0)), IsOkAndHolds(0));
+  EXPECT_THAT(proc->GetStageIndex(proc->GetStateReads(0).front()),
+              IsOkAndHolds(0));
 
   // The verifier requires that every node be in a stage before we finish.
   ASSERT_THAT(proc->AddNodeToStage(2, add), IsOkAndHolds(true));