From 207409b5d6c022dc7bc7e1cef7e176285c8d07c2 Mon Sep 17 00:00:00 2001
From: Nelson Liang <nelsonliang@google.com>
Date: Wed, 24 Jun 2026 10:21:16 -0700
Subject: [PATCH] [Explicit State Access] Decouple Next nodes in
 array_untuple_pass

PiperOrigin-RevId: 937416235
---
 xls/passes/array_untuple_pass.cc      | 28 +++++++++++++++-----
 xls/passes/array_untuple_pass_test.cc | 37 +++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/xls/passes/array_untuple_pass.cc b/xls/passes/array_untuple_pass.cc
index 80daaa685a..7c009c8d70 100644
--- a/xls/passes/array_untuple_pass.cc
+++ b/xls/passes/array_untuple_pass.cc
@@ -134,11 +134,16 @@ absl::StatusOr<absl::flat_hash_set<Node*>> FindExternalGroups(
       if (absl::c_all_of(state_read->users(), [&](Node* n) -> bool {
             if (n->Is<Next>()) {
               Next* nxt = n->As<Next>();
-              return nxt->state_read() == nxt->value() &&
-                     nxt->state_read() == state_read;
+              if (nxt->has_state_read()) {
+                return nxt->state_read() == nxt->value() &&
+                       nxt->state_read() == state_read;
+              } else {
+                return nxt->value() == state_read &&
+                       nxt->state_element() == state_read->state_element();
+              }
             }
             // TODO(nelsonliang): Handle identity state elements by retrieving
-            // all state reads and verifying all reds are identity updates.
+            // all state reads and verifying all reads are identity updates.
             return false;
           })) {
         excluded.insert(groups.Find(state_read));
@@ -352,10 +357,19 @@ class UntupleVisitor : public DfsVisitorWithDefault {
          iter::zip(iter::count(), state_read_values, update_values)) {
       XLS_RET_CHECK(state_read_node->Is<StateRead>());
       StateRead* state_read = state_read_node->As<StateRead>();
-      XLS_RETURN_IF_ERROR(proc->MakeNodeWithName<Next>(
-                                  n->loc(), state_read, value, n->predicate(),
-                                  n->label(), IdxName(n, idx))
-                              .status());
+      if (n->has_state_read()) {
+        XLS_RETURN_IF_ERROR(proc->MakeNodeWithName<Next>(
+                                    n->loc(), state_read, value, n->predicate(),
+                                    n->label(), IdxName(n, idx))
+                                .status());
+      } else {
+        StateElement* state_element = state_read->state_element();
+        XLS_RET_CHECK(state_element != nullptr);
+        XLS_RETURN_IF_ERROR(proc->MakeNodeWithName<Next>(
+                                    n->loc(), state_element, value,
+                                    n->predicate(), n->label(), IdxName(n, idx))
+                                .status());
+      }
     }
     // Remove this next from consideration.
     if (n->value() != state_read) {
diff --git a/xls/passes/array_untuple_pass_test.cc b/xls/passes/array_untuple_pass_test.cc
index e27d01f18e..9f38454c30 100644
--- a/xls/passes/array_untuple_pass_test.cc
+++ b/xls/passes/array_untuple_pass_test.cc
@@ -525,6 +525,32 @@ TEST_F(ArrayUntuplePassTest, ProcStateArrayWithInvoke) {
   EXPECT_EQ(val2_before, val2_after);
 }
 
+TEST_F(ArrayUntuplePassTest, ProcStateArrayIdentityNextWithStateElement) {
+  auto p = CreatePackage();
+  ProcBuilder pb(TestName(), p.get());
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      Value init_val,
+      ValueBuilder::Array(
+          {ValueBuilder::Tuple({Value(UBits(0, 4)), Value(UBits(0, 8))}),
+           ValueBuilder::Tuple({Value(UBits(0, 4)), Value(UBits(0, 8))})})
+          .Build());
+
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * state_elem,
+                           pb.UnreadStateElement("my_state", init_val,
+                                                 /*non_synthesizable=*/false));
+  BValue state_read = pb.StateRead(state_elem);
+
+  // Identity next (decoupled)
+  pb.Next(state_elem, state_read);
+
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * pr, pb.Build());
+  ScopedRecordIr sri(p.get());
+  EXPECT_THAT(RunPass(p.get()), IsOkAndHolds(false));
+  EXPECT_THAT(pr->StateElements(),
+              UnorderedElementsAre(m::StateElement("my_state")));
+}
+
 TEST_F(ArrayUntuplePassTest, ProcStateArrayNextWithStateElement) {
   auto p = CreatePackage();
   ProcBuilder pb(TestName(), p.get());
@@ -556,6 +582,17 @@ TEST_F(ArrayUntuplePassTest, ProcStateArrayNextWithStateElement) {
   EXPECT_THAT(pr->StateElements(),
               IsSupersetOf({m::StateElement(_, m::Type("bits[4][2]")),
                             m::StateElement(_, m::Type("bits[8][2]"))}));
+
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * se0, pr->GetStateElementByName(
+                                                   "my_state_tuple_element_0"));
+  XLS_ASSERT_OK_AND_ASSIGN(StateElement * se1, pr->GetStateElementByName(
+                                                   "my_state_tuple_element_1"));
+
+  // Verify next values are decoupled (no state_read operand)
+  EXPECT_THAT(pr->next_values(se0),
+              UnorderedElementsAre(m::NextWithStateElement(se0, _)));
+  EXPECT_THAT(pr->next_values(se1),
+              UnorderedElementsAre(m::NextWithStateElement(se1, _)));
 }
 
 void IrFuzzArrayUntuple(FuzzPackageWithArgs fuzz_package_with_args) {