Remove nullness suppression in trigger implementation

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterAllStateMachine.java

@@ -28,9 +28,6 @@
  * A {@link TriggerStateMachine} that fires and finishes once after all of its sub-triggers have
  * fired.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class AfterAllStateMachine extends TriggerStateMachine {
 
   private AfterAllStateMachine(List<TriggerStateMachine> subTriggers) {
@@ -117,7 +114,7 @@ public void onFire(TriggerContext context) throws Exception {
   @Override
   public String toString() {
     StringBuilder builder = new StringBuilder("AfterAll.of(");
-    Joiner.on(", ").appendTo(builder, subTriggers);
+    Joiner.on(", ").appendTo(builder, subTriggers());
     builder.append(")");
     return builder.toString();
   }

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterDelayFromFirstElementStateMachine.java

@@ -32,6 +32,7 @@
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.checkerframework.dataflow.qual.Pure;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 import org.joda.time.format.PeriodFormat;
@@ -45,9 +46,6 @@
  */
 // This class should be inlined to subclasses and deleted, simplifying them too
 // https://github.com/apache/beam/issues/18117
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public abstract class AfterDelayFromFirstElementStateMachine extends TriggerStateMachine {
 
   protected static final List<SerializableFunction<Instant, Instant>> IDENTITY = ImmutableList.of();
@@ -61,6 +59,7 @@ public abstract class AfterDelayFromFirstElementStateMachine extends TriggerStat
   private static final PeriodFormatter PERIOD_FORMATTER = PeriodFormat.wordBased(Locale.ENGLISH);
 
   /** To complete an implementation, return the desired time from the TriggerContext. */
+  @Pure
   public abstract @Nullable Instant getCurrentTime(TriggerStateMachine.TriggerContext context);
 
   /**

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterEachStateMachine.java

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.core.triggers;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 
 import java.util.Arrays;
@@ -41,9 +42,6 @@
  *       Repeatedly.forever(a)}, since the repeated trigger never finishes.
  * </ul>
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class AfterEachStateMachine extends TriggerStateMachine {
 
   private AfterEachStateMachine(List<TriggerStateMachine> subTriggers) {
@@ -72,7 +70,7 @@ public static TriggerStateMachine inOrder(Iterable<? extends TriggerStateMachine
   public void onElement(OnElementContext c) throws Exception {
     if (!c.trigger().isMerging()) {
       // If merges are not possible, we need only run the first unfinished subtrigger
-      c.trigger().firstUnfinishedSubTrigger().invokeOnElement(c);
+      firstUnfinishedSubTrigger(c).invokeOnElement(c);
     } else {
       // If merges are possible, we need to run all subtriggers in parallel
       for (ExecutableTriggerStateMachine subTrigger : c.trigger().subTriggers()) {
@@ -119,13 +117,12 @@ public void prefetchShouldFire(PrefetchContext c) {
 
   @Override
   public boolean shouldFire(TriggerStateMachine.TriggerContext context) throws Exception {
-    ExecutableTriggerStateMachine firstUnfinished = context.trigger().firstUnfinishedSubTrigger();
-    return firstUnfinished.invokeShouldFire(context);
+    return firstUnfinishedSubTrigger(context).invokeShouldFire(context);
   }
 
   @Override
   public void onFire(TriggerStateMachine.TriggerContext context) throws Exception {
-    context.trigger().firstUnfinishedSubTrigger().invokeOnFire(context);
+    firstUnfinishedSubTrigger(context).invokeOnFire(context);
 
     // Reset all subtriggers if in a merging context; any may be revived by merging so they are
     // all run in parallel for each pending pane.
@@ -141,12 +138,25 @@ public void onFire(TriggerStateMachine.TriggerContext context) throws Exception
   @Override
   public String toString() {
     StringBuilder builder = new StringBuilder("AfterEach.inOrder(");
-    Joiner.on(", ").appendTo(builder, subTriggers);
+    Joiner.on(", ").appendTo(builder, subTriggers());
     builder.append(")");
 
     return builder.toString();
   }
 
+  /**
+   * Returns the subtrigger that this {@link AfterEachStateMachine} is currently executing.
+   *
+   * <p>A composite trigger is never invoked once all of its subtriggers are finished, so a null
+   * here means a caller has violated that invariant.
+   */
+  private ExecutableTriggerStateMachine firstUnfinishedSubTrigger(TriggerContext context) {
+    return checkStateNotNull(
+        context.trigger().firstUnfinishedSubTrigger(),
+        "%s invoked after all of its subtriggers finished",
+        this);
+  }
+
   private void updateFinishedState(TriggerContext context) {
     context.trigger().setFinished(context.trigger().firstUnfinishedSubTrigger() == null);
   }

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachine.java

@@ -28,9 +28,6 @@
  * Create a composite {@link TriggerStateMachine} that fires once after at least one of its
  * sub-triggers have fired.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class AfterFirstStateMachine extends TriggerStateMachine {
 
   AfterFirstStateMachine(List<TriggerStateMachine> subTriggers) {
@@ -114,7 +111,7 @@ public void onFire(TriggerContext context) throws Exception {
   @Override
   public String toString() {
     StringBuilder builder = new StringBuilder("AfterFirst.of(");
-    Joiner.on(", ").appendTo(builder, subTriggers);
+    Joiner.on(", ").appendTo(builder, subTriggers());
     builder.append(")");
 
     return builder.toString();

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachine.java

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.core.triggers;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 
 import java.util.Objects;
@@ -51,9 +52,6 @@
  * AfterWatermark.pastEndOfWindow.withEarlyFirings(OnceTrigger)} or {@code
  * AfterWatermark.pastEndOfWindow.withEarlyFirings(OnceTrigger)}.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class AfterWatermarkStateMachine {
 
   private static final String TO_STRING = "AfterWatermark.pastEndOfWindow()";
@@ -77,7 +75,7 @@ public static class AfterWatermarkEarlyAndLate extends TriggerStateMachine {
 
     @SuppressWarnings("unchecked")
     private AfterWatermarkEarlyAndLate(
-        TriggerStateMachine earlyTrigger, TriggerStateMachine lateTrigger) {
+        TriggerStateMachine earlyTrigger, @Nullable TriggerStateMachine lateTrigger) {
       super(
           lateTrigger == null
               ? ImmutableList.of(earlyTrigger)
@@ -109,7 +107,11 @@ public void onElement(OnElementContext c) throws Exception {
 
       if (!c.trigger().isMerging()) {
         // If merges can never happen, we just run the unfinished subtrigger
-        c.trigger().firstUnfinishedSubTrigger().invokeOnElement(c);
+        checkStateNotNull(
+                c.trigger().firstUnfinishedSubTrigger(),
+                "%s invoked after all of its subtriggers finished",
+                this)
+            .invokeOnElement(c);
       } else {
         // If merges can happen, we run for all subtriggers because they might be
         // de-activated or re-activated

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/DefaultTriggerStateMachine.java

@@ -24,9 +24,6 @@
  * {@link RepeatedlyStateMachine#forever} and {@link AfterWatermarkStateMachine#pastEndOfWindow} for
  * more details.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class DefaultTriggerStateMachine extends TriggerStateMachine {
 
   private DefaultTriggerStateMachine() {

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/ExecutableTriggerStateMachine.java

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.core.triggers;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
@@ -29,9 +30,6 @@
  * times (both in the same trigger expression and in other trigger expressions), the {@code
  * ExecutableTrigger} wrapped around them forms a tree (only one occurrence).
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class ExecutableTriggerStateMachine implements Serializable {
 
   /** Store the index assigned to this trigger. */
@@ -60,12 +58,10 @@ private ExecutableTriggerStateMachine(TriggerStateMachine trigger, int nextUnuse
     this.trigger = checkNotNull(trigger, "trigger must not be null");
     this.triggerIndex = nextUnusedIndex++;
 
-    if (trigger.subTriggers() != null) {
-      for (TriggerStateMachine subTrigger : trigger.subTriggers()) {
-        ExecutableTriggerStateMachine subExecutable = create(subTrigger, nextUnusedIndex);
-        subTriggers.add(subExecutable);
-        nextUnusedIndex = subExecutable.firstIndexAfterSubtree;
-      }
+    for (TriggerStateMachine subTrigger : trigger.subTriggers()) {
+      ExecutableTriggerStateMachine subExecutable = create(subTrigger, nextUnusedIndex);
+      subTriggers.add(subExecutable);
+      nextUnusedIndex = subExecutable.firstIndexAfterSubtree;
     }
     firstIndexAfterSubtree = nextUnusedIndex;
   }
@@ -99,18 +95,19 @@ public boolean isCompatible(ExecutableTriggerStateMachine other) {
   }
 
   public ExecutableTriggerStateMachine getSubTriggerContaining(int index) {
-    checkNotNull(subTriggers);
     checkState(
         index > triggerIndex && index < firstIndexAfterSubtree,
         "Cannot find sub-trigger containing index not in this tree.");
     ExecutableTriggerStateMachine previous = null;
     for (ExecutableTriggerStateMachine subTrigger : subTriggers) {
       if (index < subTrigger.triggerIndex) {
-        return previous;
+        break;
       }
       previous = subTrigger;
     }
-    return previous;
+    // The bounds check above places the index within this subtree but past this trigger itself,
+    // so it always falls inside one of the sub-triggers.
+    return checkStateNotNull(previous, "No sub-trigger of %s contains index %s", trigger, index);
   }
 
   public void invokePrefetchOnElement(TriggerStateMachine.PrefetchContext c) {

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/OrFinallyStateMachine.java

@@ -23,9 +23,6 @@
 /**
  * Executes the {@code actual} trigger until it finishes or until the {@code until} trigger fires.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 class OrFinallyStateMachine extends TriggerStateMachine {
 
   private static final int ACTUAL = 0;
@@ -96,7 +93,7 @@ public void onFire(TriggerStateMachine.TriggerContext context) throws Exception
 
   @Override
   public String toString() {
-    return String.format("%s.orFinally(%s)", subTriggers.get(ACTUAL), subTriggers.get(UNTIL));
+    return String.format("%s.orFinally(%s)", subTriggers().get(ACTUAL), subTriggers().get(UNTIL));
   }
 
   private void updateFinishedState(TriggerContext c) throws Exception {

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/RepeatedlyStateMachine.java

@@ -31,9 +31,6 @@
  * <p>{@code Repeatedly.forever(someTrigger)} behaves like an infinite {@code
  * AfterEach.inOrder(someTrigger, someTrigger, someTrigger, ...)}.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class RepeatedlyStateMachine extends TriggerStateMachine {
 
   private static final int REPEATED = 0;
@@ -97,7 +94,7 @@ public void onFire(TriggerContext context) throws Exception {
 
   @Override
   public String toString() {
-    return String.format("Repeatedly.forever(%s)", subTriggers.get(REPEATED));
+    return String.format("Repeatedly.forever(%s)", subTriggers().get(REPEATED));
   }
 
   private ExecutableTriggerStateMachine getRepeated(TriggerContext context) {

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/TriggerStateMachine.java

@@ -18,6 +18,7 @@
 package org.apache.beam.runners.core.triggers;
 
 import java.io.Serializable;
+import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import org.apache.beam.runners.core.MergingStateAccessor;
@@ -27,7 +28,9 @@
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Joiner;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.checkerframework.dataflow.qual.Pure;
 import org.joda.time.Instant;
 
 /**
@@ -93,9 +96,6 @@
  * invocations of the callbacks. All important values should be persisted using state before the
  * callback returns.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public abstract class TriggerStateMachine implements Serializable {
 
   /**
@@ -130,7 +130,11 @@ public interface TriggerInfo {
     /** Returns an iterable over the unfinished sub-triggers of the current trigger. */
     Iterable<ExecutableTriggerStateMachine> unfinishedSubTriggers();
 
-    /** Returns the first unfinished sub-trigger. */
+    /**
+     * Returns the first unfinished sub-trigger, or {@code null} if all sub-triggers of the current
+     * trigger are finished.
+     */
+    @Nullable
     ExecutableTriggerStateMachine firstUnfinishedSubTrigger();
 
     /**
@@ -192,6 +196,7 @@ public abstract static class TriggerContext {
     public abstract @Nullable Instant currentSynchronizedProcessingTime();
 
     /** The current event time for the input or {@code null} if unknown. */
+    @Pure
     public abstract @Nullable Instant currentEventTime();
   }
 
@@ -346,8 +351,8 @@ public void clear(TriggerContext c) throws Exception {
     }
   }
 
-  public Iterable<TriggerStateMachine> subTriggers() {
-    return subTriggers;
+  public List<TriggerStateMachine> subTriggers() {
+    return MoreObjects.firstNonNull(subTriggers, Collections.emptyList());
   }
 
   /** Returns whether this performs the same triggering as the given {@code Trigger}. */