Add support for omitting values from the next (...) line in XLS IR procs

If using `next_value` operations, these values are now redundant; state parameters default to remaining unchanged for the next activation.

This is an all-or-nothing choice; if you want to use the `next (...)` line for any non-token value, you must provide all values, and may use no `next_value` operations.

NOTE: `next_value` operations are not yet supported in any backend, so users should not switch over yet!
PiperOrigin-RevId: 597977645

Author: ericastor

xls/ir/function_builder.cc

@@ -1129,14 +1129,16 @@ absl::StatusOr<Proc*> ProcBuilder::Build(BValue token,
         absl::StrFormat("Recurrent token of proc must be token type, is: %s.",
                         GetType(token)->ToString()));
   }
-  if (next_state.size() != state_params_.size()) {
+
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
+  if (!next_state.empty() && next_state.size() != state_params_.size()) {
     return absl::InvalidArgumentError(
         absl::StrFormat("Number of recurrent state elements given (%d) does "
                         "not equal the number of state elements in the proc "
                         "(%d)",
                         next_state.size(), state_params_.size()));
   }
-  for (int64_t i = 0; i < state_params_.size(); ++i) {
+  for (int64_t i = 0; i < next_state.size(); ++i) {
     if (GetType(next_state[i]) != GetType(GetStateParam(i))) {
       return absl::InvalidArgumentError(
           absl::StrFormat("Recurrent state type %s does not match proc "

xls/ir/function_builder.h

@@ -762,10 +762,11 @@ class ProcBuilder : public BuilderBase {
   BValue GetStateParam(int64_t index) const { return state_params_.at(index); }
 
   // Build the proc using the given BValues as the recurrent token and next
-  // state values respectively. The number of recurrent state elements in
-  // `next_state` must match the number of state parameters.
+  // state values respectively. If `next_state` is not empty, the number of
+  // recurrent state elements in `next_state` must match the number of state
+  // parameters.
   absl::StatusOr<Proc*> Build(BValue token,
-                              absl::Span<const BValue> next_state);
+                              absl::Span<const BValue> next_state = {});
 
   // Adds a state element to the proc with the given initial value. Returns the
   // newly added state parameter.

xls/ir/ir_parser.cc

@@ -1687,6 +1687,8 @@ absl::StatusOr<Parser::BodyResult> Parser::ParseBody(
             "Proc next token name @ %s  was not previously defined: \"%s\"",
             next_token_name.pos().ToHumanString(), next_token_name.value()));
       }
+
+      // TODO: Remove this once fully transitioned over to `next_value` nodes.
       std::vector<BValue> next_state;
       while (scanner_.TryDropToken(LexicalTokenType::kComma)) {
         XLS_ASSIGN_OR_RETURN(Token next_state_name,
@@ -1699,6 +1701,14 @@ absl::StatusOr<Parser::BodyResult> Parser::ParseBody(
         }
         next_state.push_back(name_to_value->at(next_state_name.value()));
       }
+      if (Proc* proc = fb->function()->AsProcOrDie();
+          !next_state.empty() && !proc->next_values().empty()) {
+        return absl::InvalidArgumentError(absl::StrFormat(
+            "Proc includes both next_value nodes (e.g., %s) and next-state "
+            "values on its 'next' line; both cannot be used at the same time.",
+            proc->next_values().front()->GetName()));
+      }
+
       XLS_RETURN_IF_ERROR(
           scanner_.DropTokenOrError(LexicalTokenType::kParenClose));
       result =

xls/ir/ir_parser_test.cc

@@ -1278,7 +1278,6 @@ proc my_proc(my_token: token, my_state: bits[32], init={42}) {
 }
 
 TEST(IrParserTest, ParseProcWithExplicitNext) {
-  // TODO(epastor): Remove the next-value for `my_state` from the last line
   const std::string input = R"(package test
 
 chan ch(bits[32], id=0, kind=streaming, ops=send_receive, flow_control=none, strictness=proven_mutually_exclusive, metadata="""""")
@@ -1289,14 +1288,33 @@ proc my_proc(my_token: token, my_state: bits[32], init={42}) {
   receive.3: (token, bits[32]) = receive(send.1, predicate=literal.2, channel=ch, id=3)
   tuple_index.4: token = tuple_index(receive.3, index=0, id=4)
   next_value.5: () = next_value(param=my_state, value=my_state, id=5)
-  next (tuple_index.4, my_state)
+  next (tuple_index.4)
 }
 )";
   ParsePackageAndCheckDump(input);
 }
 
+TEST(IrParserTest, ParseProcWithMixedNextValueStyles) {
+  const std::string input = R"(package test
+
+chan ch(bits[32], id=0, kind=streaming, ops=send_receive, flow_control=none, strictness=proven_mutually_exclusive, metadata="""""")
+
+proc my_proc(my_token: token, my_state_1: bits[32], my_state_2: bits[32], init={42, 64}) {
+  send.1: token = send(my_token, my_state_1, channel=ch, id=1)
+  literal.2: bits[1] = literal(value=1, id=2)
+  receive.3: (token, bits[32]) = receive(send.1, predicate=literal.2, channel=ch, id=3)
+  tuple_index.4: token = tuple_index(receive.3, index=0, id=4)
+  next_value.5: () = next_value(param=my_state_2, value=my_state_2, id=5)
+  next (tuple_index.4, my_state_1, my_state_2)
+}
+)";
+  EXPECT_THAT(Parser::ParsePackage(input).status(),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Proc includes both next_value nodes (e.g., "
+                                 "next_value.5) and next-state values")));
+}
+
 TEST(IrParserTest, ParseProcWithBadNextParam) {
-  // TODO(epastor): Remove the next-value for `my_state` from the last line
   const std::string input = R"(package test
 
 chan ch(bits[32], id=0, kind=streaming, ops=send_receive, flow_control=none, strictness=proven_mutually_exclusive, metadata="""""")
@@ -1307,7 +1325,7 @@ proc my_proc(my_token: token, my_state: bits[32], init={42}) {
   receive.3: (token, bits[32]) = receive(send.1, predicate=literal.2, channel=ch, id=3)
   tuple_index.4: token = tuple_index(receive.3, index=0, id=4)
   next_value.5: () = next_value(param=not_my_state, value=my_state, id=5)
-  next (tuple_index.4, my_state)
+  next (tuple_index.4)
 }
 )";
   EXPECT_THAT(Parser::ParsePackage(input).status(),
@@ -1318,7 +1336,6 @@ proc my_proc(my_token: token, my_state: bits[32], init={42}) {
 }
 
 TEST(IrParserTest, ParseProcWithBadNextValueType) {
-  // TODO(epastor): Remove the next-value for `my_state` from the last line
   const std::string input = R"(package test
 
 chan ch(bits[32], id=0, kind=streaming, ops=send_receive, flow_control=none, strictness=proven_mutually_exclusive, metadata="""""")
@@ -1329,7 +1346,7 @@ proc my_proc(my_token: token, my_state: bits[32], init={42}) {
   receive.3: (token, bits[32]) = receive(send.1, predicate=literal.2, channel=ch, id=3)
   tuple_index.4: token = tuple_index(receive.3, index=0, id=4)
   next_value.5: () = next_value(param=my_state, value=literal.2, id=5)
-  next (tuple_index.4, my_state)
+  next (tuple_index.4)
 }
 )";
   EXPECT_THAT(Parser::ParsePackage(input).status(),

xls/ir/proc.cc

@@ -85,9 +85,13 @@ std::string Proc::DumpIr() const {
     }
     absl::StrAppend(&res, "  ", node->ToString(), "\n");
   }
+
   absl::StrAppend(&res, "  next (", NextToken()->GetName());
-  for (Node* node : next_state_) {
-    absl::StrAppend(&res, ", ", node->GetName());
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
+  if (next_values_.empty()) {
+    for (Node* node : next_state_) {
+      absl::StrAppend(&res, ", ", node->GetName());
+    }
   }
   absl::StrAppend(&res, ")\n}\n");
   return res;
@@ -222,8 +226,11 @@ absl::StatusOr<Param*> Proc::ReplaceStateElement(
                         "state param %s has uses",
                         index, name(), old_param->GetName()));
   }
+
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   next_state_indices_[next_state_[index]].erase(index);
   next_state_[index] = nullptr;
+
   XLS_RETURN_IF_ERROR(RemoveNode(old_param));
 
   // Construct the new param node, and update all trackers.
@@ -243,6 +250,8 @@ absl::StatusOr<Param*> Proc::ReplaceStateElement(
         next_state.value()->GetType()->ToString(), init_value.ToString()));
   }
   init_values_[index] = init_value;
+
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   next_state_[index] = next_state.value_or(param);
   next_state_indices_[next_state_[index]].insert(index);
 
@@ -251,6 +260,8 @@ absl::StatusOr<Param*> Proc::ReplaceStateElement(
 
 absl::Status Proc::RemoveStateElement(int64_t index) {
   XLS_RET_CHECK_LT(index, GetStateElementCount());
+
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   if (index < GetStateElementCount() - 1) {
     for (auto& [_, indices] : next_state_indices_) {
       // Relabel all indices > `index`.
@@ -264,6 +275,7 @@ absl::Status Proc::RemoveStateElement(int64_t index) {
   }
   next_state_indices_[next_state_[index]].erase(index);
   next_state_.erase(next_state_.begin() + index);
+
   Param* old_param = GetStateParam(index);
   if (!old_param->users().empty()) {
     return absl::InvalidArgumentError(
@@ -294,6 +306,8 @@ absl::StatusOr<Param*> Proc::InsertStateElement(
       MakeNodeWithName<Param>(SourceInfo(), state_param_name,
                               package()->GetTypeForValue(init_value)));
   XLS_RETURN_IF_ERROR(MoveParamToIndex(param, index + 1));
+
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   if (next_state.has_value()) {
     if (!ValueConformsToType(init_value, next_state.value()->GetType())) {
       return absl::InvalidArgumentError(absl::StrFormat(
@@ -316,6 +330,7 @@ absl::StatusOr<Param*> Proc::InsertStateElement(
   }
   next_state_.insert(next_state_.begin() + index, next_state.value_or(param));
   next_state_indices_[next_state_[index]].insert(index);
+
   init_values_.insert(init_values_.begin() + index, init_value);
   return param;
 }
@@ -324,10 +339,13 @@ bool Proc::HasImplicitUse(Node* node) const {
   if (node == NextToken()) {
     return true;
   }
+
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   if (auto it = next_state_indices_.find(node);
       it != next_state_indices_.end() && !it->second.empty()) {
     return true;
   }
+
   return false;
 }
 
@@ -535,10 +553,12 @@ absl::StatusOr<Proc*> Proc::Clone(
   XLS_RETURN_IF_ERROR(
       cloned_proc->SetNextToken(original_to_clone.at(NextToken())));
 
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   for (int64_t i = 0; i < GetStateElementCount(); ++i) {
     XLS_RETURN_IF_ERROR(cloned_proc->SetNextStateElement(
         i, original_to_clone.at(GetNextStateElement(i))));
   }
+
   return cloned_proc;
 }
 

xls/ir/proc.h

@@ -97,13 +97,17 @@ class Proc : public FunctionBase {
   Node* NextToken() const { return next_token_; }
 
   // Returns the nodes holding the next recurrent state value.
+  //
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   absl::Span<Node* const> NextState() const { return next_state_; }
   Node* GetNextStateElement(int64_t index) const {
     return NextState().at(index);
   }
 
   // Return the state element indices for which the given `node` is the next
   // recurrent state value for that element.
+  //
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   absl::btree_set<int64_t> GetNextStateIndices(Node* node) const;
 
   // Returns the type of the given state element.
@@ -120,6 +124,8 @@ class Proc : public FunctionBase {
 
   // Sets the next recurrent state value for the state element of the given
   // index. Node type must match the type of the state element.
+  //
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   absl::Status SetNextStateElement(int64_t index, Node* next);
 
   // Replace all state elements with new state parameters and the given initial
@@ -133,6 +139,8 @@ class Proc : public FunctionBase {
   // than as a std::optional `next_state` argument because initializer lists do
   // not explicitly convert to std::optional<absl::Span> making callsites
   // verbose.
+  //
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   absl::Status ReplaceState(absl::Span<const std::string> state_param_names,
                             absl::Span<const Value> init_values,
                             absl::Span<Node* const> next_state);
@@ -262,9 +270,13 @@ class Proc : public FunctionBase {
   // proc.
   Node* next_token_;
   bool is_new_style_proc_;
+
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   std::vector<Node*> next_state_;
 
   // A map from the `next_state_` nodes back to the indices they control.
+  //
+  // TODO: Remove this once fully transitioned over to `next_value` nodes.
   absl::flat_hash_map<Node*, absl::btree_set<int64_t>> next_state_indices_;
 
   // All channel references in this proc. Channel references can be part of the

xls/ir/verifier.cc

@@ -2070,6 +2070,20 @@ absl::Status VerifyProc(Proc* proc, bool codegen) {
     // Verify that the order of parameters matches the state element order.
     XLS_RET_CHECK_EQ(proc->param(i + 1), proc->GetStateParam(i));
 
+    Param* param = proc->GetStateParam(i);
+    Node* next_state = proc->GetNextStateElement(i);
+    if (next_state == param) {
+      continue;
+    }
+
+    // Verify that this proc does not use `next_value` nodes.
+    if (!proc->next_values().empty()) {
+      return absl::InvalidArgumentError(absl::StrFormat(
+          "Proc %s includes both next_value nodes (e.g., %s) and next-state "
+          "values on its 'next' line; both cannot be used at the same time.",
+          proc->name(), proc->next_values().front()->GetName()));
+    }
+
     // Verify type of state param matches type of the corresponding initial
     // value and next state element.
     XLS_RET_CHECK_EQ(proc->GetStateParam(i)->GetType(),