From 216b92b438cb3796fe5e4d03fd7aaadcb87b4862 Mon Sep 17 00:00:00 2001
From: "Gregory.Shrago" <gregory.shrago@jetbrains.com>
Date: Fri, 29 Jul 2016 19:42:47 +0300
Subject: [PATCH] introduce: map, reduce, find, indexOf and flatMap

---
 .../util/containers/TreeTraverserTest.java    | 10 ++++
 .../intellij/util/containers/JBIterable.java  | 52 +++++++++++++++++++
 2 files changed, 62 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 11217739b939..372c0c071fe1 100644
--- a/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
+++ b/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
@@ -289,6 +289,16 @@ public class TreeTraverserTest extends TestCase {
     assertEquals(Arrays.asList(1, 1, 2, 3, 5, 8, 13, 21), it.toList());
   }
 
+  public void testFindIndexReduceMap() {
+    JBIterable<Integer> it = JBIterable.of(1, 2, 3, 4, 5);
+    assertEquals(15, (int)it.reduce(0, (Integer v, Integer o) -> v + o));
+    assertEquals(3, (int)it.find((o)-> o.intValue() == 3));
+    assertEquals(2, it.indexOf((o)-> o.intValue() == 3));
+    assertEquals(-1, it.indexOf((o)-> o.intValue() == 33));
+    assertEquals(Arrays.asList(1, 4, 9, 16, 25), it.map(o -> o * o).toList());
+    assertEquals(Arrays.asList(0, 1, 0, 2, 0, 3, 0, 4, 0, 5), it.flatMap(o -> ContainerUtil.list(0, o)).toList());
+  }
+
   // TreeTraversal ----------------------------------------------
 
   @NotNull
diff --git a/platform/util/src/com/intellij/util/containers/JBIterable.java b/platform/util/src/com/intellij/util/containers/JBIterable.java
index 92eb82980572..ca6b0f8a414b 100644
--- a/platform/util/src/com/intellij/util/containers/JBIterable.java
+++ b/platform/util/src/com/intellij/util/containers/JBIterable.java
@@ -430,6 +430,58 @@ public abstract class JBIterable<E> implements Iterable<E> {
     return cur;
   }
 
+  /**
+   * Perform calculation over this iterable.
+   */
+  public final <T> T reduce(@Nullable T first, @NotNull PairFunction<T, E, T> function) {
+    T cur = first;
+    for (E e : this) {
+      cur = function.fun(cur, e);
+    }
+    return cur;
+  }
+
+  /**
+   * Returns the index of the first matching element.
+   */
+  public final E find(@NotNull Condition<E> condition) {
+    return filter(condition).first();
+  }
+
+  /**
+   * Returns the index of the matching element.
+   */
+  public final int indexOf(@NotNull Condition<E> condition) {
+    int index = 0;
+    for (E e : this) {
+      if (condition.value(e)) {
+        return index;
+      }
+      index ++;
+    }
+    return -1;
+  }
+
+  /**
+   * Synonym for transform()
+   *
+   * @see JBIterable#transform(Function)
+   */
+  @NotNull
+  public final <T> JBIterable<T> map(@NotNull Function<? super E, T> function) {
+    return transform(function);
+  }
+
+  /**
+   * "Maps" and "flattens" this iterable.
+   *
+   * @see JBIterable#transform(Function)
+   */
+  @NotNull
+  public final <T> JBIterable<T> flatMap(Function<? super E, ? extends Iterable<? extends T>> function) {
+    return map(function).flatten(Function.ID);
+  }
+
   /**
    * Determines whether this iterable is empty.
    */