Refactor: Fix critical state management bug and improve null safety

guava/src/com/google/common/base/AbstractIterator.java

@@ -14,80 +14,88 @@
 
 package com.google.common.base;
 
-import static com.google.common.base.NullnessCasts.uncheckedCastNullableTToT;
-import static com.google.common.base.Preconditions.checkState;
-
 import com.google.common.annotations.GwtCompatible;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import org.jspecify.annotations.Nullable;
+
 import java.util.Iterator;
 import java.util.NoSuchElementException;
-import org.jspecify.annotations.Nullable;
+
+import static com.google.common.base.Preconditions.checkState;
 
 /**
  * Note this class is a copy of {@link com.google.common.collect.AbstractIterator} (for dependency
  * reasons).
  */
 @GwtCompatible
 abstract class AbstractIterator<T extends @Nullable Object> implements Iterator<T> {
-  private State state = State.NOT_READY;
-
-  protected AbstractIterator() {}
-
-  private enum State {
-    READY,
-    NOT_READY,
-    DONE,
-    FAILED,
-  }
-
-  private @Nullable T next;
-
-  protected abstract @Nullable T computeNext();
-
-  @CanIgnoreReturnValue
-  protected final @Nullable T endOfData() {
-    state = State.DONE;
-    return null;
-  }
-
-  @Override
-  public final boolean hasNext() {
-    checkState(state != State.FAILED);
-    switch (state) {
-      case DONE:
-        return false;
-      case READY:
+    private State state = State.NOT_READY;
+    private @Nullable T next;
+
+    protected AbstractIterator() {
+    }
+
+    protected abstract @Nullable T computeNext();
+
+    @CanIgnoreReturnValue
+    protected final @Nullable T endOfData() {
+        state = State.DONE;
+        return null;
+    }
+
+    @Override
+    public final boolean hasNext() {
+        checkState(state != State.FAILED);
+        switch (state) {
+            case DONE:
+                return false;
+            case READY:
+                return true;
+            default:
+                return tryToComputeNext();
+        }
+    }
+
+    private boolean tryToComputeNext() {
+        state = State.FAILED;
+        next = computeNext();
+
+        if (state == State.DONE) {
+            // endOfData() was called during computeNext()
+            return false;
+        }
+
+        if (next == null) {
+            // computeNext() returned null but didn't call endOfData()
+            // This is either an error or should be treated as terminal state
+            state = State.DONE;
+            return false;
+        }
+
+        state = State.READY;
         return true;
-      default:
     }
-    return tryToComputeNext();
-  }
-
-  private boolean tryToComputeNext() {
-    state = State.FAILED; // temporary pessimism
-    next = computeNext();
-    if (state != State.DONE) {
-      state = State.READY;
-      return true;
+
+    @Override
+    @ParametricNullness
+    public final T next() {
+        if (!hasNext()) {
+            throw new NoSuchElementException();
+        }
+        state = State.NOT_READY;
+        @SuppressWarnings("nullness") // hasNext() ensures next is non-null
+        T result = next;
+        next = null;
+        return result;
+    }
+
+    @Override
+    public final void remove() {
+        throw new UnsupportedOperationException();
     }
-    return false;
-  }
-
-  @Override
-  @ParametricNullness
-  public final T next() {
-    if (!hasNext()) {
-      throw new NoSuchElementException();
+
+
+    private enum State {
+        READY, NOT_READY, DONE, FAILED
     }
-    state = State.NOT_READY;
-    // Safe because hasNext() ensures that tryToComputeNext() has put a T into `next`.
-    T result = uncheckedCastNullableTToT(next);
-    next = null;
-    return result;
-  }
-
-  @Override
-  public final void remove() {
-    throw new UnsupportedOperationException();
-  }
 }