Ensure Channel-legalization can handle procs with tokens in their state performing selects.

This can cause selects containing tokens to be seen by channel-legalization which previously confused the z3 translator.

Fixes: #3738
PiperOrigin-RevId: 861807182

Author: allight

xls/passes/BUILD

@@ -3888,9 +3888,11 @@ cc_test(
         "//xls/ir:channel",
         "//xls/ir:function_builder",
         "//xls/ir:ir_parser",
+        "//xls/ir:ir_test_base",
         "//xls/ir:proc_elaboration",
         "//xls/ir:value",
         "//xls/ir:verifier",
+        "//xls/solvers:z3_ir_equivalence_testutils",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",

xls/passes/channel_legalization_pass_test.cc

@@ -44,12 +44,14 @@
 #include "xls/ir/channel.h"
 #include "xls/ir/function_builder.h"
 #include "xls/ir/ir_parser.h"
+#include "xls/ir/ir_test_base.h"
 #include "xls/ir/package.h"
 #include "xls/ir/proc_elaboration.h"
 #include "xls/ir/value.h"
 #include "xls/ir/verifier.h"
 #include "xls/passes/optimization_pass.h"
 #include "xls/passes/pass_base.h"
+#include "xls/solvers/z3_ir_equivalence_testutils.h"
 
 namespace xls {
 namespace {
@@ -98,8 +100,7 @@ struct TestParam {
   };
 };
 
-class ChannelLegalizationPassTest
-    : public TestWithParam<std::tuple<TestParam, ChannelStrictness>> {
+class ChannelLegalizationPassIrTest : public IrTestBase {
  protected:
   absl::StatusOr<bool> Run(Package* package) {
     PassResults results;
@@ -111,6 +112,11 @@ class ChannelLegalizationPassTest
   }
 };
 
+class ChannelLegalizationPassTest
+    : public testing::WithParamInterface<
+          std::tuple<TestParam, ChannelStrictness>>,
+      public ChannelLegalizationPassIrTest {};
+
 TestParam kTestParameters[] = {
     TestParam{
         .test_name = "SingleProcBackToBackDataSwitchingOps",
@@ -995,6 +1001,29 @@ TEST_P(ChannelLegalizationPassTest, EvaluatesCorrectly) {
                     .evaluate(interpreter.get(), std::get<1>(GetParam())));
 }
 
+TEST_F(ChannelLegalizationPassIrTest, LegalizeWithTokenSel) {
+  auto p = CreatePackage();
+  ProcBuilder pb(NewStyleProc{}, TestName(), p.get());
+  XLS_ASSERT_OK_AND_ASSIGN(
+      auto chan_out,
+      pb.AddOutputChannel("out", p->GetBitsType(32), ChannelKind::kStreaming,
+                          ChannelStrictness::kRuntimeMutuallyExclusive));
+  auto st = pb.StateElement("state", UBits(0, 1));
+  auto tok = pb.StateElement("tok", Value::Token());
+  auto not_st = pb.Not(st);
+  auto tok_new_a1 = pb.SendIf(chan_out, tok, st, pb.Literal(UBits(32, 32)));
+  auto tok_new_b1 = pb.SendIf(chan_out, tok, not_st, pb.Literal(UBits(12, 32)));
+  pb.Next(st, not_st);
+  // NB This can come about from fairly normal dslx see:
+  // https://github.com/google/xls/issues/3738
+  pb.Next(tok, pb.Select(st, tok_new_a1, tok_new_b1));
+  XLS_ASSERT_OK_AND_ASSIGN(auto f, pb.Build());
+  XLS_ASSERT_OK(p->SetTop(f));
+  solvers::z3::ScopedVerifyProcEquivalence svpe(f, /*activation_count=*/5);
+  ScopedRecordIr sri(p.get());
+  EXPECT_THAT(Run(p.get()), IsOkAndHolds(true));
+}
+
 INSTANTIATE_TEST_SUITE_P(
     ChannelLegalizationPassTestInstantiation, ChannelLegalizationPassTest,
     Combine(ValuesIn(kTestParameters),

xls/solvers/z3_ir_translator.cc

@@ -1427,27 +1427,31 @@ absl::StatusOr<LeafTypeTree<Z3_ast>> IrTranslator::ToLeafTypeTree(Type* type,
 
 absl::StatusOr<Z3_ast> IrTranslator::FromLeafTypeTree(
     LeafTypeTreeView<Z3_ast> ast) {
-  if (ast.type()->IsBits()) {
-    return ast.Get({});
-  }
   Type* ty = ast.type();
-  if (ty->IsArray()) {
-    std::vector<Z3_ast> elements;
-    elements.reserve(ty->AsArrayOrDie()->size());
-    for (int64_t i = 0; i < ty->AsArrayOrDie()->size(); ++i) {
-      XLS_ASSIGN_OR_RETURN(Z3_ast elem, FromLeafTypeTree(ast.AsView({i})));
-      elements.push_back(elem);
+  switch (ty->kind()) {
+    case TypeKind::kToken:
+    case TypeKind::kBits:
+      // Already a leaf value.
+      return ast.Get({});
+    case TypeKind::kTuple: {
+      std::vector<Z3_ast> elements;
+      elements.reserve(ty->AsTupleOrDie()->size());
+      for (int64_t i = 0; i < ty->AsTupleOrDie()->size(); ++i) {
+        XLS_ASSIGN_OR_RETURN(Z3_ast elem, FromLeafTypeTree(ast.AsView({i})));
+        elements.push_back(elem);
+      }
+      return CreateTuple(ty->AsTupleOrDie(), elements);
+    }
+    case TypeKind::kArray: {
+      std::vector<Z3_ast> elements;
+      elements.reserve(ty->AsArrayOrDie()->size());
+      for (int64_t i = 0; i < ty->AsArrayOrDie()->size(); ++i) {
+        XLS_ASSIGN_OR_RETURN(Z3_ast elem, FromLeafTypeTree(ast.AsView({i})));
+        elements.push_back(elem);
+      }
+      return CreateArray(ty->AsArrayOrDie(), elements);
     }
-    return CreateArray(ty->AsArrayOrDie(), elements);
-  }
-  XLS_RET_CHECK(ty->IsTuple());
-  std::vector<Z3_ast> elements;
-  elements.reserve(ty->AsTupleOrDie()->size());
-  for (int64_t i = 0; i < ty->AsTupleOrDie()->size(); ++i) {
-    XLS_ASSIGN_OR_RETURN(Z3_ast elem, FromLeafTypeTree(ast.AsView({i})));
-    elements.push_back(elem);
   }
-  return CreateTuple(ty->AsTupleOrDie(), elements);
 }
 
 absl::Status IrTranslator::HandleAndReduce(BitwiseReductionOp* and_reduce) {