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<Index> indexOfP(JsonParser<Expression<Windows>> windowsExpressionP) {
+    return productP
+        .field("kind", literalP("WindowsExpressionIndex"))
+        .field("expression", windowsExpressionP)
+        .field("index", intP)
+        .map(
+            untuple((kind, expression, index) -> new Index(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),
+        indexOfP(selfP)
     ));
   }
 

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

@@ -0,0 +1,71 @@
+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 Index implements Expression<Windows> {
+
+  public final Expression<Windows> expression;
+  //can be negative
+  public final int i;
+
+  //i must not be 0
+  public Index(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) {
+    final var intervals = this.expression.evaluate(results, bounds, environment);
+    final var itForward = intervals.iterateEqualTo(true).iterator();
+    var count = 0;
+    while(itForward.hasNext()){
+      count++;
+      final var current = itForward.next();
+      if(i == count){
+        return new Windows().set(current, true);
+      }
+    }
+    if(i < 0){
+      final var realIndex = count + i + 1;
+      final var itBackward = intervals.iterateEqualTo(true).iterator();
+      count = 0;
+      while(itBackward.hasNext()){
+        count++;
+        final var current = itBackward.next();
+        if(realIndex == count){
+          return new Windows().set(current, true);
+        }
+      }
+    }
+    return new Windows();
+  }
+
+  @Override
+  public String prettyPrint(final String prefix) {
+    return "";
+  }
+
+  @Override
+  public void extractResources(final Set<String> names) {
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (!(obj instanceof final Index 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

@@ -1335,6 +1335,35 @@ void testSpansSelectWhenTrue() {
 
     assertIterableEquals(expected, result);
   }
+  @Test
+  public void testIndex(){
+    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 Index(Supplier.of(windows), 3).evaluate(simResults, new EvaluationEnvironment());
+    final var result2 = new Index(Supplier.of(windows), -2).evaluate(simResults, new EvaluationEnvironment());
+    final var result3 = new Index(Supplier.of(windows), 5).evaluate(simResults, new EvaluationEnvironment());
+
+    final var expected = new Windows()
+        .set(Interval.between(10, Exclusive, 15, Exclusive, SECONDS), true);
+
+    assertEquivalent(expected, result);
+    assertEquivalent(expected, result2);
+    assertEquivalent(new Windows(), result3);
+  }
 
   /**
    * An expression that yields the same aliased object every time it is evaluated.

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

@@ -49,7 +49,8 @@ export enum NodeKind {
   AbsoluteInterval = 'AbsoluteInterval',
   IntervalAlias = 'IntervalAlias',
   IntervalDuration = 'IntervalDuration',
-  RollingThreshold = 'RollingThreshold'
+  RollingThreshold = 'RollingThreshold',
+  WindowsExpressionIndex = 'WindowsExpressionIndex'
 }
 
 export type Constraint = ViolationsOf | WindowsExpression | SpansExpression | ForEachActivityConstraints | RollingThreshold;
@@ -112,7 +113,9 @@ export type WindowsExpression =
   | WindowsExpressionFromSpans
   | IntervalsExpressionStarts
   | IntervalsExpressionEnds
-  | AssignGapsExpression<WindowsExpression>;
+  | AssignGapsExpression<WindowsExpression>
+    | WindowsExpressionIndex;
+
 
 export type SpansExpression =
   | SpansExpressionActivitySpan
@@ -157,6 +160,12 @@ export interface WindowsExpressionNot {
   expression: WindowsExpression;
 }
 
+export interface WindowsExpressionIndex {
+  kind: NodeKind.WindowsExpressionIndex;
+  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 a windows by its index
+   * @param i the index of the sequence of windows
+   */
+  public index(i: number): Windows {
+    return new Windows({
+      kind: AST.NodeKind.WindowsExpressionIndex,
+      expression: this.__astNode,
+      index: i
+    })
+  }
+
   /**
    * Performs the boolean Or operation on any number of Windows.
    *
@@ -1170,6 +1182,12 @@ declare global {
     public readonly __astNode: AST.WindowsExpression;
 
     /**
+     * Selects a windows by its index
+     * @param i the index of the sequence of windows
+     */
+    public index(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

@@ -20,6 +20,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.Index;
 import gov.nasa.jpl.aerie.constraints.tree.LessThan;
 import gov.nasa.jpl.aerie.constraints.tree.LessThanOrEqual;
 import gov.nasa.jpl.aerie.constraints.tree.LongerThan;
@@ -850,6 +851,19 @@ export default () => {
         )
     );
   }
+@Test
+  void testIndex() {
+    checkSuccessfulCompilation(
+        """
+          export default () => {
+            return Real.Resource("state of charge").lessThan(0.3).starts().index(2)
+          }
+        """,
+        new ViolationsOfWindows(
+            new Index(new Starts<>(new LessThan(new RealResource("state of charge"), new RealValue(0.3))), 2)
+        )
+    );
+  }
 
   @Test
   void testEnds() {