From 3156c776227cd8e2a751e64903b451162d91f6e3 Mon Sep 17 00:00:00 2001
From: Lenny Truong <leonardtruong@fb.com>
Date: Tue, 11 Oct 2022 11:56:22 -0700
Subject: [PATCH] Fix rebase bug in aggregate wireables, add tests

I chose the wrong changes for these functions when rebasing the tuple2
code.  With the aggregate_wireable_method decorator, they only need to
implement the logic for when children are elaborated (the decorator
calls the aggregate method if the has_elaborated_children check fails).
This didn't trigger a CI bug because when unwire fails, it just raises a
warning (e.g. driving an already driven value) that occurs silently in
the passing tests.  Adds a test to cover this case.
---
 magma/array.py                   |  3 ---
 magma/tuple.py                   | 13 ++++++++++---
 tests/test_aggregate_wireable.py | 18 ++++++++++++++++++
 3 files changed, 28 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_aggregate_wireable.py

diff --git a/magma/array.py b/magma/array.py
index 2bd0f2edb..8ec6f765d 100644
--- a/magma/array.py
+++ b/magma/array.py
@@ -626,9 +626,6 @@ def wire(self, o, debug_info):
     @debug_unwire
     @aggregate_wireable_method
     def unwire(self, o=None, debug_info=None, keep_wired_when_contexts=False):
-        if not self._has_elaborated_children():
-            return AggregateWireable.unwire(self, o, debug_info,
-                                            keep_wired_when_contexts)
         for i, child in self._enumerate_children():
             o_i = None if o is None else o[i]
             child.unwire(o_i, debug_info=debug_info,
diff --git a/magma/tuple.py b/magma/tuple.py
index ac62eb754..ef9825585 100644
--- a/magma/tuple.py
+++ b/magma/tuple.py
@@ -342,9 +342,16 @@ def wire(self, o, debug_info):
     @debug_unwire
     @aggregate_wireable_method
     def unwire(self, o=None, debug_info=None, keep_wired_when_contexts=False):
-        if not self._has_elaborated_children():
-            return AggregateWireable.unwire(self, o, debug_info,
-                                            keep_wired_when_contexts)
+        for k, t in self.items():
+            if o is None:
+                t.unwire(debug_info=debug_info,
+                         keep_wired_when_contexts=keep_wired_when_contexts)
+            elif o[k].is_input():
+                o[k].unwire(t, debug_info=debug_info,
+                            keep_wired_when_contexts=keep_wired_when_contexts)
+            else:
+                t.unwire(o[k], debug_info=debug_info,
+                         keep_wired_when_contexts=keep_wired_when_contexts)
 
     @aggregate_wireable_method
     def driven(self):
diff --git a/tests/test_aggregate_wireable.py b/tests/test_aggregate_wireable.py
new file mode 100644
index 000000000..162a6b8a5
--- /dev/null
+++ b/tests/test_aggregate_wireable.py
@@ -0,0 +1,18 @@
+import pytest
+import magma as m
+
+
+@pytest.mark.parametrize('T', [m.Bits[8], m.Tuple[m.Bit, m.Bits[8]]])
+def test_aggregate_wireable_unwire(T):
+    class Foo(m.Circuit):
+        io = m.IO(I=m.In(T), O=m.Out(T))
+        io.O @= io.I
+        assert io.O.value() is io.I
+        io.O.unwire(io.I)
+        assert io.O.value() is None
+
+        io.O @= io.I
+        io.O[0]  # Trigger elaboration
+        assert io.O.value() is io.I
+        io.O.unwire(io.I)
+        assert io.O.value() is None