From e445448ad6ce3e6a0bf21eb150f26866856d1e1d Mon Sep 17 00:00:00 2001
From: Mateusz Gancarz <mgancarz@antmicro.com>
Date: Tue, 27 Jan 2026 14:40:43 +0100
Subject: [PATCH] dslx: allow non-trivial final expressions in proc's config

---
 xls/dslx/frontend/parser.cc                   |  10 --
 xls/dslx/type_system/proc_typecheck_test.cc   |   4 +-
 .../typecheck_module_v2_test.cc               | 148 ++++++++++++++++++
 3 files changed, 149 insertions(+), 13 deletions(-)

diff --git a/xls/dslx/frontend/parser.cc b/xls/dslx/frontend/parser.cc
index f71265ada2..cbb1f8293b 100644
--- a/xls/dslx/frontend/parser.cc
+++ b/xls/dslx/frontend/parser.cc
@@ -3120,16 +3120,6 @@ absl::StatusOr<Function*> Parser::ParseProcConfig(
     // Implicitly nil tuple as a result.
   } else {
     Expr* final_expr = std::get<Expr*>(block->statements().back()->wrapped());
-
-    if (dynamic_cast<XlsTuple*>(final_expr) == nullptr) {
-      Span final_stmt_span =
-          ToAstNode(block->statements().back()->wrapped())->GetSpan().value();
-      return ParseErrorStatus(
-          final_stmt_span,
-          "The final expression in a Proc config must be a tuple with one "
-          "element for each Proc data member.");
-    }
-
     VLOG(5) << "ParseProcConfig; final expr: `" << final_expr->ToString()
             << "`";
   }
diff --git a/xls/dslx/type_system/proc_typecheck_test.cc b/xls/dslx/type_system/proc_typecheck_test.cc
index 4bc3b09a91..e733d4b4c4 100644
--- a/xls/dslx/type_system/proc_typecheck_test.cc
+++ b/xls/dslx/type_system/proc_typecheck_test.cc
@@ -119,9 +119,7 @@ proc entry {
   EXPECT_THAT(TypecheckV2(kProgram),
               absl_testing::StatusIs(
                   absl::StatusCode::kInvalidArgument,
-                  testing::HasSubstr(
-                      "final expression in a Proc config must be a tuple with "
-                      "one element for each Proc data member")));
+                  HasTypeMismatch("(chan<u32> in,)", "u32")));
 }
 
 TEST(TypecheckTest, RecvIfDefaultValueWrongType) {
diff --git a/xls/dslx/type_system_v2/typecheck_module_v2_test.cc b/xls/dslx/type_system_v2/typecheck_module_v2_test.cc
index 532c1feb67..5185061c2a 100644
--- a/xls/dslx/type_system_v2/typecheck_module_v2_test.cc
+++ b/xls/dslx/type_system_v2/typecheck_module_v2_test.cc
@@ -7416,6 +7416,154 @@ proc Proc {
                 HasNodeWithType("b", "chan(sN[32], dir=in)[2]"))));
 }
 
+TEST(TypecheckV2Test, ProcConfigRequireTuple) {
+  EXPECT_THAT(
+      R"(
+proc Proc {
+  input: chan<()> in;
+  config(input: chan<()> in) {
+    (input)
+  }
+  init { () }
+  next(state: ()) { () }
+}
+
+)",
+      TypecheckFails(
+          HasTypeMismatch("(chan<()> in,)", "chan<()> in")));
+}
+
+TEST(TypecheckV2Test, ProcConfigTooFewChannels) {
+  EXPECT_THAT(
+      R"(
+proc Proc {
+  input: chan<()> in;
+  output: chan<()> out;
+  config(input: chan<()> in) {
+    (input,)
+  }
+  init { () }
+  next(state: ()) { () }
+}
+
+)",
+      TypecheckFails(
+          HasSubstr("Cannot match a 2-element tuple to 1 values.")));
+}
+
+TEST(TypecheckV2Test, ProcConfigTooManyChannels) {
+  EXPECT_THAT(
+      R"(
+proc Proc {
+  req: chan<()> in;
+  resp: chan<()> out;
+  config(data_in: chan<()> in) {
+    let (resp, req) = chan<()>("io");
+    (req, resp, data_in)
+  }
+  init { () }
+  next(state: ()) { () }
+}
+
+)",
+      TypecheckFails(
+          HasSubstr("Out-of-bounds tuple index specified: 2")));
+}
+
+TEST(TypecheckV2Test, ProcConfigMismatchingChannelTypes) {
+  EXPECT_THAT(
+      R"(
+proc Proc {
+  req: chan<()> in;
+  resp: chan<()> out;
+  data_in: chan<u32> in;
+  data_out: chan<u32> out;
+  config(data_in: chan<u32> in, data_out: chan<u32> out) {
+    let (resp, req) = chan<()>("io");
+    (data_in, data_out, req, resp)
+  }
+  init { () }
+  next(state: ()) { () }
+}
+
+)",
+      TypecheckFails(HasTypeMismatch("()", "u32")));
+}
+
+TEST(TypecheckV2Test, ProcWithBranchedFinalExpression) {
+  EXPECT_THAT(
+      R"(
+const A = u32:5;
+proc Proc {
+  input: chan<()> in;
+  output: chan<()> out;
+  config() {
+    const if A == u32:1 {
+      let (first_output, first_input) = chan<()>("first");
+      (first_input, first_output)
+    } else {
+      let (second_output, second_input) = chan<()>("second");
+      (second_input, second_output)
+    }
+  }
+  init { () }
+  next(state: ()) { () }
+}
+
+)",
+      TypecheckSucceeds(
+          AllOf(HasNodeWithType("second_input", "chan((), dir=in)"),
+                HasNodeWithType("second_output", "chan((), dir=out)"))));
+}
+
+TEST(TypecheckV2Test, ProcConfigBranchedChannelTypeMismatch) {
+  EXPECT_THAT(
+      R"(
+const A = u32:4;
+proc Proc {
+  input: chan<u32> in;
+  output: chan<u32> out;
+  config() {
+    const if A == u32:5 {
+      let (first_output, first_input) = chan<u32>("first");
+      (first_input, first_output)
+    } else {
+      let (second_output, second_input) = chan<()>("second");
+      (second_input, second_output)
+    }
+  }
+  init { () }
+  next(state: ()) { () }
+}
+
+)",
+      TypecheckFails(HasTypeMismatch("()", "u32")));
+}
+
+TEST(TypecheckV2Test, ProcConfigFailedBranchedFinalExpression) {
+  EXPECT_THAT(
+      R"(
+const A = u32:5;
+proc Proc {
+  input: chan<()> in;
+  output: chan<()> out;
+  config() {
+    const if A == u32:5 {
+      let (first_output, first_input) = chan<()>("first");
+      (first_input,)
+    } else {
+      let (second_output, second_input) = chan<()>("second");
+      (second_input, second_output)
+    }
+  }
+  init { () }
+  next(state: ()) { () }
+}
+
+)",
+      TypecheckFails(HasSubstr("Cannot match a 2-element tuple to 1 values.")));
+}
+
 TEST(TypecheckV2Test, ImportParametricFunctionWithDefaultExpression) {
   constexpr std::string_view kImported = R"(
 pub fn some_function<N: u32, M: u32 = {N + 1}>() -> uN[M] { uN[M]:0 }