From 3d80887417b1537ca7d6f344b0b4d034e13c46ff Mon Sep 17 00:00:00 2001
From: whitequark <cz@m-labs.hk>
Date: Mon, 14 Jan 2019 17:04:23 +0000
Subject: [PATCH] hdl.ir: allow explicitly requesting flattening.

---
 nmigen/hdl/ir.py           | 14 ++++++++++----
 nmigen/hdl/xfrm.py         |  1 +
 nmigen/lib/io.py           |  5 ++++-
 nmigen/test/test_hdl_ir.py |  9 +++++++++
 4 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/nmigen/hdl/ir.py b/nmigen/hdl/ir.py
index a4889b1..bf66135 100644
--- a/nmigen/hdl/ir.py
+++ b/nmigen/hdl/ir.py
@@ -21,6 +21,7 @@ class Fragment:
         self.domains = OrderedDict()
         self.subfragments = []
         self.generated = OrderedDict()
+        self.flatten = False
 
     def add_ports(self, *ports, dir):
         assert dir in ("i", "o", "io")
@@ -141,7 +142,16 @@ class Fragment:
         for domain, signal in self.iter_drivers():
             add_subfrag(driver_subfrags, signal, (None, hierarchy))
 
+        flatten_subfrags = set()
         for i, (subfrag, name) in enumerate(self.subfragments):
+            if name is None:
+                name = "<unnamed #{}>".format(i)
+            subfrag_hierarchy = hierarchy + (name,)
+
+            if subfrag.flatten:
+                # Always flatten subfragments that explicitly request it.
+                flatten_subfrags.add((subfrag, subfrag_hierarchy))
+
             if isinstance(subfrag, Instance):
                 # For memories (which are subfragments, but semantically a part of superfragment),
                 # record that this fragment is driving it.
@@ -153,9 +163,6 @@ class Fragment:
                 continue
 
             # First, recurse into subfragments and let them detect driver conflicts as well.
-            if name is None:
-                name = "<unnamed #{}>".format(i)
-            subfrag_hierarchy = hierarchy + (name,)
             subfrag_drivers, subfrag_memories = \
                 subfrag._resolve_hierarchy_conflicts(subfrag_hierarchy, mode)
 
@@ -167,7 +174,6 @@ class Fragment:
 
         # Find out the set of subfragments that needs to be flattened into this fragment
         # to resolve driver-driver conflicts.
-        flatten_subfrags = set()
         def flatten_subfrags_if_needed(subfrags):
             if len(subfrags) == 1:
                 return []
diff --git a/nmigen/hdl/xfrm.py b/nmigen/hdl/xfrm.py
index 20a325c..33c534e 100644
--- a/nmigen/hdl/xfrm.py
+++ b/nmigen/hdl/xfrm.py
@@ -237,6 +237,7 @@ class FragmentTransformer:
             self.map_named_ports(fragment, new_fragment)
         else:
             new_fragment = Fragment()
+            new_fragment.flatten = fragment.flatten
         self.map_ports(fragment, new_fragment)
         self.map_subfragments(fragment, new_fragment)
         self.map_domains(fragment, new_fragment)
diff --git a/nmigen/lib/io.py b/nmigen/lib/io.py
index 1801f5f..313c8b7 100644
--- a/nmigen/lib/io.py
+++ b/nmigen/lib/io.py
@@ -41,4 +41,7 @@ class Tristate:
             i_A=self.triple.o,
             o_Y=self.io,
         )
-        return m.lower(platform)
+
+        f = m.lower(platform)
+        f.flatten = True
+        return f
diff --git a/nmigen/test/test_hdl_ir.py b/nmigen/test/test_hdl_ir.py
index 25956d4..ffe1f63 100644
--- a/nmigen/test/test_hdl_ir.py
+++ b/nmigen/test/test_hdl_ir.py
@@ -509,6 +509,15 @@ class FragmentHierarchyConflictTestCase(FHDLTestCase):
                     "top.<unnamed #1>; hierarchy will be flattened"):
             self.f1._resolve_hierarchy_conflicts(mode="warn")
 
+    def test_explicit_flatten(self):
+        self.f1 = Fragment()
+        self.f2 = Fragment()
+        self.f2.flatten = True
+        self.f1.add_subfragment(self.f2)
+
+        self.f1._resolve_hierarchy_conflicts(mode="silent")
+        self.assertEqual(self.f1.subfragments, [])
+
 
 class InstanceTestCase(FHDLTestCase):
     def setUp_cpu(self):
-- 
2.30.2