From cbc1da9cf8c2e806d16c06598e655e1cad57841a Mon Sep 17 00:00:00 2001
From: Sean Purser-Haskell <seanhaskell@google.com>
Date: Thu, 15 Feb 2024 08:33:28 -0800
Subject: [PATCH] Implement conditional breaks in switch cases.

PiperOrigin-RevId: 607343929
---
 xls/contrib/xlscc/translator.cc               | 13 +++++---
 .../xlscc/unit_tests/translator_logic_test.cc | 31 ++++++++++++++++---
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/xls/contrib/xlscc/translator.cc b/xls/contrib/xlscc/translator.cc
index 759457162a..3ac9a1e8a4 100644
--- a/xls/contrib/xlscc/translator.cc
+++ b/xls/contrib/xlscc/translator.cc
@@ -5125,12 +5125,17 @@ absl::Status Translator::GenerateIR_Stmt(const clang::Stmt* stmt,
       // We use the original condition because we only care about
       //  enclosing conditions, such as if(...) { break; }
       //  Not if(...) {return;} break;
-      if (context().full_condition_on_enter_block.node() !=
+      if (context().full_condition_on_enter_block.node() ==
           context().full_switch_cond.node()) {
-        return absl::UnimplementedError(
-            ErrorMessage(loc, "Conditional breaks are not supported"));
+        context().hit_break = true;
+      } else {
+        context().relative_break_condition =
+            context().relative_condition_bval(loc);
+
+        // Make the rest of the block no-op
+        XLS_RETURN_IF_ERROR(
+            and_condition(context().fb->Literal(xls::UBits(0, 1), loc), loc));
       }
-      context().hit_break = true;
     }
   } else if (clang::isa<clang::CompoundStmt>(stmt)) {
     PushContextGuard context_guard(*this, loc);
diff --git a/xls/contrib/xlscc/unit_tests/translator_logic_test.cc b/xls/contrib/xlscc/unit_tests/translator_logic_test.cc
index c659ed6b95..6e3a74cf88 100644
--- a/xls/contrib/xlscc/unit_tests/translator_logic_test.cc
+++ b/xls/contrib/xlscc/unit_tests/translator_logic_test.cc
@@ -955,6 +955,29 @@ TEST_F(TranslatorLogicTest, SwitchStmtWithUnrollInCase) {
   Run({{"a", 3}}, 300, content);
 }
 
+TEST_F(TranslatorLogicTest, SwitchConditionalBreakReturn) {
+  std::string_view content = R"(
+       long long my_package(long long a, long long b) {
+         long long ret;
+         switch(a) {
+           case 1:
+             ret = 100;
+             break;
+           case 2:
+             ret = 200;
+             if(b) return 55;
+           default:
+             return 300;
+         }
+         return ret;
+       })";
+
+  Run({{"a", 1}, {"b", 0}}, 100, content);
+  Run({{"a", 2}, {"b", 0}}, 300, content);
+  Run({{"a", 2}, {"b", 1}}, 55, content);
+  Run({{"a", 3}, {"b", 0}}, 300, content);
+}
+
 TEST_F(TranslatorLogicTest, SwitchConditionalBreak) {
   std::string_view content = R"(
        long long my_package(long long a, long long b) {
@@ -973,10 +996,10 @@ TEST_F(TranslatorLogicTest, SwitchConditionalBreak) {
          return ret;
        })";
 
-  ASSERT_THAT(SourceToIr(content).status(),
-              xls::status_testing::StatusIs(
-                  absl::StatusCode::kUnimplemented,
-                  testing::HasSubstr("Conditional breaks are not supported")));
+  Run({{"a", 1}, {"b", 0}}, 100, content);
+  Run({{"a", 2}, {"b", 0}}, 300, content);
+  Run({{"a", 2}, {"b", 1}}, 200, content);
+  Run({{"a", 3}, {"b", 0}}, 300, content);
 }
 
 TEST_F(TranslatorLogicTest, SwitchStmtDefaultTop) {