Add index constraints

constraints/src/main/java/gov/nasa/jpl/aerie/constraints/json/ConstraintParsers.java

@@ -280,6 +280,15 @@ static JsonParser<SpansSelectWhenTrue> spansSelectWhenTrueF(JsonParser<Expressio
               untuple((kind, alias) -> new StartOf(alias)),
               $ -> tuple(Unit.UNIT, $.activityAlias));
 
+  static JsonParser<KeepTrueSegment> keepTrueSegmentP(JsonParser<Expression<Windows>> windowsExpressionP) {
+    return productP
+        .field("kind", literalP("WindowsExpressionKeepTrueSegment"))
+        .field("expression", windowsExpressionP)
+        .field("index", intP)
+        .map(
+            untuple((kind, expression, index) -> new KeepTrueSegment(expression, index)),
+            $ -> tuple(Unit.UNIT, $.expression, $.i));
+  }
   static final JsonParser<Duration> durationP =
       longP
           .map(
@@ -580,7 +589,8 @@ private static JsonParser<Expression<Windows>> windowsExpressionF(JsonParser<Exp
         windowsFromSpansF(spansP),
         activityWindowP,
         assignGapsF(selfP),
-        shiftByF(selfP)
+        shiftByF(selfP),
+        keepTrueSegmentP(selfP)
     ));
   }
 

constraints/src/main/java/gov/nasa/jpl/aerie/constraints/tree/KeepTrueSegment.java

@@ -0,0 +1,52 @@
+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.Windows;
+
+import java.util.Objects;
+import java.util.Set;
+
+public class KeepTrueSegment implements Expression<Windows> {
+
+  public final Expression<Windows> expression;
+  public final int i;
+
+  public KeepTrueSegment(final Expression<Windows> expression, final int i) {
+    this.expression = expression;
+    this.i = i;
+  }
+
+  @Override
+  public Windows evaluate(final SimulationResults results, final Interval bounds, final EvaluationEnvironment environment) {
+    return this.expression.evaluate(results, bounds, environment).keepTrueSegment(i);
+  }
+
+  @Override
+  public String prettyPrint(final String prefix) {
+    return String.format(
+        "\n%s(%s[%d])",
+        prefix,
+        this.expression.prettyPrint(prefix + "  "),
+        i
+    );  }
+
+  @Override
+  public void extractResources(final Set<String> names) {
+    this.expression.extractResources(names);
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (!(obj instanceof final KeepTrueSegment o)) return false;
+
+    return Objects.equals(this.expression, o.expression) && this.i == o.i;
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(this.expression, this.i);
+  }
+
+}

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

@@ -1397,6 +1397,41 @@ void testSpansSelectWhenTrue() {
 
     assertIterableEquals(expected, result);
   }
+  @Test
+  public void testKeepTrueSegment(){
+    final var simResults = new SimulationResults(
+        Instant.EPOCH, Interval.between(0, 20, SECONDS),
+        List.of(),
+        Map.of(),
+        Map.of()
+    );
+
+    final var windows = new Windows()
+        .set(Interval.between(0, Inclusive, 5, Exclusive, SECONDS), true)
+        .set(Interval.between(6, Inclusive, 7, Inclusive, SECONDS), false)
+        .set(Interval.between(8, Exclusive, 9, Exclusive, SECONDS), true)
+        .set(Interval.between(10, Exclusive, 15, Exclusive, SECONDS), true)
+        .set(Interval.at(20, SECONDS), true)
+        .set(interval(22, 23, SECONDS), false)
+        .set(interval(24, 30, SECONDS), false);
+
+    final var result = new KeepTrueSegment(Supplier.of(windows), 2).evaluate(simResults, new EvaluationEnvironment());
+    final var result2 = new KeepTrueSegment(Supplier.of(windows), -2).evaluate(simResults, new EvaluationEnvironment());
+    final var result3 = new KeepTrueSegment(Supplier.of(windows), 5).evaluate(simResults, new EvaluationEnvironment());
+
+    final var expected = new Windows()
+        .set(Interval.between(0, Inclusive, 5, Exclusive, SECONDS), false)
+        .set(Interval.between(6, Inclusive, 7, Inclusive, SECONDS), false)
+        .set(Interval.between(8, Exclusive, 9, Exclusive, SECONDS), false)
+        .set(Interval.between(10, Exclusive, 15, Exclusive, SECONDS), true)
+        .set(Interval.at(20, SECONDS), false)
+        .set(interval(22, 23, SECONDS), false)
+        .set(interval(24, 30, SECONDS), false);
+
+    assertEquivalent(expected, result);
+    assertEquivalent(expected, result2);
+    assertEquivalent(expected.set(Interval.between(10, Exclusive, 15, Exclusive, SECONDS), false), result3);
+  }
 
   /**
    * An expression that yields the same aliased object every time it is evaluated.

merlin-server/constraints-dsl-compiler/src/libs/constraints-ast.ts

@@ -50,7 +50,8 @@ export enum NodeKind {
   AbsoluteInterval = 'AbsoluteInterval',
   IntervalAlias = 'IntervalAlias',
   IntervalDuration = 'IntervalDuration',
-  RollingThreshold = 'RollingThreshold'
+  RollingThreshold = 'RollingThreshold',
+  WindowsExpressionKeepTrueSegment = 'WindowsExpressionKeepTrueSegment'
 }
 
 export type Constraint = ViolationsOf | WindowsExpression | SpansExpression | ForEachActivityConstraints | RollingThreshold;
@@ -113,7 +114,9 @@ export type WindowsExpression =
   | WindowsExpressionFromSpans
   | IntervalsExpressionStarts
   | IntervalsExpressionEnds
-  | AssignGapsExpression<WindowsExpression>;
+  | AssignGapsExpression<WindowsExpression>
+    | WindowsExpressionKeepTrueSegment;
+
 
 export type SpansExpression =
   | SpansExpressionActivitySpan
@@ -159,6 +162,12 @@ export interface WindowsExpressionNot {
   expression: WindowsExpression;
 }
 
+export interface WindowsExpressionKeepTrueSegment {
+  kind: NodeKind.WindowsExpressionKeepTrueSegment;
+  expression: WindowsExpression;
+  index: number;
+}
+
 export interface IntervalsExpressionShiftEdges {
   kind: NodeKind.IntervalsExpressionShiftEdges,
   expression: IntervalsExpression,

merlin-server/constraints-dsl-compiler/src/libs/constraints-edsl-fluent-api.ts

@@ -163,6 +163,18 @@ export class Windows {
     });
   }
 
+  /**
+   * Selects the ith true window and falsifies the other true segments.
+   * @param i the index of the true segment in the sequence of windows. index(0) will point to the first element, index(-1) to the last element.
+   */
+  public keepTrueSegment(i: number): Windows {
+    return new Windows({
+      kind: AST.NodeKind.WindowsExpressionKeepTrueSegment,
+      expression: this.__astNode,
+      index: i
+    })
+  }
+
   /**
    * Performs the boolean Or operation on any number of Windows.
    *
@@ -1192,6 +1204,12 @@ declare global {
     public readonly __astNode: AST.WindowsExpression;
 
     /**
+     * Selects the ith true window and falsifies the other true segments.
+     * @param i the index of the true segment in the sequence of windows. index(0) will point to the first element, index(-1) to the last element.
+     */
+    public keepTrueSegment(i: number): Windows;
+
+      /**
      * Creates a single window.
      *
      * @param value value for the window segment.

merlin-server/src/test/java/gov/nasa/jpl/aerie/merlin/server/services/ConstraintsDSLCompilationServiceTests.java

@@ -1,6 +1,5 @@
 package gov.nasa.jpl.aerie.merlin.server.services;
 
-import gov.nasa.jpl.aerie.constraints.time.Spans;
 import gov.nasa.jpl.aerie.constraints.tree.AbsoluteInterval;
 import gov.nasa.jpl.aerie.constraints.time.Interval;
 import gov.nasa.jpl.aerie.constraints.tree.AccumulatedDuration;
@@ -20,6 +19,7 @@
 import gov.nasa.jpl.aerie.constraints.tree.ForEachActivityViolations;
 import gov.nasa.jpl.aerie.constraints.tree.GreaterThan;
 import gov.nasa.jpl.aerie.constraints.tree.GreaterThanOrEqual;
+import gov.nasa.jpl.aerie.constraints.tree.KeepTrueSegment;
 import gov.nasa.jpl.aerie.constraints.tree.LessThan;
 import gov.nasa.jpl.aerie.constraints.tree.LessThanOrEqual;
 import gov.nasa.jpl.aerie.constraints.tree.LongerThan;
@@ -851,6 +851,19 @@ export default () => {
         )
     );
   }
+@Test
+  void testKeepTrueSegment() {
+    checkSuccessfulCompilation(
+        """
+          export default () => {
+            return Real.Resource("state of charge").lessThan(0.3).starts().keepTrueSegment(2)
+          }
+        """,
+        new ViolationsOfWindows(
+            new KeepTrueSegment(new Starts<>(new LessThan(new RealResource("state of charge"), new RealValue(0.3))), 2)
+        )
+    );
+  }
 
   @Test
   void testEnds() {