Procedural constraints annotation processor + build logic

procedural/constraints/src/main/kotlin/gov/nasa/ammos/aerie/procedural/constraints/ConstraintProcedureMapper.kt

@@ -0,0 +1,10 @@
+package gov.nasa.ammos.aerie.procedural.constraints
+
+import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue
+import gov.nasa.jpl.aerie.merlin.protocol.types.ValueSchema
+
+interface ConstraintProcedureMapper<T: Constraint> {
+  fun valueSchema(): ValueSchema
+  fun serialize(procedure: T): SerializedValue
+  fun deserialize(arguments: SerializedValue): T
+}

procedural/constraints/src/main/kotlin/gov/nasa/ammos/aerie/procedural/constraints/annotations/ConstraintProcedure.kt

@@ -0,0 +1,3 @@
+package gov.nasa.ammos.aerie.procedural.constraints.annotations
+
+annotation class ConstraintProcedure

procedural/examples/foo-procedures/build.gradle

@@ -73,6 +73,52 @@ tasks.create("generateSchedulingProcedureJarTasks") {
   }
 }
 
+tasks.register('buildAllConstraintProcedureJars') {
+  group = 'ConstraintProcedureJars'
+
+  dependsOn "generateConstraintProcedureJarTasks"
+  dependsOn {
+    tasks.findAll { task -> task.name.startsWith('buildConstraintProcedureJar_') }
+  }
+}
+
+tasks.create("generateConstraintProcedureJarTasks") {
+  group = 'ConstraintProcedureJars'
+
+  final proceduresDir = findFirstMatchingBuildDir("generated/procedures")
+
+  if (proceduresDir == null) {
+    println "No procedures folder found"
+    return
+  }
+  println "Generating jar tasks for the following procedures directory: ${proceduresDir}"
+
+  final files = file(proceduresDir).listFiles()
+  if (files.length == 0) {
+    println "No procedures available within folder ${proceduresDir}"
+    return
+  }
+
+  files.toList().each { file ->
+    final nameWithoutExtension = file.name.replace(".java", "")
+    final taskName = "buildConstraintProcedureJar_${nameWithoutExtension}"
+
+    println "Generating ${taskName} task, which will build ${nameWithoutExtension}.jar"
+
+    tasks.create(taskName, ShadowJar) {
+      group = 'ConstraintProcedureJars'
+      configurations = [project.configurations.runtimeClasspath]
+      from sourceSets.main.output
+      archiveBaseName = "" // clear
+      archiveClassifier.set(nameWithoutExtension) // set output jar name
+      manifest {
+        attributes 'Main-Class': getMainClassFromGeneratedFile(file)
+      }
+      minimize()
+    }
+  }
+}
+
 private String findFirstMatchingBuildDir(String pattern) {
   String found = null
   final generatedDir = file("build/generated/sources")

procedural/examples/foo-procedures/src/main/java/gov/nasa/ammos/aerie/procedural/examples/fooprocedures/constraints/ConstFruit.java

@@ -1,20 +1,22 @@
 package gov.nasa.ammos.aerie.procedural.examples.fooprocedures.constraints;
 
-import gov.nasa.ammos.aerie.procedural.constraints.GeneratorConstraint;
+import gov.nasa.ammos.aerie.procedural.constraints.Constraint;
 import gov.nasa.ammos.aerie.procedural.constraints.Violations;
+import gov.nasa.ammos.aerie.procedural.constraints.annotations.ConstraintProcedure;
 import gov.nasa.ammos.aerie.procedural.timeline.collections.profiles.Real;
 import gov.nasa.ammos.aerie.procedural.timeline.plan.Plan;
 import gov.nasa.ammos.aerie.procedural.timeline.plan.SimulationResults;
 import org.jetbrains.annotations.NotNull;
 
-public class ConstFruit extends GeneratorConstraint {
+@ConstraintProcedure
+public record ConstFruit() implements Constraint {
   @Override
-  public void generate(@NotNull Plan plan, @NotNull SimulationResults simResults) {
+  public Violations run(@NotNull Plan plan, @NotNull SimulationResults simResults) {
     final var fruit = simResults.resource("/fruit", Real.deserializer());
 
-    violate(Violations.on(
+    return Violations.on(
         fruit.equalTo(4),
         false
-    ));
+    );
   }
 }

procedural/processor/build.gradle

@@ -15,6 +15,7 @@ dependencies {
   implementation project(':merlin-sdk')
   implementation project(':contrib')
   implementation project(':procedural:scheduling')
+  implementation project(':procedural:constraints')
   implementation 'org.apache.commons:commons-lang3:3.13.0'
   implementation 'com.squareup:javapoet:1.13.0'
 }

procedural/processor/src/main/java/gov/nasa/ammos/aerie/procedural/processor/AutoValueMappers.java

@@ -59,7 +59,7 @@ static JavaFile generateAutoValueMappers(final ClassName typeName, final Iterabl
             .addAnnotation(
                 AnnotationSpec
                     .builder(javax.annotation.processing.Generated.class)
-                    .addMember("value", "$S", SchedulingProcedureProcessor.class.getCanonicalName())
+                    .addMember("value", "$S", ProcedureProcessor.class.getCanonicalName())
                     .build())
             .addAnnotation(
                 AnnotationSpec

procedural/processor/src/main/java/gov/nasa/ammos/aerie/procedural/processor/SchedulingProcedureProcessor.java → procedural/processor/src/main/java/gov/nasa/ammos/aerie/procedural/processor/ProcedureProcessor.java

@@ -9,10 +9,13 @@
 import gov.nasa.jpl.aerie.merlin.framework.ValueMapper;
 import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue;
 import gov.nasa.jpl.aerie.merlin.protocol.types.ValueSchema;
-import gov.nasa.ammos.aerie.procedural.scheduling.ProcedureMapper;
+import gov.nasa.ammos.aerie.procedural.scheduling.SchedulingProcedureMapper;
 import gov.nasa.ammos.aerie.procedural.scheduling.annotations.SchedulingProcedure;
 import gov.nasa.ammos.aerie.procedural.scheduling.annotations.WithMappers;
 
+import gov.nasa.ammos.aerie.procedural.constraints.ConstraintProcedureMapper;
+import gov.nasa.ammos.aerie.procedural.constraints.annotations.ConstraintProcedure;
+
 import javax.annotation.processing.Completion;
 import javax.annotation.processing.Filer;
 import javax.annotation.processing.Messager;
@@ -42,8 +45,11 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
-public final class SchedulingProcedureProcessor implements Processor {
+public final class ProcedureProcessor implements Processor {
   // Effectively final, late-initialized
   private Messager messager = null;
   private Filer filer = null;
@@ -58,7 +64,11 @@ public Set<String> getSupportedOptions() {
   /** Elements marked by these annotations will be treated as processing roots. */
   @Override
   public Set<String> getSupportedAnnotationTypes() {
-    return Set.of(SchedulingProcedure.class.getCanonicalName(), WithMappers.class.getCanonicalName());
+    return Set.of(
+        SchedulingProcedure.class.getCanonicalName(),
+        ConstraintProcedure.class.getCanonicalName(),
+        WithMappers.class.getCanonicalName()
+    );
   }
 
   @Override
@@ -104,34 +114,40 @@ public boolean process(final Set<? extends TypeElement> annotations, final Round
       typeRules.addAll(parseValueMappers(factory));
     }
 
-    final var procedures = roundEnv.getElementsAnnotatedWith(SchedulingProcedure.class);
+    final var schedulingProcedures = roundEnv.getElementsAnnotatedWith(SchedulingProcedure.class);
+    final var constraintProcedures = roundEnv.getElementsAnnotatedWith(ConstraintProcedure.class);
 
     final var generatedClassName = ClassName.get(packageElement.getQualifiedName() + ".generated", "AutoValueMappers");
-    for (final var procedure : procedures) {
+    for (final var procedure : schedulingProcedures) {
+      final var procedureElement = (TypeElement) procedure;
+      typeRules.add(AutoValueMappers.recordTypeRule(procedureElement, generatedClassName));
+    }
+
+    for (final var procedure : constraintProcedures) {
       final var procedureElement = (TypeElement) procedure;
       typeRules.add(AutoValueMappers.recordTypeRule(procedureElement, generatedClassName));
     }
 
     final var generatedFiles = new ArrayList<JavaFile>();
 
-    generatedFiles.add(AutoValueMappers.generateAutoValueMappers(generatedClassName, procedures, List.of()));
+    final var allProcedures = Stream.concat(schedulingProcedures.stream(),constraintProcedures.stream()).collect(Collectors.toSet());
+
+    generatedFiles.add(AutoValueMappers.generateAutoValueMappers(generatedClassName, allProcedures, List.of()));
 
     // For each procedure, generate a file that implements Procedure, Supplier<ValueMapper>
-    for (final var procedure : procedures) {
+    for (final var procedure : schedulingProcedures) {
       final TypeName procedureType = TypeName.get(procedure.asType());
-      final ParameterizedTypeName valueMapperType = ParameterizedTypeName.get(
-          ClassName.get(ValueMapper.class),
-          procedureType);
 
-      final var valueMapperCode = new Resolver(typeUtils, elementUtils, typeRules).applyRules(new TypePattern.ClassPattern(ClassName.get(ValueMapper.class), List.of(new TypePattern.ClassPattern((ClassName) procedureType, List.of()))));
+      final var valueMapperCode = new Resolver(typeUtils, elementUtils, typeRules)
+          .applyRules(new TypePattern.ClassPattern(ClassName.get(ValueMapper.class), List.of(new TypePattern.ClassPattern((ClassName) procedureType, List.of()))));
       if (valueMapperCode.isEmpty()) throw new Error("Could not generate a valuemapper for procedure " + procedure.getSimpleName());
 
 
       generatedFiles.add(JavaFile
           .builder(generatedClassName.packageName() + ".procedures", TypeSpec
               .classBuilder(procedure.getSimpleName().toString())
               .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
-              .addSuperinterface(ParameterizedTypeName.get(ClassName.get(ProcedureMapper.class), procedureType))
+              .addSuperinterface(ParameterizedTypeName.get(ClassName.get(SchedulingProcedureMapper.class), procedureType))
               .addMethod(MethodSpec
                              .methodBuilder("valueSchema")
                              .addModifiers(Modifier.PUBLIC)
@@ -160,6 +176,52 @@ public boolean process(final Set<? extends TypeElement> annotations, final Round
           .build());
     }
 
+    // For each procedure, generate a file that implements Procedure, Supplier<ValueMapper>
+    for (final var procedure : constraintProcedures) {
+      final TypeName procedureType = TypeName.get(procedure.asType());
+
+      this.messager.printMessage(
+          Diagnostic.Kind.NOTE,
+          "Looking at: " + procedure.toString());
+
+      final var valueMapperCode = new Resolver(typeUtils, elementUtils, typeRules)
+          .applyRules(new TypePattern.ClassPattern(ClassName.get(ValueMapper.class), List.of(new TypePattern.ClassPattern((ClassName) procedureType, List.of()))));
+      if (valueMapperCode.isEmpty()) throw new Error("Could not generate a valuemapper for procedure " + procedure.getSimpleName());
+
+
+      generatedFiles.add(JavaFile
+                             .builder(generatedClassName.packageName() + ".procedures", TypeSpec
+                                 .classBuilder(procedure.getSimpleName().toString())
+                                 .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
+                                 .addSuperinterface(ParameterizedTypeName.get(ClassName.get(ConstraintProcedureMapper.class), procedureType))
+                                 .addMethod(MethodSpec
+                                                .methodBuilder("valueSchema")
+                                                .addModifiers(Modifier.PUBLIC)
+                                                .addAnnotation(Override.class)
+                                                .returns(ValueSchema.class)
+                                                .addStatement("return $L.getValueSchema()", valueMapperCode.get())
+                                                .build())
+                                 .addMethod(MethodSpec
+                                                .methodBuilder("serialize")
+                                                .addModifiers(Modifier.PUBLIC)
+                                                .addAnnotation(Override.class)
+                                                .addParameter(procedureType, "procedure")
+                                                .returns(SerializedValue.class)
+                                                .addStatement("return $L.serializeValue(procedure)", valueMapperCode.get())
+                                                .build())
+                                 .addMethod(MethodSpec
+                                                .methodBuilder("deserialize")
+                                                .addModifiers(Modifier.PUBLIC)
+                                                .addAnnotation(Override.class)
+                                                .addParameter(SerializedValue.class, "value")
+                                                .returns(procedureType)
+                                                .addStatement("return $L.deserializeValue(value).getSuccessOrThrow(e -> new $T(e))", valueMapperCode.get(), RuntimeException.class)
+                                                .build())
+                                 .build())
+                             .skipJavaLangImports(true)
+                             .build());
+    }
+
     for (final var generatedFile : generatedFiles) {
       this.messager.printMessage(
           Diagnostic.Kind.NOTE,

procedural/processor/src/main/resources/META-INF/services/javax.annotation.processing.Processor

@@ -1 +1 @@
-gov.nasa.ammos.aerie.procedural.processor.SchedulingProcedureProcessor
+gov.nasa.ammos.aerie.procedural.processor.ProcedureProcessor

procedural/scheduling/src/main/kotlin/gov/nasa/ammos/aerie/procedural/scheduling/ProcedureMapper.kt → procedural/scheduling/src/main/kotlin/gov/nasa/ammos/aerie/procedural/scheduling/SchedulingProcedureMapper.kt

@@ -3,7 +3,7 @@ package gov.nasa.ammos.aerie.procedural.scheduling
 import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue
 import gov.nasa.jpl.aerie.merlin.protocol.types.ValueSchema
 
-interface ProcedureMapper<T: Goal> {
+interface SchedulingProcedureMapper<T: Goal> {
   fun valueSchema(): ValueSchema
   fun serialize(procedure: T): SerializedValue
   fun deserialize(arguments: SerializedValue): T

scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/ProcedureLoader.java

@@ -1,7 +1,6 @@
 package gov.nasa.jpl.aerie.scheduler;
 
-import gov.nasa.jpl.aerie.merlin.protocol.types.ValueSchema;
-import gov.nasa.ammos.aerie.procedural.scheduling.ProcedureMapper;
+import gov.nasa.ammos.aerie.procedural.scheduling.SchedulingProcedureMapper;
 
 import java.io.IOException;
 import java.net.MalformedURLException;
@@ -12,19 +11,19 @@
 import java.util.jar.JarFile;
 
 public final class ProcedureLoader {
-  public static ProcedureMapper<?> loadProcedure(final Path path)
+  public static SchedulingProcedureMapper<?> loadProcedure(final Path path)
   throws ProcedureLoadException
   {
     final var className = getImplementingClassName(path);
     final var classLoader = new URLClassLoader(new URL[] {pathToUrl(path)});
 
     try {
       final var pluginClass$ = classLoader.loadClass(className);
-      if (!ProcedureMapper.class.isAssignableFrom(pluginClass$)) {
+      if (!SchedulingProcedureMapper.class.isAssignableFrom(pluginClass$)) {
         throw new ProcedureLoadException(path);
       }
 
-      return (ProcedureMapper<?>) pluginClass$.getConstructor().newInstance();
+      return (SchedulingProcedureMapper<?>) pluginClass$.getConstructor().newInstance();
     } catch (final ReflectiveOperationException ex) {
       throw new ProcedureLoadException(path, ex);
     }
@@ -58,7 +57,7 @@ private ProcedureLoadException(final Path path, final Throwable cause) {
       super(
           String.format(
               "No implementation found for `%s` at path `%s`",
-              ProcedureMapper.class.getSimpleName(),
+              SchedulingProcedureMapper.class.getSimpleName(),
               path),
           cause);
     }

scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/Procedure.java

@@ -2,7 +2,7 @@
 
 import gov.nasa.jpl.aerie.merlin.driver.MissionModel;
 import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue;
-import gov.nasa.ammos.aerie.procedural.scheduling.ProcedureMapper;
+import gov.nasa.ammos.aerie.procedural.scheduling.SchedulingProcedureMapper;
 import gov.nasa.ammos.aerie.procedural.scheduling.plan.Edit;
 import gov.nasa.jpl.aerie.scheduler.DirectiveIdGenerator;
 import gov.nasa.jpl.aerie.scheduler.ProcedureLoader;
@@ -45,7 +45,7 @@ public void run(
       final SimulationFacade simulationFacade,
       final DirectiveIdGenerator idGenerator
   ) {
-    final ProcedureMapper<?> procedureMapper;
+    final SchedulingProcedureMapper<?> procedureMapper;
     try {
       procedureMapper = ProcedureLoader.loadProcedure(jarPath);
     } catch (ProcedureLoader.ProcedureLoadException e) {

scheduler-server/src/main/java/gov/nasa/jpl/aerie/scheduler/server/services/SpecificationService.java

@@ -1,6 +1,6 @@
 package gov.nasa.jpl.aerie.scheduler.server.services;
 
-import gov.nasa.ammos.aerie.procedural.scheduling.ProcedureMapper;
+import gov.nasa.ammos.aerie.procedural.scheduling.SchedulingProcedureMapper;
 import gov.nasa.jpl.aerie.scheduler.ProcedureLoader;
 import gov.nasa.jpl.aerie.scheduler.server.exceptions.NoSuchSchedulingGoalException;
 import gov.nasa.jpl.aerie.scheduler.server.exceptions.NoSuchSpecificationException;
@@ -40,7 +40,7 @@ public void refreshSchedulingProcedureParameterTypes(long goalId, long revision)
         // Do nothing
       }
       case GoalType.JAR jar -> {
-        final ProcedureMapper<?> mapper;
+        final SchedulingProcedureMapper<?> mapper;
         try {
           mapper = ProcedureLoader.loadProcedure(Path.of("/usr/src/app/merlin_file_store", jar.path().toString()));
         } catch (ProcedureLoader.ProcedureLoadException e) {