From c1b4f4713a34cbfb68163c98e963ad4de83040f6 Mon Sep 17 00:00:00 2001
From: "Gregory.Shrago" <gregory.shrago@jetbrains.com>
Date: Thu, 16 Feb 2017 20:50:48 +0300
Subject: [PATCH] make flatten() support nulls

---
 .../com/intellij/util/containers/TreeTraverserTest.java      | 5 +++++
 .../util/src/com/intellij/util/containers/JBIterable.java    | 3 ++-
 2 files changed, 7 insertions(+), 1 deletion(-)

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 9c8386ff1031..bd6828081e3c 100644
--- a/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
+++ b/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
@@ -304,6 +304,11 @@ public class TreeTraverserTest extends TestCase {
     }
   }
 
+  public void testFlatten() {
+    JBIterable<Integer> it = JBIterable.generate(1, INCREMENT).take(3).flatten(i -> i % 2 == 0 ? null : JBIterable.of(i - 1, i));
+    assertEquals(Arrays.asList(0, 1, 2, 3), it.toList());
+  }
+
   public void testStatefulFilter() {
     JBIterable<Integer> it = JBIterable.generate(1, INCREMENT).take(5).filter(new JBIterable.StatefulFilter<Integer>() {
       int prev;
diff --git a/platform/util/src/com/intellij/util/containers/JBIterable.java b/platform/util/src/com/intellij/util/containers/JBIterable.java
index dfd862eae595..41eaa3b19ba0 100644
--- a/platform/util/src/com/intellij/util/containers/JBIterable.java
+++ b/platform/util/src/com/intellij/util/containers/JBIterable.java
@@ -391,7 +391,8 @@ public abstract class JBIterable<E> implements Iterable<E> {
           public T nextImpl() {
             if (cur != null && cur.hasNext()) return cur.next();
             if (!iterator.hasNext()) return stop();
-            cur = fun.fun(iterator.next()).iterator();
+            Iterable<? extends T> next = fun.fun(iterator.next());
+            cur = next == null ? null : next.iterator();
             return skip();
           }
         };