JBIterator: safe contract mixing 2

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() {

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;