sim.core: use asyncgen hooks to call aclose()

Async event loops are responsible for cleaning up async generators. This is
a best-effort implementation to call aclose() when possible.

Reverts #1590 and closes #1638.

Author: agrif

amaranth/sim/_async.py

@@ -362,15 +362,12 @@ async def until(self, condition: ValueLike):
                 raise TypeError(f"The shape of a condition may only be `signed` or `unsigned`, "
                                 f"not {shape!r}")
         tick = self.sample(condition).__aiter__()
-        try:
-            done = False
-            while not done:
-                clk, rst, *values, done = await tick.__anext__()
-                if rst:
-                    raise DomainReset
-            return tuple(values)
-        finally:
-            await tick.aclose()
+        done = False
+        while not done:
+            clk, rst, *values, done = await tick.__anext__()
+            if rst:
+                raise DomainReset
+        return tuple(values)
 
     async def repeat(self, count: int):
         """Repeat this trigger a specific number of times.
@@ -403,15 +400,12 @@ async def repeat(self, count: int):
         if count <= 0:
             raise ValueError(f"Repeat count must be a positive integer, not {count!r}")
         tick = self.__aiter__()
-        try:
-            for _ in range(count):
-                clk, rst, *values = await tick.__anext__()
-                if rst:
-                    raise DomainReset
-                assert clk
-            return tuple(values)
-        finally:
-            await tick.aclose()
+        for _ in range(count):
+            clk, rst, *values = await tick.__anext__()
+            if rst:
+                raise DomainReset
+            assert clk
+        return tuple(values)
 
     def _collect_trigger(self):
         clk_polarity = (1 if self._domain.clk_edge == "pos" else 0)

amaranth/sim/core.py

@@ -1,4 +1,6 @@
+from contextlib import contextmanager
 import inspect
+import sys
 import warnings
 
 from .._utils import deprecated
@@ -341,6 +343,33 @@ def run_until(self, deadline, *, run_passive=None):
         while self._engine.now < deadline.femtoseconds:
             self.advance()
 
+    @contextmanager
+    def _replace_asyncgen_hooks(self):
+        # Async generators require hooks for lifetime management. Replace existing hooks with ours
+        # for the duration of this context manager.
+
+        def firstiter(agen):
+            # Prevent any outer event loop from seeing this generator.
+            pass
+
+        def finalizer(agen):
+            # Generators can't be closed if they are currently running.
+            if not agen.ag_running:
+                # Try to run aclose() once, but skip it if it awaits anything.
+                try:
+                    coroutine = agen.aclose()
+                    coroutine.send(None)
+                except StopIteration:
+                    # Success
+                    return
+
+        old_hooks = sys.get_asyncgen_hooks()
+        sys.set_asyncgen_hooks(firstiter=firstiter, finalizer=finalizer)
+        try:
+            yield
+        finally:
+            sys.set_asyncgen_hooks(*old_hooks)
+
     def advance(self):
         """Advance the simulation.
 
@@ -356,7 +385,8 @@ def advance(self):
         :py:`False` otherwise.
         """
         self._running = True
-        return self._engine.advance()
+        with self._replace_asyncgen_hooks():
+            return self._engine.advance()
 
     def write_vcd(self, vcd_file, gtkw_file=None, *, traces=(), fs_per_delta=0):
         # `fs_per_delta`` is not currently documented; it is not clear if we want to expose