From: whitequark <cz@m-labs.hk>
Date: Wed, 10 Apr 2019 00:23:11 +0000 (+0000)
Subject: hdl.xfrm: allow using FragmentTransformer on any elaboratable.
X-Git-Url: https://git.libre-soc.org/?a=commitdiff_plain;h=d64ffeffe463fb9bb5ca19f47f2cf70558beb66e;p=nmigen.git

hdl.xfrm: allow using FragmentTransformer on any elaboratable.

Fixes #29.
---

diff --git a/nmigen/hdl/ir.py b/nmigen/hdl/ir.py
index 50b23d0..079bba2 100644
--- a/nmigen/hdl/ir.py
+++ b/nmigen/hdl/ir.py
@@ -18,7 +18,7 @@ class Fragment:
     def get(obj, platform):
         if isinstance(obj, Fragment):
             return obj
-        if hasattr(obj, "elaborate"):
+        elif hasattr(obj, "elaborate"):
             frag = obj.elaborate(platform)
         else:
             raise AttributeError("Object '{!r}' cannot be elaborated".format(obj))
diff --git a/nmigen/hdl/xfrm.py b/nmigen/hdl/xfrm.py
index e5f415e..7c51a61 100644
--- a/nmigen/hdl/xfrm.py
+++ b/nmigen/hdl/xfrm.py
@@ -13,6 +13,7 @@ from .rec import *
 __all__ = ["ValueVisitor", "ValueTransformer",
            "StatementVisitor", "StatementTransformer",
            "FragmentTransformer",
+           "TransformedElaboratable",
            "DomainRenamer", "DomainLowerer",
            "SampleDomainInjector", "SampleLowerer",
            "SwitchCleaner", "LHSGroupAnalyzer", "LHSGroupFilter",
@@ -277,7 +278,36 @@ class FragmentTransformer:
         return new_fragment
 
     def __call__(self, value):
-        return self.on_fragment(value)
+        if isinstance(value, Fragment):
+            return self.on_fragment(value)
+        elif isinstance(value, TransformedElaboratable):
+            value._transforms_.append(self)
+            return value
+        elif hasattr(value, "elaborate"):
+            value = TransformedElaboratable(value)
+            value._transforms_.append(self)
+            return value
+        else:
+            raise AttributeError("Object '{!r}' cannot be elaborated".format(value))
+
+
+class TransformedElaboratable:
+    def __init__(self, elaboratable):
+        assert hasattr(elaboratable, "elaborate")
+
+        # Fields prefixed and suffixed with underscore to avoid as many conflicts with the inner
+        # object as possible, since we're forwarding attribute requests to it.
+        self._elaboratable_ = elaboratable
+        self._transforms_   = []
+
+    def __getattr__(self, attr):
+        return getattr(self._elaboratable_, attr)
+
+    def elaborate(self, platform):
+        fragment = Fragment.get(self._elaboratable_, platform)
+        for transform in self._transforms_:
+            fragment = transform(fragment)
+        return fragment
 
 
 class DomainRenamer(FragmentTransformer, ValueTransformer, StatementTransformer):
diff --git a/nmigen/test/test_hdl_xfrm.py b/nmigen/test/test_hdl_xfrm.py
index f983ab7..94af1f8 100644
--- a/nmigen/test/test_hdl_xfrm.py
+++ b/nmigen/test/test_hdl_xfrm.py
@@ -486,3 +486,49 @@ class CEInserterTestCase(FHDLTestCase):
             )
         )
         """)
+
+
+class _MockElaboratable:
+    def __init__(self):
+        self.s1 = Signal()
+
+    def elaborate(self, platform):
+        f = Fragment()
+        f.add_statements(
+            self.s1.eq(1)
+        )
+        f.add_driver(self.s1, "sync")
+        return f
+
+
+class TransformedElaboratableTestCase(FHDLTestCase):
+    def setUp(self):
+        self.c1 = Signal()
+        self.c2 = Signal()
+
+    def test_getattr(self):
+        e = _MockElaboratable()
+        te = CEInserter(self.c1)(e)
+
+        self.assertIs(te.s1, e.s1)
+
+    def test_composition(self):
+        e = _MockElaboratable()
+        te1 = CEInserter(self.c1)(e)
+        te2 = ResetInserter(self.c2)(te1)
+
+        self.assertIsInstance(te1, TransformedElaboratable)
+        self.assertIs(te1, te2)
+
+        f = Fragment.get(te2, None)
+        self.assertRepr(f.statements, """
+        (
+            (eq (sig s1) (const 1'd1))
+            (switch (sig c1)
+                (case 0 (eq (sig s1) (sig s1)))
+            )
+            (switch (sig c2)
+                (case 1 (eq (sig s1) (const 1'd0)))
+            )
+        )
+        """)