[AERIE-1930] Include awaiting execution state activity tasks in results

examples/banananation/src/main/java/gov/nasa/jpl/aerie/banananation/activities/DecomposingSpawnActivity.java

@@ -0,0 +1,43 @@
+package gov.nasa.jpl.aerie.banananation.activities;
+
+import gov.nasa.jpl.aerie.banananation.Mission;
+import gov.nasa.jpl.aerie.merlin.framework.annotations.ActivityType;
+import gov.nasa.jpl.aerie.merlin.framework.annotations.ActivityType.EffectModel;
+import gov.nasa.jpl.aerie.merlin.framework.annotations.Export.Parameter;
+
+import static gov.nasa.jpl.aerie.banananation.generated.ActivityActions.spawn;
+import static gov.nasa.jpl.aerie.banananation.generated.ActivityActions.call;
+import static gov.nasa.jpl.aerie.merlin.framework.ModelActions.*;
+import static gov.nasa.jpl.aerie.merlin.protocol.types.Duration.SECOND;
+
+public final class DecomposingSpawnActivity {
+  @ActivityType("DecomposingSpawnParent")
+  public static final class DecomposingSpawnParentActivity {
+    @Parameter
+    public String label = "unlabeled";
+
+    @EffectModel
+    public void run(final Mission mission) {
+      spawn(new DecomposingSpawnChildActivity(1));
+      delay(1, SECOND);
+      spawn(new DecomposingSpawnChildActivity(2));
+    }
+  }
+
+  @ActivityType("DecomposingSpawnChild")
+  public static final class DecomposingSpawnChildActivity {
+    @Parameter
+    public int counter = 0;
+
+    public DecomposingSpawnChildActivity() {}
+
+    public DecomposingSpawnChildActivity(final int counter) {
+      this.counter = counter;
+    }
+
+    @EffectModel
+    public void run(final Mission mission) {
+      delay(2, SECOND);
+    }
+  }
+}

examples/banananation/src/main/java/gov/nasa/jpl/aerie/banananation/package-info.java

@@ -15,6 +15,8 @@
 @WithActivityType(DecomposingActivity.ParentActivity.class)
 @WithActivityType(DecomposingActivity.ChildActivity.class)
 @WithActivityType(DecomposingActivity.GrandchildActivity.class)
+@WithActivityType(DecomposingSpawnActivity.DecomposingSpawnParentActivity.class)
+@WithActivityType(DecomposingSpawnActivity.DecomposingSpawnChildActivity.class)
 @WithActivityType(BakeBananaBreadActivity.class)
 @WithActivityType(BananaNapActivity.class)
 
@@ -25,6 +27,7 @@
 import gov.nasa.jpl.aerie.banananation.activities.BiteBananaActivity;
 import gov.nasa.jpl.aerie.banananation.activities.ChangeProducerActivity;
 import gov.nasa.jpl.aerie.banananation.activities.DecomposingActivity;
+import gov.nasa.jpl.aerie.banananation.activities.DecomposingSpawnActivity;
 import gov.nasa.jpl.aerie.banananation.activities.GrowBananaActivity;
 import gov.nasa.jpl.aerie.banananation.activities.LineCountBananaActivity;
 import gov.nasa.jpl.aerie.banananation.activities.ParameterTestActivity;

examples/banananation/src/test/java/gov/nasa/jpl/aerie/banananation/SimulatedActivityTest.java

@@ -1,11 +1,19 @@
 package gov.nasa.jpl.aerie.banananation;
 
+import gov.nasa.jpl.aerie.merlin.driver.MissionModel;
 import gov.nasa.jpl.aerie.merlin.driver.SerializedActivity;
+import gov.nasa.jpl.aerie.merlin.driver.engine.SimulationEngine;
+import gov.nasa.jpl.aerie.merlin.driver.timeline.TemporalEventSource;
+import gov.nasa.jpl.aerie.merlin.protocol.driver.Topic;
 import gov.nasa.jpl.aerie.merlin.protocol.types.Duration;
 import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue;
 import org.apache.commons.lang3.tuple.Pair;
 import org.junit.jupiter.api.Test;
 
+import java.time.Instant;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashSet;
 import java.util.Map;
 
 import static gov.nasa.jpl.aerie.merlin.protocol.types.Duration.SECOND;
@@ -45,4 +53,53 @@ public void testUnspecifiedArgInSimulatedActivity() {
       assertTrue(act.arguments().containsKey("growingDuration"));
     });
   }
+
+  /** This test is a response to not accounting for all Task ExecutionStates
+   * when collecting activities into the results object. This indirectly tests that portion
+   * of {@link gov.nasa.jpl.aerie.merlin.driver.engine.SimulationEngine#computeResults(
+   * SimulationEngine, Instant, Duration, Topic, TemporalEventSource, MissionModel) computeResults()}
+   *
+   * The schedule in this test, results produces Tasks in all three of the states,
+   * {@link gov.nasa.jpl.aerie.merlin.driver.engine.SimulationEngine.ExecutionState.AwaitingChildren AwaitingChildren},
+   * {@link gov.nasa.jpl.aerie.merlin.driver.engine.SimulationEngine.ExecutionState.InProgress InProgress}, and
+   * {@link gov.nasa.jpl.aerie.merlin.driver.engine.SimulationEngine.ExecutionState.Terminated Terminated}.
+   */
+  @Test
+  public void testCollectAllActivitiesInResults() {
+    final var schedule = SimulationUtility.buildSchedule(
+        Pair.of(
+            duration(0, SECONDS),
+            new SerializedActivity("PeelBanana", Map.of())),
+        Pair.of(
+            duration(0, SECONDS),
+            new SerializedActivity("GrowBanana", Map.of(
+                "quantity", SerializedValue.of(1),
+                "growingDuration", SerializedValue.of(Duration.SECOND.times(3).in(Duration.MICROSECONDS))
+            ))),
+        Pair.of(
+            duration(0, SECONDS),
+            new SerializedActivity("DecomposingSpawnParent", Map.of())));
+
+    final var simDuration = duration(2, SECONDS);
+
+    final var simulationResults = SimulationUtility.simulate(schedule, simDuration);
+
+    assertEquals(2, simulationResults.simulatedActivities.size());
+
+    var simulatedActivityTypes = new HashSet<String>();
+    simulationResults.simulatedActivities.forEach( (id, act) -> simulatedActivityTypes.add(act.type()));
+    Collection<String> expectedSimulated = new HashSet<>(
+        Arrays.asList("PeelBanana", "DecomposingSpawnChild"));
+
+    assertEquals(simulatedActivityTypes, expectedSimulated);
+
+    assertEquals(3, simulationResults.unfinishedActivities.size());
+
+    var unfinishedActivityTypes = new HashSet<String>();
+    simulationResults.unfinishedActivities.forEach( (id, act) -> unfinishedActivityTypes.add(act.type()));
+
+    Collection<String> expectedUnfinished = new HashSet<>(
+        Arrays.asList("GrowBanana", "DecomposingSpawnChild", "DecomposingSpawnParent"));
+    assertEquals(unfinishedActivityTypes, expectedUnfinished);
+  }
 }

merlin-driver/src/main/java/gov/nasa/jpl/aerie/merlin/driver/engine/SimulationEngine.java

@@ -531,6 +531,18 @@ public static SimulationResults computeResults(
             activityChildren.getOrDefault(activityId, Collections.emptyList()),
             (activityParents.containsKey(activityId)) ? Optional.empty() : Optional.of(activityId)
         ));
+      } else if (state instanceof ExecutionState.AwaitingChildren<?> e){
+        final var inputAttributes = taskInfo.input().get(task.id());
+        unfinishedActivities.put(activityId, new UnfinishedActivity(
+            inputAttributes.getTypeName(),
+            inputAttributes.getArguments(),
+            startTime.plus(e.startOffset().in(Duration.MICROSECONDS), ChronoUnit.MICROS),
+            activityParents.get(activityId),
+            activityChildren.getOrDefault(activityId, Collections.emptyList()),
+            (activityParents.containsKey(activityId)) ? Optional.empty() : Optional.of(activityId)
+        ));
+      } else {
+        throw new Error("Unexpected subtype of %s: %s".formatted(ExecutionState.class, state.getClass()));
       }
     });