Merge pull request #1027 from NASA-AMMOS/feat/spans-shiftBy

Implement ShiftEdges for Spans

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"))
@@ -335,15 +347,15 @@ static JsonParser<Transition> transitionP(JsonParser<ProfileExpression<?>> profi
               $ -> tuple(Unit.UNIT, $.value(), $.interval())
           );
 
-  static JsonParser<ShiftWindowsEdges> shiftWindowsEdgesF(JsonParser<Expression<Windows>> windowsExpressionP) {
+  static <I extends IntervalContainer<I>> JsonParser<ShiftEdges<I>> shiftEdgesF(final JsonParser<Expression<I>> intervalExpressionP) {
     return productP
-        .field("kind", literalP("WindowsExpressionShiftBy"))
-        .field("windowExpression", windowsExpressionP)
+        .field("kind", literalP("IntervalsExpressionShiftEdges"))
+        .field("expression", intervalExpressionP)
         .field("fromStart", durationExprP)
         .field("fromEnd", durationExprP)
         .map(
-            untuple((kind, windowsExpression, fromStart, fromEnd) -> new ShiftWindowsEdges(windowsExpression, fromStart, fromEnd)),
-            $ -> tuple(Unit.UNIT, $.windows, $.fromStart, $.fromEnd));
+            untuple((kind, expr, fromStart, fromEnd) -> new ShiftEdges<I>(expr, fromStart, fromEnd)),
+            $ -> tuple(Unit.UNIT, $.expression, $.fromStart, $.fromEnd));
   }
   static final JsonParser<EndOf> endOfP =
       productP
@@ -562,7 +574,7 @@ private static JsonParser<Expression<Windows>> windowsExpressionF(JsonParser<Exp
         andF(selfP),
         orF(selfP),
         notF(selfP),
-        shiftWindowsEdgesF(selfP),
+        shiftEdgesF(selfP),
         startsF(selfP),
         endsF(selfP),
         windowsFromSpansF(spansP),
@@ -595,11 +607,13 @@ private static JsonParser<Expression<Spans>> spansExpressionF(JsonParser<Express
           spansIntervalP,
           startsF(selfP),
           endsF(selfP),
+          shiftEdgesF(selfP),
           splitF(selfP),
           splitF(windowsP),
           spansFromWindowsF(windowsP),
           forEachActivitySpansF(selfP),
-          activitySpanP
+          activitySpanP,
+          spansSelectWhenTrueF(selfP, windowsP)
           ));
   }
 

constraints/src/main/java/gov/nasa/jpl/aerie/constraints/time/IntervalContainer.java

@@ -9,4 +9,6 @@ public interface IntervalContainer<T extends IntervalContainer<T>> {
   LinearProfile accumulatedDuration(final Duration unit);
   T starts();
   T ends();
+  T shiftEdges(final Duration fromStart, final Duration fromEnd);
+  T select(final Interval... intervals);
 }

constraints/src/main/java/gov/nasa/jpl/aerie/constraints/time/Spans.java

@@ -215,6 +215,17 @@ public Spans ends() {
     return this.map($ -> Interval.at($.end));
   }
 
+  @Override
+  public Spans shiftEdges(final Duration fromStart, final Duration fromEnd) {
+    return this.map($ -> Interval.between($.start.plus(fromStart), $.startInclusivity, $.end.plus(fromEnd), $.endInclusivity))
+        .filter($ -> !$.isEmpty());
+  }
+
+  @Override
+  public Spans select(final Interval... intervals) {
+    return this.flatMap($ -> Arrays.stream(intervals).map(selection -> Interval.intersect($, selection)));
+  }
+
   @Override
   public boolean equals(final Object obj) {
     if (!(obj instanceof final Spans spans)) return false;

constraints/src/main/java/gov/nasa/jpl/aerie/constraints/time/Windows.java

@@ -342,6 +342,7 @@ public Windows filterByDuration(Duration minDur, Duration maxDur) {
    * @param fromEnd duration to shift true -> false falling edges
    * @return a new Windows
    */
+  @Override
   public Windows shiftEdges(Duration fromStart, Duration fromEnd) {
     final var builder = IntervalMap.<Boolean>builder();
 
@@ -622,6 +623,7 @@ public Windows unset(final List<Interval> intervals) {
   }
 
   /** Delegated to {@link IntervalMap#select(Interval...)} */
+  @Override
   public Windows select(final Interval... intervals) {
     return new Windows(segments.select(intervals));
   }

constraints/src/main/java/gov/nasa/jpl/aerie/constraints/tree/ShiftWindowsEdges.java → constraints/src/main/java/gov/nasa/jpl/aerie/constraints/tree/ShiftEdges.java

@@ -3,25 +3,25 @@
 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.Windows;
+import gov.nasa.jpl.aerie.constraints.time.IntervalContainer;
 import gov.nasa.jpl.aerie.merlin.protocol.types.Duration;
 
 import java.util.Objects;
 import java.util.Set;
 
-public final class ShiftWindowsEdges implements Expression<Windows> {
-  public final Expression<Windows> windows;
+public final class ShiftEdges<I extends IntervalContainer<I>> implements Expression<I> {
+  public final Expression<I> expression;
   public final Expression<Duration> fromStart;
   public final Expression<Duration> fromEnd;
 
-  public ShiftWindowsEdges(final Expression<Windows> left, final Expression<Duration> fromStart, final Expression<Duration> fromEnd) {
-    this.windows = left;
+  public ShiftEdges(final Expression<I> left, final Expression<Duration> fromStart, final Expression<Duration> fromEnd) {
+    this.expression = left;
     this.fromStart = fromStart;
     this.fromEnd = fromEnd;
   }
 
   @Override
-  public Windows evaluate(final SimulationResults results, final Interval bounds, final EvaluationEnvironment environment) {
+  public I evaluate(final SimulationResults results, final Interval bounds, final EvaluationEnvironment environment) {
     final var shiftRising = this.fromStart.evaluate(results, bounds, environment);
     final var shiftFalling = this.fromEnd.evaluate(results, bounds, environment);
 
@@ -32,38 +32,38 @@ public Windows evaluate(final SimulationResults results, final Interval bounds,
         bounds.endInclusivity
     );
 
-    final var windows = this.windows.evaluate(results, newBounds, environment);
-    return windows.shiftEdges(shiftRising, shiftFalling).select(bounds);
+    final var intervals = this.expression.evaluate(results, newBounds, environment);
+    return intervals.shiftEdges(shiftRising, shiftFalling).select(bounds);
   }
 
   @Override
   public void extractResources(final Set<String> names) {
-    this.windows.extractResources(names);
+    this.expression.extractResources(names);
   }
 
   @Override
   public String prettyPrint(final String prefix) {
     return String.format(
-        "\n%s(shiftWindowsEdges %s by %s %s)",
+        "\n%s(shiftEdges %s by %s %s)",
         prefix,
-        this.windows.prettyPrint(prefix + "  "),
+        this.expression.prettyPrint(prefix + "  "),
         this.fromStart.toString(),
         this.fromEnd.toString()
     );
   }
 
   @Override
   public boolean equals(Object obj) {
-    if (!(obj instanceof ShiftWindowsEdges)) return false;
-    final var o = (ShiftWindowsEdges)obj;
+    if (!(obj instanceof ShiftEdges)) return false;
+    final var o = (ShiftEdges<?>)obj;
 
-    return Objects.equals(this.windows, o.windows) &&
+    return Objects.equals(this.expression, o.expression) &&
            Objects.equals(this.fromStart, o.fromStart) &&
            Objects.equals(this.fromEnd, o.fromEnd);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(this.windows, this.fromStart, this.fromEnd);
+    return Objects.hash(this.expression, this.fromStart, this.fromEnd);
   }
 }

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 + "  ")
+    );
+  }
+}

constraints/src/test/java/gov/nasa/jpl/aerie/constraints/tree/ASTTests.java

@@ -378,7 +378,7 @@ public void testExpandBy() {
     final var expandByFromStart = Duration.of(-1, SECONDS);
     final var expandByFromEnd = Duration.of(0, SECONDS);
 
-    final var result = new ShiftWindowsEdges(Supplier.of(left), Supplier.of(expandByFromStart), Supplier.of(expandByFromEnd)).evaluate(simResults, new EvaluationEnvironment());
+    final var result = new ShiftEdges<>(Supplier.of(left), Supplier.of(expandByFromStart), Supplier.of(expandByFromEnd)).evaluate(simResults, new EvaluationEnvironment());
 
     final var expected = new Windows()
         .set(Interval.between(-1, Inclusive, 7, Inclusive, SECONDS), true)
@@ -410,7 +410,7 @@ public void testShrink() {
     final var clampFromStart = Duration.of(1, SECONDS);
     final var clampFromEnd = Duration.negate(Duration.of(1, SECONDS));
 
-    final var result = new ShiftWindowsEdges(Supplier.of(left), Supplier.of(clampFromStart), Supplier.of(clampFromEnd)).evaluate(simResults, new EvaluationEnvironment());
+    final var result = new ShiftEdges<>(Supplier.of(left), Supplier.of(clampFromStart), Supplier.of(clampFromEnd)).evaluate(simResults, new EvaluationEnvironment());
 
     final var expected = new Windows()
         .set(Interval.between(1, Inclusive, 4, Exclusive, SECONDS), true)
@@ -1092,15 +1092,15 @@ public void testShiftWindowsEdgesBoundsAdjustment() {
 
     final var crossingStartOfPlan = new Windows(false).set(Interval.between(-1, 1, SECONDS), true);
 
-    final var result1 = new ShiftWindowsEdges(
+    final var result1 = new ShiftEdges<>(
         Supplier.of(crossingStartOfPlan),
         Supplier.of(Duration.ZERO),
         Supplier.of(Duration.of(10, SECONDS))
     ).evaluate(simResults);
     final var expected1 = new Windows(false).set(Interval.between(-1, 11, SECONDS), true).select(simResults.bounds);
     assertIterableEquals(expected1, result1);
 
-    final var result2 = new ShiftWindowsEdges(
+    final var result2 = new ShiftEdges<>(
         Supplier.of(crossingStartOfPlan),
         Supplier.of(Duration.of(-10, SECONDS)),
         Supplier.of(Duration.ZERO)
@@ -1110,15 +1110,15 @@ public void testShiftWindowsEdgesBoundsAdjustment() {
 
     final var crossingEndOfPlan = new Windows(false).set(Interval.between(19, 21, SECONDS), true);
 
-    final var result3 = new ShiftWindowsEdges(
+    final var result3 = new ShiftEdges<>(
         Supplier.of(crossingEndOfPlan),
         Supplier.of(Duration.ZERO),
         Supplier.of(Duration.of(10, SECONDS))
     ).evaluate(simResults);
     final var expected3 = new Windows(false).set(Interval.between(19, 20, SECONDS), true).select(simResults.bounds);
     assertIterableEquals(expected3, result3);
 
-    final var result4 = new ShiftWindowsEdges(
+    final var result4 = new ShiftEdges<>(
         Supplier.of(crossingEndOfPlan),
         Supplier.of(Duration.of(-10, SECONDS)),
         Supplier.of(Duration.ZERO)
@@ -1251,6 +1251,91 @@ public void tesRollingThresholdDeficit() {
     assertEquals(expected2, result2);
   }
 
+  @Test
+  void testSpansShiftEdges() {
+    final var simResults = new SimulationResults(
+        Instant.EPOCH, Interval.between(0, 20, SECONDS),
+        List.of(),
+        Map.of(),
+        Map.of()
+    );
+
+    final var result1 = new ShiftEdges<>(
+        Supplier.of(new Spans(
+            Interval.between(0, 1, SECONDS),
+            Interval.between(0, 2, SECONDS),
+            Interval.between(0, Inclusive, 2, Exclusive, SECONDS),
+            Interval.between(0, 3, SECONDS)
+        )),
+        Supplier.of(Duration.of(1, SECONDS)),
+        Supplier.of(Duration.of(-1, SECONDS))
+    ).evaluate(simResults);
+
+    final var expected1 = new Spans(
+        Interval.at(1, SECONDS),
+        Interval.between(1, 2, SECONDS)
+    );
+
+    assertIterableEquals(expected1, result1);
+
+    final var result2 = new ShiftEdges<>(
+        Supplier.of(new Spans(
+            Interval.between(0, 1, SECONDS),
+            Interval.between(0, 2, SECONDS),
+            Interval.between(0, Inclusive, 2, Exclusive, SECONDS),
+            Interval.between(0, 3, SECONDS)
+        )),
+        Supplier.of(Duration.of(3, SECONDS)),
+        Supplier.of(Duration.of(4, SECONDS))
+    ).evaluate(simResults);
+
+    final var expected2 = new Spans(
+        Interval.between(3, 5, SECONDS),
+        Interval.between(3, 6, SECONDS),
+        Interval.between(3, Inclusive, 6, Exclusive, SECONDS),
+        Interval.between(3, 7, SECONDS)
+    );
+
+    assertIterableEquals(expected2, result2);
+  }
+
+  @Test
+  void testSpansSelectWhenTrue() {
+    final var simResults = new SimulationResults(
+        Instant.EPOCH, Interval.between(0, 20, SECONDS),
+        List.of(),
+        Map.of(),
+        Map.of()
+    );
+
+    final var result = new SpansSelectWhenTrue(
+        Supplier.of(new Spans(
+            Interval.between(0, 1, SECONDS), // fully inside
+            Interval.between(3, 4, SECONDS), // fully outside
+            Interval.between(5, 7, SECONDS), // half outside
+            Interval.between(10, 14, SECONDS) // split across multiple
+        )),
+        Supplier.of(new Windows(false).set(
+            List.of(
+                Interval.between(0, 1, SECONDS),
+                Interval.between(6, 8, SECONDS),
+                Interval.between(9, 11, SECONDS),
+                Interval.between(13, 15, SECONDS)
+            ),
+            true
+        ))
+    ).evaluate(simResults);
+
+    final var expected = new Spans(
+        Interval.between(0, 1, SECONDS),
+        Interval.between(6, 7, SECONDS),
+        Interval.between(10, 11, SECONDS),
+        Interval.between(13, 14, SECONDS)
+    );
+
+    assertIterableEquals(expected, result);
+  }
+
   /**
    * An expression that yields the same aliased object every time it is evaluated.
    */