From 2553f667173a52146e21553a352c6aeb09e667b5 Mon Sep 17 00:00:00 2001
From: whitequark <cz@m-labs.hk>
Date: Sat, 25 May 2019 20:09:26 +0000
Subject: [PATCH] Consider Instances a part of containing fragment for use-def
 purposes.

Fixes #70.
---
 nmigen/hdl/ir.py           | 43 ++++++++++++++++++++------------------
 nmigen/test/test_hdl_ir.py | 16 ++++++++++++--
 2 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/nmigen/hdl/ir.py b/nmigen/hdl/ir.py
index 7884e1d..c462e7e 100644
--- a/nmigen/hdl/ir.py
+++ b/nmigen/hdl/ir.py
@@ -366,31 +366,34 @@ class Fragment:
 
         # Collect all signals we're driving (on LHS of statements), and signals we're using
         # (on RHS of statements, or in clock domains).
-        if isinstance(self, Instance):
-            for port_name, (value, dir) in self.named_ports.items():
-                if dir == "i":
-                    add_uses(value._rhs_signals())
-                if dir == "o":
-                    add_defs(value._lhs_signals())
-                if dir == "io":
-                    add_io(value)
-        else:
-            for stmt in self.statements:
-                add_uses(stmt._rhs_signals())
-                add_defs(stmt._lhs_signals())
+        for stmt in self.statements:
+            add_uses(stmt._rhs_signals())
+            add_defs(stmt._lhs_signals())
 
-            for domain, _ in self.iter_sync():
-                cd = self.domains[domain]
-                add_uses(cd.clk)
-                if cd.rst is not None:
-                    add_uses(cd.rst)
+        for domain, _ in self.iter_sync():
+            cd = self.domains[domain]
+            add_uses(cd.clk)
+            if cd.rst is not None:
+                add_uses(cd.rst)
 
         # Repeat for subfragments.
         for subfrag, name in self.subfragments:
-            parent[subfrag] = self
-            level [subfrag] = level[self] + 1
+            if isinstance(subfrag, Instance):
+                for port_name, (value, dir) in subfrag.named_ports.items():
+                    if dir == "i":
+                        subfrag.add_ports(value._rhs_signals(), dir=dir)
+                        add_uses(value._rhs_signals())
+                    if dir == "o":
+                        subfrag.add_ports(value._lhs_signals(), dir=dir)
+                        add_defs(value._lhs_signals())
+                    if dir == "io":
+                        subfrag.add_ports(value, dir=dir)
+                        add_io(value)
+            else:
+                parent[subfrag] = self
+                level [subfrag] = level[self] + 1
 
-            subfrag._prepare_use_def_graph(parent, level, uses, defs, ios, top)
+                subfrag._prepare_use_def_graph(parent, level, uses, defs, ios, top)
 
     def _propagate_ports(self, ports, all_undef_as_ports):
         # Take this fragment graph:
diff --git a/nmigen/test/test_hdl_ir.py b/nmigen/test/test_hdl_ir.py
index 26ddfd5..e2426cf 100644
--- a/nmigen/test/test_hdl_ir.py
+++ b/nmigen/test/test_hdl_ir.py
@@ -561,6 +561,8 @@ class InstanceTestCase(FHDLTestCase):
             o_data=Cat(self.datal, self.datah),
             io_pins=self.pins
         )
+        self.wrap = Fragment()
+        self.wrap.add_subfragment(self.inst)
 
     def test_init(self):
         self.setUp_cpu()
@@ -572,7 +574,7 @@ class InstanceTestCase(FHDLTestCase):
 
     def test_prepare(self):
         self.setUp_cpu()
-        f = self.inst.prepare()
+        f = self.wrap.prepare()
         sync_clk = f.domains["sync"].clk
         self.assertEqual(f.ports, SignalDict([
             (sync_clk, "i"),
@@ -582,7 +584,7 @@ class InstanceTestCase(FHDLTestCase):
 
     def test_prepare_explicit_ports(self):
         self.setUp_cpu()
-        f = self.inst.prepare(ports=[self.rst, self.stb])
+        f = self.wrap.prepare(ports=[self.rst, self.stb])
         sync_clk = f.domains["sync"].clk
         sync_rst = f.domains["sync"].rst
         self.assertEqual(f.ports, SignalDict([
@@ -592,3 +594,13 @@ class InstanceTestCase(FHDLTestCase):
             (self.stb, "o"),
             (self.pins, "io"),
         ]))
+
+    def test_prepare_slice_in_port(self):
+        s = Signal(2)
+        f = Fragment()
+        f.add_subfragment(Instance("foo", o_O=s[0]))
+        f.add_subfragment(Instance("foo", o_O=s[1]))
+        fp = f.prepare(ports=[s], ensure_sync_exists=False)
+        self.assertEqual(fp.ports, SignalDict([
+            (s, "o"),
+        ]))
-- 
2.30.2