From 2e1a38e0aa6aa06ab2ac16e755da2be7ce1ad6b9 Mon Sep 17 00:00:00 2001
From: jacobdweightman <jacob@risczero.com>
Date: Tue, 18 Feb 2025 15:21:02 -0600
Subject: [PATCH] Fix semantics of if without else

---
 zirgen/circuit/rv32im/v2/dsl/inst_control.zir |  4 ---
 zirgen/dsl/parser.cpp                         | 18 ++++++++----
 zirgen/dsl/test/parser.cpp                    |  8 +++--
 zirgen/dsl/test/repro/if_without_else.zir     | 29 +++++++++++++++++++
 4 files changed, 47 insertions(+), 12 deletions(-)
 create mode 100644 zirgen/dsl/test/repro/if_without_else.zir

diff --git a/zirgen/circuit/rv32im/v2/dsl/inst_control.zir b/zirgen/circuit/rv32im/v2/dsl/inst_control.zir
index 96a8e899..ccf575c0 100644
--- a/zirgen/circuit/rv32im/v2/dsl/inst_control.zir
+++ b/zirgen/circuit/rv32im/v2/dsl/inst_control.zir
@@ -90,10 +90,6 @@ component ControlSuspend(cycle: Reg, input: InstInput) {
       global termA0high := Reg(0);
       global termA1low := Reg(0);
       global termA1high := Reg(0);
-    } else {
-      // Do nothing
-      // TODO: Apparently if we don't have this clause, the compiler generates
-      // a mux where the 'else' case fails at runtime.
     };
 
     // Begin page out
diff --git a/zirgen/dsl/parser.cpp b/zirgen/dsl/parser.cpp
index 6efd00ca..3722a68e 100644
--- a/zirgen/dsl/parser.cpp
+++ b/zirgen/dsl/parser.cpp
@@ -796,13 +796,21 @@ Switch::Ptr Parser::parseConditional() {
 
   cases.push_back(parseBlock());
   selectors.push_back(condition);
+
+  // Build else clause, implicit empty one if none is provided
+  Expression::Vec subArgs;
+  subArgs.push_back(make_shared<Literal>(location, 1));
+  subArgs.push_back(condition);
+  selectors.push_back(
+      make_shared<Construct>(location, make_shared<Ident>(location, "Sub"), subArgs));
+
   if (lexer.takeTokenIf(tok_else)) {
     cases.push_back(parseBlock());
-    Expression::Vec subArgs;
-    subArgs.push_back(make_shared<Literal>(location, 1));
-    subArgs.push_back(condition);
-    selectors.push_back(
-        make_shared<Construct>(location, make_shared<Ident>(location, "Sub"), subArgs));
+  } else {
+    Ident::Ptr component = make_shared<Ident>(location, "Component");
+    auto trivial = make_shared<Construct>(location, std::move(component), Expression::Vec{});
+    Block::Ptr block = make_shared<Block>(location, Statement::Vec{}, std::move(trivial));
+    cases.push_back(block);
   }
   Expression::Ptr selector = make_shared<ArrayLiteral>(location, std::move(selectors));
   return make_shared<Switch>(location, std::move(selector), std::move(cases), false);
diff --git a/zirgen/dsl/test/parser.cpp b/zirgen/dsl/test/parser.cpp
index 0ecd74e2..a0ddfa5c 100644
--- a/zirgen/dsl/test/parser.cpp
+++ b/zirgen/dsl/test/parser.cpp
@@ -1,4 +1,4 @@
-// Copyright 2024 RISC Zero, Inc.
+// Copyright 2025 RISC Zero, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -145,8 +145,10 @@ TEST(parser, conditional_if_else) {
 TEST(parser, conditional_if) {
   TestParser parser("if (x) { y }");
 
-  // desuggars to [x] -> ({ y })"
-  auto expected = Switch(ArrayLiteral(arr(Ident("x"))), arr(Block({}, Ident("y"))));
+  // desugars to [x, 1-x] -> ({ y }, {})"
+  auto expected =
+      Switch(ArrayLiteral(arr(Ident("x"), Construct(Ident("Sub"), arr(Literal(1), Ident("x"))))),
+             arr(Block({}, Ident("y")), Block({}, Construct(Ident("Component"), {}))));
   EXPECT_EQ(*parser.parseExpression(), *expected);
 }
 
diff --git a/zirgen/dsl/test/repro/if_without_else.zir b/zirgen/dsl/test/repro/if_without_else.zir
new file mode 100644
index 00000000..12799280
--- /dev/null
+++ b/zirgen/dsl/test/repro/if_without_else.zir
@@ -0,0 +1,29 @@
+// RUN: zirgen %s --emit=zhl 2>&1 | FileCheck %s
+// RUN: zirgen %s --test --test-cycles=4 2>&1 | FileCheck %s --check-prefix=EXEC
+
+// This test covers a bug where an `if` without an `else` resulted in a mux with
+// only one arm. This is incorrect, as these three expressions should all be
+// semantically equivalent:
+
+// if (x) {a}
+// if (x) {a} else {}
+// [x, 1-x] -> ({a}, {})
+
+extern IsFirstCycle() : Val;
+
+component Top() {
+  x := NondetReg(IsFirstCycle());
+  if (x) {
+    Log("First cycle")
+  }
+}
+
+// CHECK: %[[arr:[0-9]+]] = zhl.array[%{{[0-9]+}}, %{{[0-9]+}}]
+// CHECK: %{{[0-9]+}} = zhl.switch %[[arr]] -> {
+
+// EXEC: [0] Log: First cycle
+// EXEC-NOT: Log: First cycle
+
+test {
+  Top()
+}