From 82eac480cf72b9eb8eb6df8ea442345b668f17fe Mon Sep 17 00:00:00 2001
From: "Gregory.Shrago" <gregory.shrago@jetbrains.com>
Date: Tue, 8 Aug 2017 00:10:00 +0300
Subject: [PATCH] JBIterator: rollback prepended take(N) counters if nextImpl()
 votes SKIP

---
 .../com/intellij/util/containers/TreeTraverserTest.java     | 6 ++++++
 .../util/src/com/intellij/util/containers/JBIterator.java   | 6 ++++++
 2 files changed, 12 insertions(+)

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 f5ef63a9395d..e9b2e31094d4 100644
--- a/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
+++ b/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
@@ -250,6 +250,12 @@ public class TreeTraverserTest extends TestCase {
     assertEquals(new Integer(11), it.first());
   }
 
+  public void testFlattenSkipTake() {
+    assertEquals(1, JBIterable.of(1).flatMap(o -> JBIterable.of(o)).take(1).take(1).take(1).size());
+    assertEquals((Integer)1, JBIterable.of(1).flatMap(o -> JBIterable.of(o, o + 1)).take(2).take(1).get(0));
+    assertEquals((Integer)2, JBIterable.of(1).flatMap(o -> JBIterable.of(o, o + 1)).skip(1).take(1).get(0));
+  }
+
   public void testRangeWithSkipAndTake() {
     Condition<Integer> cond = i -> Math.abs(i - 10) <= 5;
     JBIterable<Integer> it = JBIterable.generate(1, INCREMENT).skipWhile(not(cond)).takeWhile(cond);
diff --git a/platform/util/src/com/intellij/util/containers/JBIterator.java b/platform/util/src/com/intellij/util/containers/JBIterator.java
index eff41eead1ba..630188c294d7 100644
--- a/platform/util/src/com/intellij/util/containers/JBIterator.java
+++ b/platform/util/src/com/intellij/util/containers/JBIterator.java
@@ -151,6 +151,12 @@ public abstract class JBIterator<E> implements Iterator<E> {
       if (myNext == Do.STOP) return;
       if (myNext == Do.SKIP) {
         o = myNext = Do.INIT;
+        if (op.impl == null) {
+          // rollback all prepended takeWhile conditions if nextImpl() votes SKIP
+          for (Op op2 = myFirstOp; op2.impl instanceof CountDown; op2 = op2.nextOp) {
+            ((CountDown)op2.impl).cur ++;
+          }
+        }
         op = null;
       }
     }