Initial commit: only works for plain int typed arguments

1 files changed, 72 insertions, 0 deletions

diff --git a/src/main/java/au/id/zancanaro/RoseTree.java b/src/main/java/au/id/zancanaro/RoseTree.java
new file mode 100644
index 0000000..2a83ad2
--- /dev/null
+++ b/src/main/java/au/id/zancanaro/RoseTree.java
@@ -0,0 +1,72 @@
+package au.id.zancanaro;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.function.Function;
+
+public class RoseTree<T> {
+    private final T value;
+    private final Iterable<RoseTree<T>> children;
+
+    public RoseTree(T value, Iterable<RoseTree<T>> children) {
+        this.value = value;
+        this.children = children;
+    }
+
+    public T getValue() {
+        return value;
+    }
+
+    public Iterator<RoseTree<T>> getChildren() {
+        return children.iterator();
+    }
+
+    public static <T> RoseTree<T> pure(T value) {
+        return new RoseTree<>(value, Collections.emptyList());
+    }
+
+    public static <T> RoseTree<T> join(RoseTree<RoseTree<T>> tree) {
+        return new RoseTree<>(
+                tree.getValue().getValue(),
+                () -> Iterators.concat(
+                        Iterators.mappingIterator(RoseTree::join, tree.children.iterator()),
+                        tree.getValue().children.iterator()));
+    }
+
+    public static <T> Iterator<RoseTree<T>[]> permutations(RoseTree<T>[] trees) {
+        return Iterators.flatten(Iterators.rangeIterator(trees.length, index ->
+                Iterators.mappingIterator(child -> {
+                    @SuppressWarnings("unchecked")
+                    RoseTree<T>[] result = (RoseTree<T>[]) new RoseTree[trees.length];
+                    for (int i = 0; i < trees.length; ++i) {
+                        result[i] = (i == index ? child : trees[i]);
+                    }
+                    return result;
+                }, trees[index].getChildren())
+        ));
+    }
+
+    public static <T, R> RoseTree<R> zip(Function<T[], R> fn, RoseTree<T>[] trees) {
+        @SuppressWarnings("unchecked")
+        T[] heads = (T[]) new Object[trees.length];
+        for (int i = 0; i < trees.length; ++i) {
+            heads[i] = trees[i].getValue();
+        }
+        return new RoseTree<>(
+                fn.apply(heads),
+                () -> Iterators.mappingIterator(
+                        roses -> RoseTree.zip(fn, roses),
+                        RoseTree.permutations(trees)));
+    }
+
+    public <R> RoseTree<R> fmap(Function<T, R> f) {
+        return new RoseTree<>(
+                f.apply(this.value),
+                () -> Iterators.mappingIterator(tree -> tree.fmap(f), this.children.iterator()));
+    }
+
+    public <R> RoseTree<R> flatmap(Function<T, RoseTree<R>> f) {
+        return RoseTree.join(this.fmap(f));
+    }
+}