From e0b54b417e60ad6d59c97c21a02be5887435de27 Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Sat, 3 Aug 2019 18:19:40 +0000
Subject: [PATCH] hdl.ir: allow adding more than one domain in missing domain
 callback.

This is useful for injecting complex power-on reset logic.
---
 nmigen/hdl/ir.py           | 22 ++++++++++------------
 nmigen/test/test_hdl_ir.py | 23 +++++++++++++++++++++--
 2 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/nmigen/hdl/ir.py b/nmigen/hdl/ir.py
index 370b255..6dd5d08 100644
--- a/nmigen/hdl/ir.py
+++ b/nmigen/hdl/ir.py
@@ -363,23 +363,21 @@ class Fragment:
                 if value is None:
                     raise DomainError("Domain '{}' is used but not defined".format(domain_name))
                 if type(value) is ClockDomain:
-                    domain = value
-                    # Only expose ports on clock domains returned directly, i.e. not as a part of
-                    # a fragment driving that domain.
-                    new_domains.append(domain)
+                    self.add_domains(value)
+                    # And expose ports on the newly added clock domain, since it is added directly
+                    # and there was no chance to add any logic driving it.
+                    new_domains.append(value)
                 else:
                     new_fragment = Fragment.get(value, platform=None)
-                    if new_fragment.domains.keys() != {domain_name}:
+                    if domain_name not in new_fragment.domains:
+                        defined = new_fragment.domains.keys()
                         raise DomainError(
-                            "Fragment returned by missing domain callback should define exactly "
-                            "one domain '{}', but defines domain(s) {}."
-                            .format(domain_name,
-                                    ", ".join("'{}'".format(n)
-                                              for n in new_fragment.domains.keys())))
+                            "Fragment returned by missing domain callback does not define "
+                            "requested domain '{}' (defines {})."
+                            .format(domain_name, ", ".join("'{}'".format(n) for n in defined)))
                     new_fragment.flatten = True
                     self.add_subfragment(new_fragment)
-                    domain = new_fragment.domains[domain_name]
-                self.add_domains(domain)
+                    self.add_domains(new_fragment.domains.values())
         return new_domains
 
     def _propagate_domains(self, missing_domain):
diff --git a/nmigen/test/test_hdl_ir.py b/nmigen/test/test_hdl_ir.py
index d6c43be..c89e8d7 100644
--- a/nmigen/test/test_hdl_ir.py
+++ b/nmigen/test/test_hdl_ir.py
@@ -421,6 +421,25 @@ class FragmentDomainsTestCase(FHDLTestCase):
         ])
         self.assertTrue(f2.flatten)
 
+    def test_propagate_create_missing_fragment_many_domains(self):
+        s1 = Signal()
+        f1 = Fragment()
+        f1.add_driver(s1, "sync")
+
+        cd_por  = ClockDomain("por")
+        cd_sync = ClockDomain("sync")
+        f2 = Fragment()
+        f2.add_domains(cd_por, cd_sync)
+
+        new_domains = f1._propagate_domains(missing_domain=lambda name: f2)
+        self.assertEqual(f1.domains.keys(), {"sync", "por"})
+        self.assertEqual(f2.domains.keys(), {"sync", "por"})
+        self.assertEqual(f1.domains["sync"], f2.domains["sync"])
+        self.assertEqual(new_domains, [])
+        self.assertEqual(f1.subfragments, [
+            (f2, None)
+        ])
+
     def test_propagate_create_missing_fragment_wrong(self):
         s1 = Signal()
         f1 = Fragment()
@@ -430,8 +449,8 @@ class FragmentDomainsTestCase(FHDLTestCase):
         f2.add_domains(ClockDomain("foo"))
 
         with self.assertRaises(DomainError,
-                msg="Fragment returned by missing domain callback should define exactly "
-                    "one domain 'sync', but defines domain(s) 'foo'."):
+                msg="Fragment returned by missing domain callback does not define requested "
+                    "domain 'sync' (defines 'foo')."):
             f1._propagate_domains(missing_domain=lambda name: f2)
 
 
-- 
2.30.2