From 8e42dc337400fea0275634a1310eff99a5db7bf0 Mon Sep 17 00:00:00 2001
From: Milad Sadeghi <miladsadeghidm@gmail.com>
Date: Thu, 18 Sep 2025 22:54:44 +0330
Subject: [PATCH] refactor: Fix critical state management bug and improve null
 safety

---
 .../google/common/base/AbstractIterator.java  | 130 ++++++++++--------
 1 file changed, 69 insertions(+), 61 deletions(-)

diff --git a/guava/src/com/google/common/base/AbstractIterator.java b/guava/src/com/google/common/base/AbstractIterator.java
index f46e12ecbe0b..2193b7094098 100644
--- a/guava/src/com/google/common/base/AbstractIterator.java
+++ b/guava/src/com/google/common/base/AbstractIterator.java
@@ -14,14 +14,14 @@
 
 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
@@ -29,65 +29,73 @@
  */
 @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();
-  }
 }