Added new tree traversal option, and some minor changes

3 files changed, 68 insertions, 2 deletions

diff --git a/BJC-Utils2/src/main/java/bjc/utils/funcdata/ITree.java b/BJC-Utils2/src/main/java/bjc/utils/funcdata/ITree.java
index 624c9d6..7f28682 100644
--- a/BJC-Utils2/src/main/java/bjc/utils/funcdata/ITree.java
+++ b/BJC-Utils2/src/main/java/bjc/utils/funcdata/ITree.java
@@ -114,8 +114,8 @@ public interface ITree<ContainedType> {
 	 *            The function to use to transform tokens
 	 * @return A tree with the token types transformed
 	 */
-	public <MappedType> ITree<MappedType> transformTree(
-			Function<ContainedType, MappedType> transformer);
+	public <MappedType> ITree<MappedType>
+			transformTree(Function<ContainedType, MappedType> transformer);
 
 	/**
 	 * Perform an action on each part of the tree
@@ -142,4 +142,17 @@ public interface ITree<ContainedType> {
 	public <MappedType> ITree<MappedType> rebuildTree(
 			Function<ContainedType, MappedType> leafTransformer,
 			Function<ContainedType, MappedType> operatorTransformer);
+
+	/**
+	 * Do a top-down transform of the tree
+	 * 
+	 * @param transformPicker
+	 *            The function to use to pick how to progress
+	 * @param transformer
+	 *            The function used to transform picked subtrees
+	 * @return The tree with the transform applied to picked subtrees
+	 */
+	public ITree<ContainedType> topDownTransform(
+			Function<ContainedType, TopDownTransformResult> transformPicker,
+			UnaryOperator<ITree<ContainedType>> transformer);
 }
diff --git a/BJC-Utils2/src/main/java/bjc/utils/funcdata/TopDownTransformResult.java b/BJC-Utils2/src/main/java/bjc/utils/funcdata/TopDownTransformResult.java
new file mode 100644
index 0000000..9aec2b9
--- /dev/null
+++ b/BJC-Utils2/src/main/java/bjc/utils/funcdata/TopDownTransformResult.java
@@ -0,0 +1,22 @@
+package bjc.utils.funcdata;
+
+/**
+ * Represents the results for doing a top-down transform of a tree
+ * 
+ * @author ben
+ *
+ */
+public enum TopDownTransformResult {
+	/**
+	 * Do not do anything to this node, and ignore it's children
+	 */
+	SKIP,
+	/**
+	 * Transform this node, and don't touch its children
+	 */
+	TRANSFORM,
+	/**
+	 * Ignore this node, and traverse its children
+	 */
+	PASSTHROUGH;
+}
diff --git a/BJC-Utils2/src/main/java/bjc/utils/funcdata/Tree.java b/BJC-Utils2/src/main/java/bjc/utils/funcdata/Tree.java
index 5ee4200..03a1f93 100644
--- a/BJC-Utils2/src/main/java/bjc/utils/funcdata/Tree.java
+++ b/BJC-Utils2/src/main/java/bjc/utils/funcdata/Tree.java
@@ -262,4 +262,35 @@ public class Tree<ContainedType> implements ITree<ContainedType> {
 			});
 		}
 	}
+
+	@Override
+	public ITree<ContainedType> topDownTransform(
+			Function<ContainedType, TopDownTransformResult> transformPicker,
+			UnaryOperator<ITree<ContainedType>> transformer) {
+		TopDownTransformResult transformResult =
+				transformPicker.apply(data);
+
+		switch (transformResult) {
+			case PASSTHROUGH:
+				ITree<ContainedType> result = new Tree<>(data);
+
+				if (hasChildren) {
+					children.forEach((child) -> {
+						result.addChild(child.topDownTransform(
+								transformPicker, transformer));
+					});
+				}
+
+				return result;
+			case SKIP:
+				return this;
+			case TRANSFORM:
+				return transformer.apply(this);
+			default:
+				throw new IllegalArgumentException(
+						"Recieved unknown transform result "
+								+ transformResult);
+
+		}
+	}
 }