summaryrefslogtreecommitdiff
path: root/src/main/java/bjc/data/FlatMapIterator.java
diff options
context:
space:
mode:
Diffstat (limited to 'src/main/java/bjc/data/FlatMapIterator.java')
-rw-r--r--src/main/java/bjc/data/FlatMapIterator.java52
1 files changed, 52 insertions, 0 deletions
diff --git a/src/main/java/bjc/data/FlatMapIterator.java b/src/main/java/bjc/data/FlatMapIterator.java
new file mode 100644
index 0000000..47bc06c
--- /dev/null
+++ b/src/main/java/bjc/data/FlatMapIterator.java
@@ -0,0 +1,52 @@
+package bjc.data;
+
+import java.util.*;
+import java.util.function.Function;
+
+/**
+ * Create an iterator that applies a flat-map function
+ *
+ * @author bjcul
+ *
+ * @param <T1> The source type
+ * @param <T2> The output type
+ *
+ */
+public class FlatMapIterator<T1, T2> implements Iterator<T2> {
+ private Function<T1, Iterator<T2>> action;
+
+ private Iterator<T1> source;
+ private Iterator<T2> output;
+ /**
+ * Create a new flat-map iterator
+ *
+ * @param source The source iterator
+ * @param action The action to apply to the iterator
+ */
+ public FlatMapIterator(Iterator<T1> source, Function<T1, Iterator<T2>> action) {
+ this.source = source;
+ this.action = action;
+
+ }
+
+ @Override
+ public boolean hasNext() {
+ if (output != null && output.hasNext()) return true;
+ while (source.hasNext()) {
+ output = action.apply(source.next());
+ if (output.hasNext()) return true;
+ }
+ return false;
+ }
+
+ @Override
+ public T2 next() {
+ if (output != null && output.hasNext()) return output.next();
+ while (source.hasNext()) {
+ output = action.apply(source.next());
+ if (output.hasNext()) return output.next();
+ }
+ throw new NoSuchElementException();
+ }
+
+}