Repository: phanrahan/magma

Issue:
Magma protocol losing type under ite
Bit.ite does not "cast" protocols back to their protocol type

Specifically:
```Python
assert type(proto0) == type(proto1)
v = Bit().ite(proto0, proto1)
assert type(v) == type(proto0)._to_magma_
```
where I would expect:
```Python
assert type(proto0) == type(proto1)
v = Bit().ite(proto0, proto1)
assert type(v) == type(proto0)
```

Discussion:

Patch:
diff --git a/magma/bit.py b/magma/bit.py
index 6e3cbd73..e67cef11 100644
--- a/magma/bit.py
+++ b/magma/bit.py
@@ -20,7 +20,7 @@ from magma.debug import debug_wire
 from magma.family import get_family
 from magma.interface import IO
 from magma.language_utils import primitive_to_python
-from magma.protocol_type import magma_type
+from magma.protocol_type import magma_type, MagmaProtocol
 from magma.operator_utils import output_only
 
 
@@ -161,7 +161,13 @@ class Bit(Digital, AbstractBit, metaclass=DigitalMeta):
         # Note: coreir flips t/f cases
         # self._mux monkey patched in magma/primitives/mux.py to avoid circular
         # dependency
-        return self._mux([f_branch, t_branch], self)
+        r = self._mux([f_branch, t_branch], self)
+
+        # cast protocol values back to protocol
+        if type(t_branch) == type(f_branch) and isinstance(t_branch, MagmaProtocol):
+            r = type(t_branch)._from_magma_value_(r)
+
+        return r
 
     def __bool__(self) -> bool:
         raise NotImplementedError("Converting magma bit to bool not supported")


Test patch:
diff --git a/tests/test_type/test_magma_protocol.py b/tests/test_type/test_magma_protocol.py
index eda13967..2ed9c768 100644
--- a/tests/test_type/test_magma_protocol.py
+++ b/tests/test_type/test_magma_protocol.py
@@ -5,56 +5,69 @@ from magma.testing import check_files_equal
 import sys
 import os
 
+_CACHE = {}
+class FooMeta(m.MagmaProtocolMeta):
+    def _to_magma_(cls):
+        # Need way to retrieve underlying magma type
+        return cls.T
+
+    def _qualify_magma_(cls, direction: m.Direction):
+        # Need way to create a new version (e.g. give me a Foo with the
+        # underlying type qualified to be an input)
+        return cls[cls.T.qualify(direction)]
+
+    def _flip_magma_(cls):
+        # Need way to create a new version (e.g. give me a Foo with the
+        # underlying type qualified to be an input)
+        return cls[cls.T.flip()]
+
+    def _from_magma_value_(cls, val: m.Type):
+        # Need a way to create an instance of Foo from a value, this just
+        # dispatches to the __init__ logic, but you could define any
+        # custom behavior here
+        return cls(val)
+
+    def __getitem__(cls, T):
+        return _CACHE.setdefault(T, type(cls)(f"Foo{T}", (cls, ), {"T": T}))
+
+class Foo(m.MagmaProtocol, metaclass=FooMeta):
+    def __init__(self, val: Optional[m.Type] = None):
+        if val is None:
+            val = self.T()
+        self._val = val
+
+    def _get_magma_value_(self):
+        # Need way to access underlying magma value
+        return self._val
+
+    def non_standard_operation(self):
+        v0 = self._val << 2
+        v1 = m.bits(self._val[0], len(self.T)) << 1
+        return Foo(v0 | v1 | m.bits(self._val[0], len(self.T)))
 
-def test_foo_type_magma_protocol():
-    class FooMeta(m.MagmaProtocolMeta):
-        def _to_magma_(cls):
-            # Need way to retrieve underlying magma type
-            return cls.T
-
-        def _qualify_magma_(cls, direction: m.Direction):
-            # Need way to create a new version (e.g. give me a Foo with the
-            # underlying type qualified to be an input)
-            return cls[cls.T.qualify(direction)]
-
-        def _flip_magma_(cls):
-            # Need way to create a new version (e.g. give me a Foo with the
-            # underlying type qualified to be an input)
-            return cls[cls.T.flip()]
-
-        def _from_magma_value_(cls, val: m.Type):
-            # Need a way to create an instance of Foo from a value, this just
-            # dispatches to the __init__ logic, but you could define any
-            # custom behavior here
-            return cls(val)
-
-        def __getitem__(cls, T):
-            return type(cls)(f"Foo{T}", (cls, ), {"T": T})
-
-    class Foo(m.MagmaProtocol, metaclass=FooMeta):
-        def __init__(self, val: Optional[m.Type] = None):
-            if val is None:
-                val = self.T()
-            self._val = val
-
-        def _get_magma_value_(self):
-            # Need way to access underlying magma value
-            return self._val
-
-        def non_standard_operation(self):
-            v0 = self._val << 2
-            v1 = m.bits(self._val[0], len(self.T)) << 1
-            return Foo(v0 | v1 | m.bits(self._val[0], len(self.T)))
 
+def test_foo_type_magma_protocol():
     @m.sequential2()
     class Bar:
         def __init__(self):
             self.reg = m.Register(m.Bits[1])()
 
-        def __call__(self, foo: Foo[m.Bits[8]]) -> m.Bits[8]:
+        def __call__(self, foo: Foo[m.Bits[8]]) -> Foo[m.Bits[8]]:
             self.reg = self.reg
             return foo.non_standard_operation()
 
     m.compile("build/test_foo_magma_protocol", Bar)
     assert check_files_equal(__file__, f"build/test_foo_magma_protocol.v",
                              f"gold/test_foo_magma_protocol.v")
+
+
+def test_ite():
+    @m.sequential2()
+    class Bar:
+        def __call__(self, cond: m.Bit, i0: Foo[m.Bits[8]], i1: Foo[m.Bits[8]]) -> Foo[m.Bits[8]]:
+            if cond:
+                foo = i0
+            else:
+                foo = i1
+
+            return foo.non_standard_operation()