Add @AllChildren annotation

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

@@ -3,6 +3,7 @@
 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.ActivityType.AllChildren;
 import gov.nasa.jpl.aerie.merlin.framework.annotations.Export.Parameter;
 
 import static gov.nasa.jpl.aerie.banananation.generated.ActivityActions.call;
@@ -15,6 +16,7 @@ public static final class ParentActivity {
     @Parameter
     public String label = "unlabeled";
 
+    @AllChildren(children = {"child"})
     @EffectModel
     public void run(final Mission mission) {
       call(mission, new ChildActivity(1));
@@ -34,6 +36,7 @@ public ChildActivity(final int counter) {
       this.counter = counter;
     }
 
+    @ActivityType.AllChildren(children = {"grandchild"})
     @EffectModel
     public void run(final Mission mission) {
       call(mission, new GrandchildActivity(1));
@@ -54,6 +57,7 @@ public GrandchildActivity(final int counter) {
     }
 
     @EffectModel
+    @AllChildren(children = {})
     public void run(final Mission mission) {
       delay(6*24, HOURS);
     }

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

@@ -6,8 +6,7 @@
 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.framework.ModelActions.delay;
 import static gov.nasa.jpl.aerie.merlin.protocol.types.Duration.SECOND;
 
 public final class DecomposingSpawnActivity {

merlin-framework-processor/src/main/java/gov/nasa/jpl/aerie/merlin/processor/MissionModelParser.java

@@ -70,6 +70,8 @@ public MissionModelRecord parseMissionModel(final PackageElement missionModelEle
       activityTypes.add(this.parseActivityType(missionModelElement, activityTypeElement));
     }
 
+    verifyChildrenNames(activityTypes);
+
     return new MissionModelRecord(
         missionModelElement,
         topLevelModel.type,
@@ -79,6 +81,21 @@ public MissionModelRecord parseMissionModel(final PackageElement missionModelEle
         activityTypes);
   }
 
+  private void verifyChildrenNames(final List<ActivityTypeRecord> activityTypeRecords)
+  throws InvalidMissionModelException
+  {
+    final var allActivityNames = activityTypeRecords.stream().map(ActivityTypeRecord::name).collect(Collectors.toSet());
+    for(final var activityTypeRecord:activityTypeRecords){
+      if(activityTypeRecord.effectModel().isPresent() && activityTypeRecord.effectModel().get().children().isPresent()){
+        for(final var childName: activityTypeRecord.effectModel().get().children().get()){
+          if(!allActivityNames.contains(childName)){
+            throw new InvalidMissionModelException(childName + " has been declared as a child of "+ activityTypeRecord.name() + " with the @AllChildren annotation but it is not a valid activity type name.");
+          }
+        }
+      }
+    }
+  }
+
   private record MissionModelTypeRecord(
       TypeElement type,
       boolean expectsPlanStart,
@@ -550,6 +567,7 @@ private Optional<EffectModelRecord> getActivityEffectModel(final TypeElement act
   {
     Optional<String> fixedDuration = Optional.empty();
     Optional<String> parameterizedDuration = Optional.empty();
+    Optional<String[]> children = Optional.empty();
     for (final var element: activityTypeElement.getEnclosedElements()) {
       if (element.getAnnotation(ActivityType.FixedDuration.class) != null) {
         if (fixedDuration.isPresent()) throw new InvalidMissionModelException(
@@ -581,6 +599,13 @@ private Optional<EffectModelRecord> getActivityEffectModel(final TypeElement act
         );
 
         parameterizedDuration = Optional.of(executableElement.getSimpleName().toString());
+      } else if (element.getAnnotation(ActivityType.AllChildren.class) != null) {
+        if (children.isPresent()) throw new InvalidMissionModelException(
+            "AllChildren annotation cannot be applied multiple times in one activity type."
+        );
+        if (!(element instanceof ExecutableElement executableElement)) throw new InvalidMissionModelException(
+            "AllChildren method annotation must be an executable element.");
+        children = Optional.of(element.getAnnotation(ActivityType.AllChildren.class).children());
       }
     }
 
@@ -609,7 +634,7 @@ private Optional<EffectModelRecord> getActivityEffectModel(final TypeElement act
           ? Optional.<TypeMirror>empty()
           : Optional.of(returnType);
 
-      return Optional.of(new EffectModelRecord(element.getSimpleName().toString(), executorAnnotation.value(), nonVoidReturnType, durationParameter, fixedDuration, parameterizedDuration));
+      return Optional.of(new EffectModelRecord(element.getSimpleName().toString(), executorAnnotation.value(), nonVoidReturnType, durationParameter, fixedDuration, parameterizedDuration, children));
     }
 
     return Optional.empty();

merlin-framework-processor/src/main/java/gov/nasa/jpl/aerie/merlin/processor/generator/MissionModelGenerator.java

@@ -46,6 +46,7 @@
 import javax.lang.model.util.Types;
 import javax.tools.Diagnostic;
 import java.time.Instant;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -265,7 +266,7 @@ public JavaFile generateSchedulerModel(final MissionModelRecord missionModel) {
             .classBuilder(typeName)
             .addAnnotation(
                 AnnotationSpec
-                    .builder(javax.annotation.processing.Generated.class)
+                    .builder(Generated.class)
                     .addMember("value", "$S", MissionModelProcessor.class.getCanonicalName())
                     .build())
             .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
@@ -327,6 +328,23 @@ public JavaFile generateSchedulerModel(final MissionModelRecord missionModel) {
                            )
                            .returns(Duration.class)
                            .build())
+            .addMethod(MethodSpec
+                           .methodBuilder("getChildren")
+                           .addModifiers(Modifier.PUBLIC)
+                           .addAnnotation(Override.class)
+                           .returns(ParameterizedTypeName.get(ClassName.get(Map.class), ClassName.get(String.class), ParameterizedTypeName.get(List.class, String.class)))
+                           .addStatement("final var result = new $T()", ParameterizedTypeName.get(ClassName.get(HashMap.class), ClassName.get(String.class), ParameterizedTypeName.get(List.class, String.class)))
+                           .addCode(
+                               missionModel
+                                   .activityTypes()
+                                   .stream()
+                                   .map(
+                                       activityTypeRecord ->
+                                           getChildrenStatement(missionModel, activityTypeRecord))
+                                   .reduce((x, y) -> x.add("$L", y.build()))
+                                   .orElse(CodeBlock.builder()).build())
+                           .addStatement("return result")
+                           .build())
 
             .build();
 
@@ -336,6 +354,26 @@ public JavaFile generateSchedulerModel(final MissionModelRecord missionModel) {
         .build();
   }
 
+  private CodeBlock.Builder getChildrenStatement(final MissionModelRecord missionModel, final ActivityTypeRecord activityTypeRecord){
+    String[] children = null;
+    //if the effect model is empty, an activity cannot have children
+    if(activityTypeRecord.effectModel().isEmpty()) children = new String[]{};
+    else if(activityTypeRecord.effectModel().get().children().isPresent()){
+      //if children have been declared with an annotation, record those
+      children = activityTypeRecord.effectModel().get().children().get();
+    } else {
+      //if children have not been declared with an annotation, assume the activity might generate all activity types
+      children = missionModel.activityTypes().stream().map(ActivityTypeRecord::name).toArray(String[]::new);
+    }
+    final var childrenString = String.join(",", Arrays.stream(children).map(child-> "\""+child + "\"").toList());
+
+    return CodeBlock
+        .builder()
+        .addStatement("result.put(\"$L\", $L)",
+                      activityTypeRecord.name(),
+                      CodeBlock.of("$T.of($N)", List.class, childrenString));
+  }
+
   /** Generate `ActivityActions` class. */
   public JavaFile generateActivityActions(final MissionModelRecord missionModel) {
     final var typeName = missionModel.getActivityActionsName();

merlin-framework-processor/src/main/java/gov/nasa/jpl/aerie/merlin/processor/metamodel/EffectModelRecord.java

@@ -11,6 +11,7 @@ public record EffectModelRecord(
     Optional<TypeMirror> returnType,
     Optional<String> durationParameter,
     Optional<String> fixedDurationExpr,
-    Optional<String> parametricDuration
+    Optional<String> parametricDuration,
+    Optional<String[]> children
 ) {
 }

merlin-framework/src/main/java/gov/nasa/jpl/aerie/merlin/framework/annotations/ActivityType.java

@@ -134,4 +134,10 @@ enum Executor { Threaded, Replaying }
   @Retention(RetentionPolicy.CLASS)
   @Target(ElementType.METHOD)
   @interface ParametricDuration {}
+
+  @Retention(RetentionPolicy.CLASS)
+  @Target(ElementType.METHOD)
+  @interface AllChildren {
+    String[] children();
+  }
 }

merlin-sdk/src/main/java/gov/nasa/jpl/aerie/merlin/protocol/model/SchedulerModel.java

@@ -4,10 +4,12 @@
 import gov.nasa.jpl.aerie.merlin.protocol.types.DurationType;
 import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue;
 
+import java.util.List;
 import java.util.Map;
 
 public interface SchedulerModel {
   Map<String, DurationType> getDurationTypes();
   SerializedValue serializeDuration(final Duration duration);
   Duration deserializeDuration(final SerializedValue serializedValue);
+  Map<String, List<String>> getChildren();
 }

scheduler-driver/src/test/java/gov/nasa/jpl/aerie/scheduler/AllChildrenTest.java

@@ -0,0 +1,29 @@
+package gov.nasa.jpl.aerie.scheduler;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.HashSet;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class AllChildrenTest {
+  @Test
+  public void testAllChildrenAnnotation(){
+    final var bananaSchedulerModel = SimulationUtility.getBananaSchedulerModel();
+    final var children = bananaSchedulerModel.getChildren();
+    final var bananaMissionModel = SimulationUtility.getBananaMissionModel();
+    final var allActivityTypes = bananaMissionModel.getDirectiveTypes().directiveTypes().keySet();
+    //there is one key per activity type in the children map
+    allActivityTypes.forEach(at -> assertTrue(children.containsKey(at)));
+    //the only declared child of parent is the one present in the children map
+    assertEquals(List.of("child"), children.get("parent"));
+    //the declared absence of children of grandchild is reported in the children map
+    assertEquals(List.of(), children.get("grandchild"));
+    //an activity without effect model does not have any children
+    assertEquals(List.of(), children.get("ParameterTest"));
+    //an activity with an effect model but with no @AllChildren annotation will result in all activity types of the mission model being reported as its potential children
+    assertEquals(new HashSet<>(allActivityTypes.stream().toList()), new HashSet<>(children.get("BiteBanana")));
+  }
+}