[CompileGuard] Exclude default from keys

Author: rsetaluri

magma/compile_guard.py

@@ -140,15 +140,15 @@ def _is_simple_type(T: Kind) -> bool:
 
 class _CompileGuardSelect(Generator2):
     def __init__(self, T: Kind, keys: Tuple[str]):
+        assert keys  # not empty
         self.T = T
         self.keys = keys
-        assert len(self.keys) > 1
         self.io = IO(
-            **{f"I{i}": In(T) for i in range(len(self.keys))},
+            **{f"I{i}": In(T) for i in range(len(self.keys) + 1)},
             O=Out(T),
         )
         T_str = type_to_sanitized_string(T)
-        self.name = f"CompileGuardSelect_{T_str}_{'_'.join(keys)}"
+        self.name = f"CompileGuardSelect_{T_str}_{'_'.join(keys)}_default"
         self.primitive = True
 
     def elaborate(self):
@@ -161,17 +161,9 @@ def elaborate(self):
             raise TypeError(f"Unsupported type: {self.T}")
         self.verilog = ""
         for i, key in enumerate(self.keys):
-            if i == 0:
-                stmt = f"`ifdef {key}"
-            elif key == "default":
-                assert i == (len(self.keys) - 1)
-                stmt = "`else"
-            else:
-                stmt = f"`elsif {key}"
-            self.verilog += f"""\
-{stmt}
-    assign O = I{i};
-"""
+            pred = f"`ifdef {key}" if i == 0 else f"`elsif {key}"
+            self.verilog += f"{pred}\n    assign O = I{i};\n"
+        self.verilog += f"`else\n    assign O = I{len(self.keys)};\n"
         self.verilog += "`endif"
 
 
@@ -185,9 +177,7 @@ def compile_guard_select(**kwargs):
         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
-    T, _ = infer_mux_type(list(kwargs.values()))
+    values = tuple(kwargs.values()) + (default,)
+    T, values = infer_mux_type(values)
     Select = _CompileGuardSelect(T, tuple(kwargs.keys()))
-    return Select()(*kwargs.values())
+    return Select()(*values)