Tweak tree utils

1 files changed, 39 insertions, 4 deletions

diff --git a/base/src/main/java/bjc/utils/funcutils/TreeUtils.java b/base/src/main/java/bjc/utils/funcutils/TreeUtils.java
index d525773..7f54f73 100644
--- a/base/src/main/java/bjc/utils/funcutils/TreeUtils.java
+++ b/base/src/main/java/bjc/utils/funcutils/TreeUtils.java
@@ -1,11 +1,10 @@
 package bjc.utils.funcutils;
 
 import java.util.LinkedList;
-import java.util.function.Predicate;
+import java.util.function.*;
 
-import bjc.data.ITree;
-import bjc.funcdata.FunctionalList;
-import bjc.funcdata.IList;
+import bjc.data.*;
+import bjc.funcdata.*;
 
 /**
  * Implements various utilities for trees.
@@ -56,4 +55,40 @@ public class TreeUtils {
 			path.removeLast();
 		}
 	}
+	
+	/**
+	 * Performs 'variable substitution' or something along those lines on a tree.
+	 * 
+	 * @param <ContainedType> The type of element contained in the tree.
+	 * @param tree The tree to do expansion in.
+	 * @param marker The function to mark which nodes should be expanded.
+	 * @param expander The function to expand nodes.
+	 * @return A transformed copy of the tree.
+	 */
+	public static <ContainedType> ITree<ContainedType> substitute(
+			ITree<ContainedType> tree,
+			Predicate<ContainedType> marker,
+			Function<ContainedType, ITree<ContainedType>> expander) {
+		tree.topDownTransform((contents) -> {
+			if (marker.test(contents)) return TopDownTransformResult.TRANSFORM;
+			else                       return TopDownTransformResult.PASSTHROUGH;
+		}, (node) -> {
+			return expander.apply(node.getHead());
+		});
+		return tree;
+	}
+	
+	/**
+	 * Performs 'variable substitution' or something along those lines on a tree.
+	 * 
+	 * @param <ContainedType> The type of element contained in the tree.
+	 * @param tree The tree to do expansion in.
+	 * @param environment A map which contains the variables to substitute.
+	 * @return A transformed copy of the tree.
+	 */
+	public static <ContainedType> ITree<ContainedType> substitute(
+			ITree<ContainedType> tree,
+			IMap<ContainedType, ITree<ContainedType>> environment) {
+		return substitute(tree, environment::containsKey, environment::get);
+	}
 }