package au.id.zancanaro.javacheck; import java.io.IOException; import java.io.Writer; import java.util.*; import java.util.function.Function; import java.util.function.Predicate; public class RoseTree { private final T value; private final Iterable> children; public RoseTree(T value, Iterable> children) { this.value = value; this.children = children; } public T getValue() { return value; } public Iterator> getChildren() { return children.iterator(); } public static RoseTree pure(T value) { return new RoseTree<>(value, Collections.emptyList()); } public static RoseTree join(RoseTree> tree) { return new RoseTree<>( tree.getValue().getValue(), () -> Iterators.concat( Iterators.mappingIterator(RoseTree::join, tree.children.iterator()), tree.getValue().children.iterator())); } private static Iterator[]> permutations(RoseTree[] trees) { return Iterators.flatten( Iterators.rangeIterator(trees.length, index -> Iterators.mappingIterator(child -> { @SuppressWarnings("unchecked") RoseTree[] result = (RoseTree[]) 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 RoseTree zip(Function, R> fn, RoseTree[] trees) { @SuppressWarnings("unchecked") List heads = new ArrayList(trees.length); for (RoseTree tree : trees) { heads.add(tree.getValue()); } return new RoseTree<>( fn.apply(heads), () -> Iterators.mappingIterator( roses -> RoseTree.zip(fn, roses), RoseTree.permutations(trees))); } private static Iterator[]> removeEach(RoseTree[] trees) { return Iterators.concat( Iterators.rangeIterator(trees.length, index -> { @SuppressWarnings("unchecked") RoseTree[] result = (RoseTree[]) new RoseTree[trees.length - 1]; for (int i = 0; i < trees.length - 1; ++i) { result[i] = trees[(i >= index ? i + 1 : i)]; } return result; }), permutations(trees)); } public static RoseTree shrink(Function, R> fn, RoseTree[] trees) { @SuppressWarnings("unchecked") List heads = new ArrayList(trees.length); for (RoseTree tree : trees) { heads.add(tree.getValue()); } return new RoseTree<>( fn.apply(heads), () -> Iterators.mappingIterator( roses -> RoseTree.shrink(fn, roses), RoseTree.removeEach(trees))); } public RoseTree map(Function f) { return new RoseTree<>( f.apply(this.value), () -> Iterators.mappingIterator(tree -> tree.map(f), this.children.iterator())); } public RoseTree flatMap(Function> f) { return RoseTree.join(this.map(f)); } public RoseTree filter(Predicate predicate) { if (predicate.test(this.getValue())) { return new RoseTree<>( this.getValue(), () -> Iterators.mappingIterator(tree -> tree.filter(predicate), Iterators.filteringIterator( tree -> predicate.test(tree.getValue()), this.getChildren()))); } else { throw new IllegalArgumentException("Current value doesn't match predicate: whoops!"); } } @SuppressWarnings("unused") public void print(Writer output) throws IOException { print(output, Object::toString); } @SuppressWarnings("unused") public void print(Writer output, Function toString) throws IOException { output.write(toString.apply(this.getValue())); output.write('['); for (RoseTree child : children) { child.print(output, toString); } output.write(']'); output.flush(); } }