From 80c8f8796a0e4e94d00a2607b7a11ad90fa0fa14 Mon Sep 17 00:00:00 2001
From: "Gregory.Shrago" <gregory.shrago@jetbrains.com>
Date: Fri, 22 Apr 2016 01:46:16 +0300
Subject: [PATCH] JBIterator: safe contract mixing 2

---
 .../util/containers/TreeTraverserTest.java         |  8 ++++++++
 .../com/intellij/util/containers/JBIterator.java   | 14 +++++++++-----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java b/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
index 47695d38983f..2e02c0959709 100644
--- a/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
+++ b/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
@@ -208,6 +208,14 @@ public class TreeTraverserTest extends TestCase {
     assertFalse(it.hasNext());
   }
 
+  public void testIteratorContractsCurrent() {
+    JBIterator<Integer> it = JBIterator.from(JBIterable.of(1).iterator());
+    assertTrue(it.advance());
+    assertEquals(new Integer(1), it.current());
+    assertFalse(it.hasNext());
+    assertEquals(new Integer(1), it.current());
+  }
+
   // JBIterable ----------------------------------------------
 
   public void testAppend() {
diff --git a/platform/util/src/com/intellij/util/containers/JBIterator.java b/platform/util/src/com/intellij/util/containers/JBIterator.java
index 16f5486310e7..19617bf7720d 100644
--- a/platform/util/src/com/intellij/util/containers/JBIterator.java
+++ b/platform/util/src/com/intellij/util/containers/JBIterator.java
@@ -81,12 +81,16 @@ public abstract class JBIterator<E> implements Iterator<E> {
    */
   protected abstract E nextImpl();
 
+  /**
+   * Called right after the new current value is set.
+   */
+  protected void currentChanged() { }
+
   /**
    * Notifies the iterator that there's no more elements.
    */
   protected final E stop() {
     myNext = STOP;
-    myCurrent = null;
     return null;
   }
 
@@ -95,7 +99,6 @@ public abstract class JBIterator<E> implements Iterator<E> {
    */
   protected final E skip() {
     myNext = SKIP;
-    myCurrent = null;
     return null;
   }
 
@@ -107,10 +110,12 @@ public abstract class JBIterator<E> implements Iterator<E> {
 
   @Override
   public final E next() {
+    myCurrent = NONE;
     peekNext();
     if (myNext == STOP) throw new NoSuchElementException();
     myCurrent = myNext;
     myNext = NONE;
+    currentChanged();
     return (E)myCurrent;
   }
 
@@ -123,21 +128,20 @@ public abstract class JBIterator<E> implements Iterator<E> {
     if (myNext == STOP) return false;
     myCurrent = myNext;
     myNext = NONE;
+    currentChanged();
     return true;
   }
 
   /**
    * Returns the current element if any; otherwise throws exception.
    */
-  public E current() {
+  public final E current() {
     if (myCurrent == NONE) throw new NoSuchElementException();
-    if (myNext == STOP) throw new NoSuchElementException();
     return (E)myCurrent;
   }
 
   private void peekNext() {
     if (myNext != NONE) return;
-    myNext = null;
     Object o = nextImpl();
     if (myNext == STOP) return;
     Op op = myFirstOp.nextOp;