From f24dc012e3cd55d28783545fbfaca77315ebe194 Mon Sep 17 00:00:00 2001
From: "Gregory.Shrago" <gregory.shrago@jetbrains.com>
Date: Tue, 1 Dec 2015 00:45:55 +0300
Subject: [PATCH] JBIterable: introduce get(index)

---
 .../util/containers/TreeTraverserTest.java       | 11 +++++++++++
 .../com/intellij/util/containers/JBIterable.java | 16 ++++++++++++++++
 2 files changed, 27 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 e921579bb52a..daa631d312c1 100644
--- a/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
+++ b/platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java
@@ -177,6 +177,17 @@ public class TreeTraverserTest extends TestCase {
     assertEquals(Arrays.asList(1, 2, 3, 4), it.toList());
   }
 
+  public void testGetAt() {
+    JBIterable<Integer> it = JBIterable.of(1, 2, 3, 4);
+    assertEquals((Integer)4, it.get(3));
+    assertNull(it.get(4));
+    assertNull(it.get(5));
+    JBIterable<Integer> it2 = JBIterable.generate(1, INCREMENT).take(4);
+    assertEquals((Integer)4, it2.get(3));
+    assertNull(it2.get(4));
+    assertNull(it2.get(5));
+  }
+
   public void testFilterTransformTakeWhile() {
     JBIterable<Integer> it = JBIterable.generate(1, INCREMENT).filter(IS_ODD).transform(SQUARE).takeWhile(LESS_THAN(100));
     assertEquals(Arrays.asList(1, 9, 25, 49, 81), it.toList());
diff --git a/platform/util/src/com/intellij/util/containers/JBIterable.java b/platform/util/src/com/intellij/util/containers/JBIterable.java
index aa82ccda6540..f154ba558eb6 100644
--- a/platform/util/src/com/intellij/util/containers/JBIterable.java
+++ b/platform/util/src/com/intellij/util/containers/JBIterable.java
@@ -202,6 +202,9 @@ public abstract class JBIterable<E> implements Iterable<E> {
    * Returns the number of elements in this iterable.
    */
   public final int size() {
+    if (myIterable instanceof Collection) {
+      return ((Collection)myIterable).size();
+    }
     int count = 0;
     for (E ignored : myIterable) {
       count++;
@@ -223,6 +226,16 @@ public abstract class JBIterable<E> implements Iterable<E> {
     return false;
   }
 
+  /**
+   * Returns element at index if it is present; otherwise {@code null}
+   */
+  public final E get(int index) {
+    if (myIterable instanceof List) {
+      return index >= ((List)myIterable).size() ? null : ((List<E>)myIterable).get(index);
+    }
+    return skip(index).first();
+  }
+
   /**
    * Returns a {@code JBIterable} whose iterators traverse first the elements of this iterable,
    * followed by those of {@code other}. The iterators are not polled until necessary.
@@ -412,6 +425,9 @@ public abstract class JBIterable<E> implements Iterable<E> {
    * Determines whether this iterable is empty.
    */
   public final boolean isEmpty() {
+    if (myIterable instanceof Collection) {
+      return ((Collection)myIterable).isEmpty();
+    }
     return !myIterable.iterator().hasNext();
   }