JBIterator: rollback prepended take(N) counters if nextImpl() votes SKIP

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);

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;
       }
     }