From dc5c2444d431a74962794e2db96b432073fe6941 Mon Sep 17 00:00:00 2001
From: rsetaluri <raj.setaluri@gmail.com>
Date: Wed, 26 Jul 2023 11:09:28 -0700
Subject: [PATCH] [CompileGuard] Fix compile guard select key parsing

See #1291. Fixes two-key case (previously unsupported), and adds and refactors
tests.
---
 magma/compile_guard.py                        |  6 +-
 .../gold/test_compile_guard_select_two_keys.v | 68 +++++++++++++++++++
 tests/test_compile_guard.py                   | 45 ++++++++++--
 3 files changed, 109 insertions(+), 10 deletions(-)
 create mode 100644 tests/gold/test_compile_guard_select_two_keys.v

diff --git a/magma/compile_guard.py b/magma/compile_guard.py
index 60903ba8b..3a9184e8a 100644
--- a/magma/compile_guard.py
+++ b/magma/compile_guard.py
@@ -172,9 +172,9 @@ def compile_guard_select(**kwargs):
     try:
         default = kwargs.pop("default")
     except KeyError:
-        raise ValueError("Expected default argument") from None
-    if not (len(kwargs) > 1):
-        raise ValueError("Expected at least one key besides default")
+        raise KeyError("Expected default argument") from None
+    if not kwargs:  # kwargs is empty
+        raise KeyError("Expected at least one key besides default")
     # We rely on insertion order to make the default the last element for the
     # generated if/elif/else code.
     kwargs["default"] = default
diff --git a/tests/gold/test_compile_guard_select_two_keys.v b/tests/gold/test_compile_guard_select_two_keys.v
new file mode 100644
index 000000000..b14c0f922
--- /dev/null
+++ b/tests/gold/test_compile_guard_select_two_keys.v
@@ -0,0 +1,68 @@
+module coreir_reg #(
+    parameter width = 1,
+    parameter clk_posedge = 1,
+    parameter init = 1
+) (
+    input clk,
+    input [width-1:0] in,
+    output [width-1:0] out
+);
+  reg [width-1:0] outReg=init;
+  wire real_clk;
+  assign real_clk = clk_posedge ? clk : ~clk;
+  always @(posedge real_clk) begin
+    outReg <= in;
+  end
+  assign out = outReg;
+endmodule
+
+module Register (
+    input I,
+    output O,
+    input CLK
+);
+wire [0:0] reg_P1_inst0_out;
+coreir_reg #(
+    .clk_posedge(1'b1),
+    .init(1'h0),
+    .width(1)
+) reg_P1_inst0 (
+    .clk(CLK),
+    .in(I),
+    .out(reg_P1_inst0_out)
+);
+assign O = reg_P1_inst0_out[0];
+endmodule
+
+module CompileGuardSelect_Bit_COND1_default (
+    input I0,
+    input I1,
+    output O
+);
+`ifdef COND1
+    assign O = I0;
+`else
+    assign O = I1;
+`endif
+endmodule
+
+module _Top (
+    input I,
+    output O,
+    input CLK
+);
+wire Register_inst0_O;
+wire magma_Bit_xor_inst0_out;
+CompileGuardSelect_Bit_COND1_default CompileGuardSelect_Bit_COND1_default_inst0 (
+    .I0(Register_inst0_O),
+    .I1(I),
+    .O(O)
+);
+Register Register_inst0 (
+    .I(magma_Bit_xor_inst0_out),
+    .O(Register_inst0_O),
+    .CLK(CLK)
+);
+assign magma_Bit_xor_inst0_out = I ^ 1'b1;
+endmodule
+
diff --git a/tests/test_compile_guard.py b/tests/test_compile_guard.py
index 7906d00da..b7dc381cb 100644
--- a/tests/test_compile_guard.py
+++ b/tests/test_compile_guard.py
@@ -253,20 +253,50 @@ class _Top(m.Circuit):
 
 
 def test_compile_guard_select_basic():
+
     class _Top(m.Circuit):
         io = m.IO(I=m.In(m.Bit), O=m.Out(m.Bit)) + m.ClockIO()
-
         x = m.Register(m.Bit)()(io.I ^ 1)
         y = m.Register(m.Bit)()(io.I)
-
-        io.O @= m.compile_guard_select(
-            COND1=x, COND2=y, default=io.I
-        )
+        io.O @= m.compile_guard_select(COND1=x, COND2=y, default=io.I)
 
     basename = "test_compile_guard_select_basic"
     m.compile(f"build/{basename}", _Top, inline=True)
     assert m.testing.check_files_equal(
-        __file__, f"build/{basename}.v", f"gold/{basename}.v")
+        __file__, f"build/{basename}.v", f"gold/{basename}.v"
+    )
+
+
+def test_compile_guard_select_two_keys():
+
+    class _Top(m.Circuit):
+        io = m.IO(I=m.In(m.Bit), O=m.Out(m.Bit)) + m.ClockIO()
+        x = m.Register(m.Bit)()(io.I ^ 1)
+        io.O @= m.compile_guard_select(COND1=x, default=io.I)
+
+    basename = "test_compile_guard_select_two_keys"
+    m.compile(f"build/{basename}", _Top, inline=True)
+    assert m.testing.check_files_equal(
+        __file__, f"build/{basename}.v", f"gold/{basename}.v"
+    )
+
+
+def test_compile_guard_select_no_default():
+    with pytest.raises(KeyError):
+
+        class _Top(m.Circuit):
+            io = m.IO(I=m.In(m.Bit), O=m.Out(m.Bit)) + m.ClockIO()
+            x = m.Register(m.Bit)()(io.I ^ 1)
+            io.O @= m.compile_guard_select(COND1=x)
+
+
+def test_compile_guard_select_only_default():
+    with pytest.raises(KeyError):
+
+        class _Top(m.Circuit):
+            io = m.IO(I=m.In(m.Bit), O=m.Out(m.Bit)) + m.ClockIO()
+            x = m.Register(m.Bit)()(io.I ^ 1)
+            io.O @= m.compile_guard_select(default=x)
 
 
 def test_compile_guard_select_complex_type():
@@ -277,7 +307,8 @@ def make_top():
         class _Top(m.Circuit):
             io = m.IO(I0=m.In(T), I1=m.In(T), O=m.Out(T))
             io.O @= m.compile_guard_select(
-                COND1=io.I0, COND2=io.I1, default=io.I0)
+                COND1=io.I0, COND2=io.I1, default=io.I0
+            )
 
     with pytest.raises(TypeError):
         make_top()