From 5bbe4d4e42f9605c45cbca7066222fef5b5eee77 Mon Sep 17 00:00:00 2001
From: Chris Jester-Young <cky@cky.nz>
Date: Mon, 6 Apr 2020 20:19:10 -0400
Subject: [PATCH] Make `Iterators.peekingIterator()` reuse
 `AbstractIterator.peek()`.

Currently, `Iterators.peekingIterator()` will treat `AbstractIterator`
just like any other iterator, and will call `next()` to get the element,
which it then stashes. This is wasteful, since `AbstractIterator` already
provides a `peek()` method that is designed to be compatible with the
`PeekingIterator` interface, so use it instead.

`AbstractIterator` subclasses `UnmodifiableIterator`, so we can do
likewise with the `PeekingIterator` forwarder. The resulting forwarder
is much more lightweight than the default `Iterators.peekingIterator()`
implementation.
---
 .../common/collect/PeekingIteratorTest.java   | 14 +++++++
 .../com/google/common/collect/Iterators.java  | 40 +++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/guava-tests/test/com/google/common/collect/PeekingIteratorTest.java b/guava-tests/test/com/google/common/collect/PeekingIteratorTest.java
index 0f79985bcf39..5d6a80f5a0e4 100644
--- a/guava-tests/test/com/google/common/collect/PeekingIteratorTest.java
+++ b/guava-tests/test/com/google/common/collect/PeekingIteratorTest.java
@@ -257,4 +257,18 @@ private void assertNextThrows(Iterator<?> iterator) {
     } catch (ThrowsAtEndException expected) {
     }
   }
+
+  private static class UniqueObjectIterator extends AbstractIterator<Object> {
+    @Override
+    protected Object computeNext() {
+      return new Object();
+    }
+  }
+
+  public void testPeekingAbstractIteratorDoesntAdvancePrematurely() {
+    Iterator<Object> iterator = new UniqueObjectIterator();
+    PeekingIterator<Object> peeking1 = Iterators.peekingIterator(iterator);
+    PeekingIterator<Object> peeking2 = Iterators.peekingIterator(iterator);
+    assertSame(peeking1.peek(), peeking2.peek());
+  }
 }
diff --git a/guava/src/com/google/common/collect/Iterators.java b/guava/src/com/google/common/collect/Iterators.java
index f655c279de53..4c3bcf5a28b9 100644
--- a/guava/src/com/google/common/collect/Iterators.java
+++ b/guava/src/com/google/common/collect/Iterators.java
@@ -1154,6 +1154,30 @@ public E peek() {
     }
   }
 
+  private static class AbstractIteratorPeekingImpl<E> extends UnmodifiableIterator<E>
+      implements PeekingIterator<E> {
+    private final AbstractIterator<? extends E> iterator;
+
+    public AbstractIteratorPeekingImpl(AbstractIterator<? extends E> iterator) {
+      this.iterator = checkNotNull(iterator);
+    }
+
+    @Override
+    public boolean hasNext() {
+      return iterator.hasNext();
+    }
+
+    @Override
+    public E next() {
+      return iterator.next();
+    }
+
+    @Override
+    public E peek() {
+      return iterator.peek();
+    }
+  }
+
   /**
    * Returns a {@code PeekingIterator} backed by the given iterator.
    *
@@ -1198,6 +1222,22 @@ public static <T> PeekingIterator<T> peekingIterator(Iterator<? extends T> itera
       PeekingImpl<T> peeking = (PeekingImpl<T>) iterator;
       return peeking;
     }
+    if (iterator instanceof AbstractIteratorPeekingImpl) {
+      // Safe to cast <? extends T> to <T> because AbstractIteratorPeekingImpl
+      // only uses T covariantly (and cannot be subclassed to add non-covariant
+      // uses).
+      @SuppressWarnings("unchecked")
+      AbstractIteratorPeekingImpl<T> peeking = (AbstractIteratorPeekingImpl<T>) iterator;
+      return peeking;
+    }
+    return peekingImpl(iterator);
+  }
+
+  /* Captures the iterator's element type so the downcast does not cause warnings. */
+  private static <T, E extends T> PeekingIterator<T> peekingImpl(Iterator<E> iterator) {
+    if (iterator instanceof AbstractIterator) {
+      return new AbstractIteratorPeekingImpl<T>((AbstractIterator<E>) iterator);
+    }
     return new PeekingImpl<T>(iterator);
   }