package com.github.nikita_volkov.java.iterators;

import java.util.Iterator;
import java.util.function.Function;

public final class FlatmappingIterator<a, b> implements Iterator<b> {

  private final Iterator<a> initialIterator;
  private final Function<a, Iterator<b>> projection;

  private Iterator<b> nextIterator;
  private b next;

  public FlatmappingIterator(Iterator<a> initialIterator, Function<a, Iterator<b>> projection) {
    this.initialIterator = initialIterator;
    this.projection = projection;
    preiterate();
  }

  private void preiterate() {
    if (nextIterator != null && nextIterator.hasNext()) {
      next = nextIterator.next();
    } else if (initialIterator.hasNext()) {
      nextIterator = projection.apply(initialIterator.next());
      preiterate();
    } else {
      next = null;
    }
  }

  @Override
  public boolean hasNext() {
    return next != null;
  }

  @Override
  public b next() {
    b result = next;
    preiterate();
    return result;
  }

}

@@ -564,19 +564,23 @@
 ) %7B%0A    
-if 
+while (
 (nextIte
@@ -589,62 +589,121 @@
 tor 
-!
+=
 = null 
-&& nextIterator.hasNext()) %7B%0A      next = next
+%7C%7C !nextIterator.hasNext()) && initialIterator.hasNext()) %7B%0A      nextIterator = projection.apply(initial
 Iter
@@ -717,32 +717,26 @@
 xt()
+)
 ;%0A    
-%7D else if (initial
+  if (!next
 Iter
@@ -751,32 +751,34 @@
 Next()) %7B%0A      
+  
 nextIterator = p
@@ -780,32 +780,64 @@
 r = 
-projection.apply(initial
+null;%0A      %7D%0A    %7D%0A    if (nextIterator != null && next
 Iter
@@ -845,33 +845,51 @@
 tor.
-n
+hasN
 ext())
-;
+ %7B
 %0A      
-preiterate
+next = nextIterator.next
 ();%0A

Update FlatmappingIterator