await registered actions in Simulator

lib/src/modules/conditional.dart

@@ -369,13 +369,22 @@ class Sequential extends _Always {
             // once the clocks are stable, execute the contents of the FF
             _execute();
             _pendingExecute = false;
-          }).catchError(
-            test: (error) => error is Exception,
-            // ignore: avoid_types_on_closure_parameters
-            (Object err, StackTrace stackTrace) {
-              Simulator.throwException(err as Exception, stackTrace);
-            },
-          ));
+          }).catchError(test: (error) => error is Exception,
+              // ignore: avoid_types_on_closure_parameters
+              (Object err, StackTrace stackTrace) {
+            Simulator.throwException(err as Exception, stackTrace);
+          }).catchError(test: (error) => error is StateError,
+              // ignore: avoid_types_on_closure_parameters
+              (Object err, StackTrace stackTrace) {
+            // This could be a result of the `Simulator` being reset, causing
+            // the stream to `close` before `first` occurs.
+            if (!Simulator.simulationHasEnded) {
+              // If the `Simulator` is still running, rethrow immediately.
+
+              // ignore: only_throw_errors
+              throw err;
+            }
+          }));
         }
         _pendingExecute = true;
       });

lib/src/simulator.dart

@@ -73,8 +73,8 @@ class Simulator {
       _pendingTimestamps.isNotEmpty || _injectedActions.isNotEmpty;
 
   /// Sorted storage for pending functions to execute at appropriate times.
-  static final SplayTreeMap<int, List<void Function()>> _pendingTimestamps =
-      SplayTreeMap<int, List<void Function()>>();
+  static final SplayTreeMap<int, List<dynamic Function()>> _pendingTimestamps =
+      SplayTreeMap<int, List<dynamic Function()>>();
 
   /// Functions to be executed as soon as possible by the [Simulator].
   ///
@@ -170,7 +170,9 @@ class Simulator {
   }
 
   /// Registers an abritrary [action] to be executed at [timestamp] time.
-  static void registerAction(int timestamp, void Function() action) {
+  ///
+  /// The [action], if it returns a [Future], will be `await`ed.
+  static void registerAction(int timestamp, dynamic Function() action) {
     if (timestamp <= _currentTimestamp) {
       throw Exception('Cannot add timestamp "$timestamp" in the past.'
           '  Current time is ${Simulator.time}');
@@ -225,9 +227,9 @@ class Simulator {
 
     _currentTimestamp = nextTimeStamp;
 
-    await tickExecute(() {
+    await tickExecute(() async {
       for (final func in _pendingTimestamps[nextTimeStamp]!) {
-        func();
+        await func();
       }
     });
     _pendingTimestamps.remove(_currentTimestamp);
@@ -242,7 +244,7 @@ class Simulator {
   }
 
   /// Performs the actual execution of a collection of actions for a [tick()].
-  static Future<void> tickExecute(void Function() toExecute) async {
+  static Future<void> tickExecute(dynamic Function() toExecute) async {
     _phase = SimulatorPhase.beforeTick;
 
     // useful for flop sampling
@@ -252,7 +254,7 @@ class Simulator {
 
     // useful for things that need to trigger every tick without other input
     _startTickController.add(null);
-    toExecute();
+    await toExecute();
 
     _phase = SimulatorPhase.clkStable;
 

test/simulator_test.dart

@@ -51,7 +51,9 @@ void main() {
 
   test('simulator reset waits for simulation to complete', () async {
     Simulator.registerAction(100, Simulator.endSimulation);
-    Simulator.registerAction(100, Simulator.reset);
+    Simulator.registerAction(100, () {
+      unawaited(Simulator.reset());
+    });
     Simulator.registerAction(100, () => true);
     await Simulator.run();
   });
@@ -76,4 +78,28 @@ void main() {
     await Simulator.run();
     expect(endOfSimActionExecuted, isTrue);
   });
+
+  test('simulator waits for async registered actions to complete', () async {
+    var registeredActionExecuted = false;
+    Simulator.registerAction(100, () => true);
+    Simulator.registerAction(50, () async {
+      await Future<void>.delayed(const Duration(microseconds: 10));
+      registeredActionExecuted = true;
+    });
+    await Simulator.run();
+    expect(registeredActionExecuted, isTrue);
+  });
+
+  test('simulator waits for async injected actions to complete', () async {
+    var injectedActionExecuted = false;
+    Simulator.registerAction(100, () => true);
+    Simulator.registerAction(50, () async {
+      Simulator.injectAction(() async {
+        await Future<void>.delayed(const Duration(microseconds: 10));
+        injectedActionExecuted = true;
+      });
+    });
+    await Simulator.run();
+    expect(injectedActionExecuted, isTrue);
+  });
 }