Implement Spans.selectWhenTrue

NASA-AMMOS · Sep 18, 2023 · c7ca72f · c7ca72f
1 parent c698611
commit c7ca72f
Show file tree

Hide file tree

Showing 4 changed files with 87 additions and 2 deletions.
diff --git a/constraints/src/main/java/gov/nasa/jpl/aerie/constraints/json/ConstraintParsers.java b/constraints/src/main/java/gov/nasa/jpl/aerie/constraints/json/ConstraintParsers.java
@@ -260,6 +260,18 @@ static JsonParser<Transition> transitionP(JsonParser<ProfileExpression<?>> profi
           .map(
               untuple((kind, alias) -> new ActivitySpan(alias)),
               $ -> tuple(Unit.UNIT, $.activityAlias));
+
+  static JsonParser<SpansSelectWhenTrue> spansSelectWhenTrueF(JsonParser<Expression<Spans>> spansP, JsonParser<Expression<Windows>> windowsP) {
+    return productP
+        .field("kind", literalP("SpansSelectWhenTrue"))
+        .field("spansExpression", spansP)
+        .field("windowsExpression", windowsP)
+        .map(
+            untuple((kind, spans, windows) -> new SpansSelectWhenTrue(spans, windows)),
+            $ -> tuple(Unit.UNIT, $.spans(), $.windows())
+        );
+  }
+
   static final JsonParser<StartOf> startOfP =
       productP
           .field("kind", literalP("WindowsExpressionStartOf"))
@@ -600,7 +612,8 @@ private static JsonParser<Expression<Spans>> spansExpressionF(JsonParser<Express
           splitF(windowsP),
           spansFromWindowsF(windowsP),
           forEachActivitySpansF(selfP),
-          activitySpanP
+          activitySpanP,
+          spansSelectWhenTrueF(selfP, windowsP)
           ));
   }
 

diff --git a/constraints/src/main/java/gov/nasa/jpl/aerie/constraints/tree/SpansSelectWhenTrue.java b/constraints/src/main/java/gov/nasa/jpl/aerie/constraints/tree/SpansSelectWhenTrue.java
@@ -0,0 +1,38 @@
+package gov.nasa.jpl.aerie.constraints.tree;
+
+import gov.nasa.jpl.aerie.constraints.model.EvaluationEnvironment;
+import gov.nasa.jpl.aerie.constraints.model.SimulationResults;
+import gov.nasa.jpl.aerie.constraints.time.Interval;
+import gov.nasa.jpl.aerie.constraints.time.Spans;
+import gov.nasa.jpl.aerie.constraints.time.Windows;
+
+import java.util.ArrayList;
+import java.util.Set;
+
+public record SpansSelectWhenTrue(Expression<Spans> spans, Expression<Windows> windows) implements Expression<Spans> {
+
+  @Override
+  public Spans evaluate(SimulationResults results, final Interval bounds, EvaluationEnvironment environment) {
+    final var spans = this.spans.evaluate(results, bounds, environment);
+    final var windows = this.windows.evaluate(results, bounds, environment);
+    final var trueSegments = new ArrayList<Interval>();
+    for (final var window: windows.iterateEqualTo(true)) trueSegments.add(window);
+    return spans.select(trueSegments.toArray(new Interval[] {}));
+  }
+
+  @Override
+  public void extractResources(final Set<String> names) {
+    this.spans.extractResources(names);
+    this.windows.extractResources(names);
+  }
+
+  @Override
+  public String prettyPrint(final String prefix) {
+    return String.format(
+        "\n%s(spans-select-when-true %s %s)",
+        prefix,
+        this.spans.prettyPrint(prefix + "  "),
+        this.windows.prettyPrint(prefix + "  ")
+    );
+  }
+}
diff --git a/merlin-server/constraints-dsl-compiler/src/libs/constraints-ast.ts b/merlin-server/constraints-dsl-compiler/src/libs/constraints-ast.ts
@@ -28,6 +28,7 @@ export enum NodeKind {
   SpansExpressionFromWindows = 'SpansExpressionFromWindows',
   SpansExpressionSplit = 'SpansExpressionSplit',
   SpansExpressionInterval = 'SpansExpressionInterval',
+  SpansSelectWhenTrue = 'SpansSelectWhenTrue',
   ExpressionEqual = 'ExpressionEqual',
   ExpressionNotEqual = 'ExpressionNotEqual',
   RealProfileLessThan = 'RealProfileLessThan',
@@ -121,7 +122,14 @@ export type SpansExpression =
   | IntervalsExpressionShiftEdges
   | SpansExpressionFromWindows
   | ForEachActivitySpans
-  | SpansExpressionInterval;
+  | SpansExpressionInterval
+  | SpansSelectWhenTrue;
+
+export interface SpansSelectWhenTrue {
+  kind: NodeKind.SpansSelectWhenTrue,
+  spansExpression: SpansExpression
+  windowsExpression: WindowsExpression
+}
 
 export type IntervalsExpression =
   | WindowsExpression

diff --git a/merlin-server/constraints-dsl-compiler/src/libs/constraints-edsl-fluent-api.ts b/merlin-server/constraints-dsl-compiler/src/libs/constraints-edsl-fluent-api.ts
@@ -538,6 +538,22 @@ export class Spans {
       expression: expression(new ActivityInstance(activityType, alias)).__astNode,
     });
   }
+
+  /**
+   * Selects only spans that occur during a true segment, removing those that don't.
+   *
+   * Spans that only partially overlap with a true segment will be truncated, and spans
+   * that overlap with multiple true segments will be split.
+   *
+   * @param windows
+   */
+  public selectWhenTrue(windows: Windows): Spans {
+    return new Spans({
+      kind: AST.NodeKind.SpansSelectWhenTrue,
+      spansExpression: this.__astNode,
+      windowsExpression: windows.__astNode
+    });
+  }
 }
 
 /**
@@ -1379,6 +1395,16 @@ declare global {
      * @param unit unit of time to count. Does not need to be a round unit (i.e. can be 1.5 minutes, if you want).
      */
     public accumulatedDuration(unit: AST.Duration): Real;
+
+    /**
+     * Selects only spans that occur during a true segment, removing those that don't.
+     *
+     * Spans that only partially overlap with a true segment will be truncated, and spans
+     * that overlap with multiple true segments will be split.
+     *
+     * @param windows
+     */
+    public selectWhenTrue(windows: Windows): Spans;
   }
 
   /**