Add new top-down transform type

3 files changed, 12 insertions, 1 deletions

diff --git a/BJC-Utils2/src/main/java/bjc/utils/data/TopDownTransformIterator.java b/BJC-Utils2/src/main/java/bjc/utils/data/TopDownTransformIterator.java
index 3886264..4c5ec8a 100644
--- a/BJC-Utils2/src/main/java/bjc/utils/data/TopDownTransformIterator.java
+++ b/BJC-Utils2/src/main/java/bjc/utils/data/TopDownTransformIterator.java
@@ -8,6 +8,8 @@ import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
+import static bjc.utils.data.TopDownTransformResult.*;
+
 public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<ContainedType>> {
 	private Function<ContainedType, TopDownTransformResult> picker;
 	private BiFunction<ITree<ContainedType>, 
@@ -83,6 +85,9 @@ public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<C
 				case TRANSFORM:
 					done = true;
 					return transform.apply(preParent, this::addYield);
+				case RTRANSFORM:
+					preParent = transform.apply(preParent, this::addYield);
+					break;
 				case PUSHDOWN:
 					if(preParent.getChildrenCount() != 0) {
 						for(int i = 0; i < preParent.getChildrenCount(); i++) {
@@ -115,7 +120,7 @@ public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<C
 					throw new IllegalArgumentException("Unknown result type " + res);
 			}
 
-			initial = false;
+			if(res != RTRANSFORM) initial = false;
 		}
 
 		if(curYield != null) {
diff --git a/BJC-Utils2/src/main/java/bjc/utils/data/TopDownTransformResult.java b/BJC-Utils2/src/main/java/bjc/utils/data/TopDownTransformResult.java
index e72cdf6..0db9482 100644
--- a/BJC-Utils2/src/main/java/bjc/utils/data/TopDownTransformResult.java
+++ b/BJC-Utils2/src/main/java/bjc/utils/data/TopDownTransformResult.java
@@ -16,6 +16,10 @@ public enum TopDownTransformResult {
 	 */
 	TRANSFORM,
 	/**
+	 * Transform this node, then do a top-down transform on the result
+	 */
+	RTRANSFORM,
+	/**
 	 * Ignore this node, and traverse its children
 	 */
 	PASSTHROUGH,
diff --git a/BJC-Utils2/src/main/java/bjc/utils/data/Tree.java b/BJC-Utils2/src/main/java/bjc/utils/data/Tree.java
index 90e6793..757ffae 100644
--- a/BJC-Utils2/src/main/java/bjc/utils/data/Tree.java
+++ b/BJC-Utils2/src/main/java/bjc/utils/data/Tree.java
@@ -204,6 +204,8 @@ public class Tree<ContainedType> implements ITree<ContainedType> {
 				return this;
 			case TRANSFORM:
 				return transformer.apply(this);
+			case RTRANSFORM:
+				return transformer.apply(this).topDownTransform(transformPicker, transformer);
 			case PUSHDOWN:
 				result = new Tree<>(data);