Optimize '.after' method

Author: dykov

src/main/java/io/github/dgroup/enumerable4j/Enumerable.java

@@ -220,32 +220,41 @@ default X reduce(X idn, BinaryOperator<X> opr) {
      * @return The enumerable.
      */
     default Enumerable<X> after(Predicate<X> prd) {
-        return this.after(prd, this.size() - 1);
+        return this.after(prd, this.size());
     }
 
     /**
      * Returns an enumerable containing a certain number of elements of enumerable
      *  after the first one which corresponds the condition.
-     * If no predicate (null) is given, then 'this' is returned instead.
+     * If no predicate (null) is given, then 'this' of the specified size is returned instead.
      * @param prd The function to match element after which enumerable elements should be returned.
      * @param size The number of elements the enumerable should be limited to.
      * @return The enumerable.
      */
     default Enumerable<X> after(Predicate<X> prd, long size) {
         final Enumerable<X> out;
-        if (prd == null) {
-            out = this;
+        if (size < 0) {
+            throw new IllegalArgumentException(Long.toString(size));
+        } else if (size == 0) {
+            out = new Empty<>();
+        } else if (prd == null) {
+            out = new Linked<>(this.stream().limit(size).collect(Collectors.toList()));
         } else {
-            int skip = 0;
+            boolean found = false;
+            out = new Linked<>();
+            long cnt = size;
             for (final X elem : this) {
-                ++skip;
-                if (prd.test(elem)) {
-                    break;
+                if (!found && prd.negate().test(elem)) {
+                    continue;
+                } else if (!found) {
+                    found = true;
+                    continue;
+                }
+                if (cnt > 0) {
+                    --cnt;
+                    out.add(elem);
                 }
             }
-            out = new Linked<>(
-                this.stream().skip(skip).limit(size).collect(Collectors.toList())
-            );
         }
         return out;
     }

src/test/java/io/github/dgroup/enumerable4j/AfterTest.java

@@ -24,6 +24,8 @@
 
 package io.github.dgroup.enumerable4j;
 
+import org.hamcrest.collection.IsEmptyIterable;
+import org.hamcrest.core.AllOf;
 import org.junit.jupiter.api.Test;
 import org.llorllale.cactoos.matchers.Assertion;
 import org.llorllale.cactoos.matchers.HasSize;
@@ -44,43 +46,46 @@ void nullPredicate() {
         new Assertion<>(
             "In case null-function the self enumerable is expected",
             new Linked<>(1, 2, 3, 4).after(null),
-            new HasValues<>(1, 2, 3, 4)
+            new AllOf<>(
+                new HasSize(4),
+                new HasValues<>(1, 2, 3, 4)
+            )
         ).affirm();
     }
 
     @Test
     void allAfter() {
         new Assertion<>(
-            "All elements after the first one which corresponds the condition",
-            new Linked<>(1, 2, 3, 4, 5, 6, 7).after(n -> n > 2),
-            new HasValues<>(4, 5, 6, 7)
+            "Returns elements after the first one which corresponds the condition",
+            new Linked<>(1, 2, 3, 4, 5).after(n -> n > 1),
+            new AllOf<>(
+                new HasSize(3),
+                new HasValues<>(3, 4, 5)
+            )
         ).affirm();
     }
 
     @Test
-    void firstThreeAfter() {
+    void firstTwoAfter() {
         new Assertion<>(
-            "First 3 elements after the first one which corresponds the condition",
-            new Linked<>(1, 2, 3, 4, 5, 6, 7).after(n -> n > 2, 3),
-            new HasValues<>(4, 5, 6)
-        ).affirm();
-    }
-
-    @Test
-    void allStringsAfter() {
-        new Assertion<>(
-            "All string elements after the first one which corresponds the condition",
-            new Linked<>("a", "b", "c", "d", "e").after(s -> s.equals("c")),
-            new HasValues<>("d", "e")
+            "Returns first 2 elements after the first one which corresponds the condition",
+            new Linked<>("a", "b", "c", "d", "e").after(s -> s.equals("b"), 2),
+            new AllOf<>(
+                new HasSize(2),
+                new HasValues<>("c", "d")
+            )
         ).affirm();
     }
 
     @Test
     void firstTenAfter() {
         new Assertion<>(
             "The specified size can be greater than an actual enumerable size",
-            new Linked<>(1, 2, 3, 4, 5).after(n -> n == 3, 10),
-            new HasValues<>(4, 5)
+            new Linked<>(1, 2, 3, 4, 5).after(n -> n == 2, 10),
+        new AllOf<>(
+            new HasSize(3),
+            new HasValues<>(3, 4, 5)
+        )
         ).affirm();
     }
 
@@ -89,16 +94,49 @@ void allAfterFirst() {
         new Assertion<>(
             "The first element of enumerable corresponds the condition",
             new Linked<>(1, 2, 3, 4, 5).after(n -> n < 10),
+        new AllOf<>(
+            new HasSize(4),
             new HasValues<>(2, 3, 4, 5)
+        )
         ).affirm();
     }
 
     @Test
     void noneMatch() {
         new Assertion<>(
-            "No one elements corresponds the condition, an empty collection is returned",
+            "No one element corresponds the condition, an empty collection is returned",
             new Linked<>(1, 2, 3).after(n -> n < 0),
-            new HasSize(0)
+            new IsEmptyIterable<>()
+        ).affirm();
+    }
+
+    @Test
+    void zeroSize() {
+        new Assertion<>(
+            "If the size value is 0, an empty collection is returned",
+            new Linked<>(1, 2, 3).after(n -> n < 10, 0),
+            new IsEmptyIterable<>()
+        ).affirm();
+    }
+
+    @Test
+    void nullPredicateWithSize() {
+        new Assertion<>(
+            "In case null-function and specified size a enumerable of the size is expected",
+            new Linked<>(1, 2, 3, 4).after(null, 2),
+            new AllOf<>(
+                new HasSize(2),
+                new HasValues<>(1, 2)
+            )
+        ).affirm();
+    }
+
+    @Test
+    void nullPredicateAndZeroSize() {
+        new Assertion<>(
+            "If the size value is 0, an empty collection is returned",
+            new Linked<>(1, 2, 3).after(null, 0),
+            new IsEmptyIterable<>()
         ).affirm();
     }