Augmented iterators for a post on my blog

AugmentedIterable.java

public interface AugmentedIterable<T> extends Iterable<T> {

  default <R> R foldLeft(R init, BiFunction<? super R, ? super T, ? extends R> f) {
    R out = init;
    for (T elem : this) {
      out = f.apply(out, elem);
    }
    return out;
  }

  default <R> AugmentedIterable<R> map(Function<? super T, ? extends R> f) {
    final Iterable<T> origIterable = this;
    return () -> new Iterator<R>() {
      final Iterator<T> origIter = origIterable.iterator();

      @Override
      public boolean hasNext() {
        return origIter.hasNext();
      }

      @Override
      public R next() {
        return f.apply(origIter.next());
      }
    };
  }

  default AugmentedIterable<T> filter(Predicate<? super T> p) {
    final Iterable<T> origIterable = this;
    return () -> new Iterator<T>() {
      final Iterator<T> origIter = origIterable.iterator();
      T next;
      boolean hasNext = false;

      @Override
      public boolean hasNext() {
        if (hasNext) return true;
        while (origIter.hasNext()) {
          next = origIter.next();
          if (p.test(next)) {
            hasNext = true;
            return true;
          }
        }
        return false;
      }

      @Override
      public T next() {
        if (hasNext) {
          hasNext = false;
          return next;
        }
        throw new NoSuchElementException();
      }
    };
  }

  default <R> AugmentedIterable<R>
  flatMap(Function<? super T, ? extends Iterable<? extends R>> f) {
    final Iterable<T> origIterable = this;
    return () -> new Iterator<R>() {
      final Iterator<T> origIter = origIterable.iterator();
      Iterator<? extends R> nestedIterator = null;

      @Override
      public boolean hasNext() {
        while (nestedIterator == null || !nestedIterator.hasNext()) {
          if (origIter.hasNext()) {
            nestedIterator = f.apply(origIter.next()).iterator();
          }
        }
        return (nestedIterator != null && nestedIterator.hasNext());
      }

      @Override
      public R next() {
        return nestedIterator.next();
      }
    };
  }

  default AugmentedIterable<T> take(int n) {
    final Iterable<T> origIterable = this;
    return () -> new Iterator<T>() {
      final Iterator<T> origIter = origIterable.iterator();
      private int pos = 0;

      @Override
      public boolean hasNext() {
        return pos < n && origIter.hasNext();
      }

      @Override
      public T next() {
        pos++;
        return origIter.next();
      }
    };
  }

  default AugmentedIterable<T> drop(int n) {
    final Iterable<T> origIterable = this;
    return () -> {
      final Iterator<T> origIter = origIterable.iterator();
      for (int i = 0; i < n; i++) {
        if (origIter.hasNext()) {
          origIter.next();
        } else {
          break;
        }
      }
      return origIter;
    };
  }

  // I'm using this for testing, but it's also generally useful
  default List<T> toList() {
    final List<T> out = new ArrayList<>();
    forEach(out::add);
    return out;
  }
}

AugmentedIterableTest.java

public class AugmentedIterableTest {
  @Test
  public void testFoldLeft() throws Exception {
    assertEquals(10,
      augment(Arrays.asList(1, 2, 3, 4)).foldLeft(0, (a, b) -> a + b).intValue());
  }

  @Test
  public void testMap() throws Exception {
    assertEquals(Arrays.asList(2, 4, 6, 8),
      augment(Arrays.asList(1, 2, 3, 4)).map((a) -> a * 2).toList());
  }

  @Test
  public void testFilter() throws Exception {
    assertEquals(Arrays.asList(2, 4, 6, 8),
      augment(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8))
        .filter((a) -> a % 2 == 0).toList());
  }

  @Test
  public void testTake() throws Exception {
    assertEquals(Arrays.asList(1, 2, 3, 4),
      augment(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8)).take(4).toList());
  }

  @Test
  public void testDrop() throws Exception {
    assertEquals(Arrays.asList(5, 6, 7, 8),
      augment(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8)).drop(4).toList());
  }
  @Test
  public void testLazyMap() throws Exception {
    // Confirm that evaluation is lazy
    final AtomicInteger counter = new AtomicInteger(0);
    Function<Integer, Integer> doubleWithSideEffect = (a) -> {
      counter.incrementAndGet();
      return a * 2;
    };
    Predicate<Integer> isOddWithSideEffect = (a) -> {
      counter.incrementAndGet();
      return a % 2 == 1;
    };
    AugmentedIterable<Integer> result =
      augment(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8))
        .filter(isOddWithSideEffect)
        .map(doubleWithSideEffect);
    assertEquals(0, counter.get());
    assertEquals(Arrays.asList(2, 6, 10, 14), result.toList());
    // 8 calls to the predicate and 4 calls to mapping function
    assertEquals(8 + 4, counter.get());
  }

  // A simple take on the rather dubious straw man argument given at
  // http://cafe.elharo.com/programming/java-programming/why-functional-programming-in-java-is-dangerous/
  @Test
  public void testSquaresOfIntegers() throws Exception {
    // An infinite iterable that generates positive integers
    AugmentedIterable<Integer> integers = () -> new Iterator<Integer>() {
      int i = 1;
      @Override
      public boolean hasNext() {
        return true;
      }

      @Override
      public Integer next() {
        return i++;
      }
    };
    assertEquals(Arrays.asList(1, 4, 9, 16, 25, 36, 49, 64),
      integers.map((a) -> a * a).take(8).toList());
  }
}

AugmentedIterableUtils.java

public final class AugmentedIterableUtils {
  private AugmentedIterableUtils() {
  }

  public static <T> AugmentedIterable<T> augment(Iterable<T> iterable) {
    return iterable::iterator;
  }
}

AugmentedLists.java

public class AugmentedLinkedList<T> extends LinkedList<T>
    implements AugmentedIterable<T> {
  // Expose constructors from superclass
  public AugmentedLinkedList() {
  }

  public AugmentedLinkedList(final Collection<? extends T> c) {
    super(c);
  }
}

public class AugmentedArrayList<T> extends ArrayList<T>
    implements AugmentedIterable<T> {
  // Expose constructors from superclass
  public AugmentedArrayList(final int initialCapacity) {
    super(initialCapacity);
  }
  
  public AugmentedArrayList() {
  }

  public AugmentedArrayList(final Collection<? extends T> c) {
    super(c);
  }
}