modified rebind_bitvector to support modifiers

leonardt · Aug 2, 2019 · 699ca30 · 699ca30
1 parent 87af791
commit 699ca30
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 10 deletions.
diff --git a/hwtypes/adt_util.py b/hwtypes/adt_util.py
@@ -2,12 +2,17 @@
 from .bit_vector_abc import AbstractBitVectorMeta, AbstractBitVector
 
 from .util import _issubclass
-from hwtypes.modifiers import get_all_modifiers
+from hwtypes.modifiers import unwrap_modifier, wrap_modifier, is_modified
 
 def rebind_bitvector(
         adt,
         bv_type_0: AbstractBitVectorMeta,
-        bv_type_1: AbstractBitVectorMeta):
+        bv_type_1: AbstractBitVectorMeta,
+        keep_modifiers=False):
+    if keep_modifiers and is_modified(adt):
+        unmod, mods = unwrap_modifier(adt)
+        return wrap_modifier(rebind_bitvector(unmod,bv_type_0,bv_type_1,True),mods)
+
     if _issubclass(adt, bv_type_0):
         if adt.is_sized:
             return bv_type_1[adt.size]
@@ -16,22 +21,23 @@ def rebind_bitvector(
     elif isinstance(adt, BoundMeta):
         new_adt = adt
         for field in adt.fields:
-            new_field = rebind_bitvector(field, bv_type_0, bv_type_1)
+            new_field = rebind_bitvector(field, bv_type_0, bv_type_1,keep_modifiers)
             new_adt = new_adt.rebind(field, new_field)
         return new_adt
     else:
         return adt
 
+def rebind_keep_modifiers(adt, A, B):
+    if is_modified(adt):
+        unmod, mods = unwrap_modifier(adt)
+        return wrap_modifier(rebind_keep_modifiers(unmod,A,B),mods)
 
-def rebind_keep_modifiers(adt, A, B, rebind_sub_types=False):
-    if A is B or (rebind_sub_types and _issubclass(adt, B)):
-        for mod in get_all_modifiers(A):
-            B = mod(B)
+    if _issubclass(adt,A):
         return B
     elif isinstance(adt, BoundMeta):
         new_adt = adt
         for field in adt.fields:
-            new_field = rebind_keep_modifiers(field, A, B, rebind_sub_types)
+            new_field = rebind_keep_modifiers(field, A, B)
             new_adt = new_adt.rebind(field, new_field)
         return new_adt
     else:

diff --git a/tests/test_rebind.py b/tests/test_rebind.py
@@ -1,10 +1,11 @@
 import pytest
 
 from hwtypes.adt import Product, Sum, Enum, Tuple
-from hwtypes.adt_util import rebind_bitvector
+from hwtypes.adt_util import rebind_bitvector, rebind_keep_modifiers
 from hwtypes.bit_vector import AbstractBitVector, BitVector, AbstractBit, Bit
-from hwtypes.smt_bit_vector import SMTBit
+from hwtypes.smt_bit_vector import SMTBit, SMTBitVector
 from hwtypes.util import _issubclass
+from hwtypes.modifiers import make_modifier
 
 class A: pass
 class B: pass
@@ -147,3 +148,15 @@ class A(Product):
 
     A_smt = A.rebind(AbstractBit, SMTBit, True)
     assert A_smt.a is SMTBit
+
+def test_rebind_mod():
+    M = make_modifier("M")
+    class A(Product):
+        a=M(Bit)
+        b=M(BitVector[4])
+
+    A_smt = rebind_bitvector(A,AbstractBitVector, SMTBitVector, True)
+    A_smt = rebind_keep_modifiers(A_smt, AbstractBit, SMTBit)
+    assert A_smt.b == M(SMTBitVector[4])
+    assert A_smt.a == M(SMTBit)
+test_rebind_mod()