1 files changed, 51 insertions, 16 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..5a1d9c8 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.
@@ -16,14 +15,13 @@ public class TreeUtils {
 	/**
 	 * Convert a tree into a list of outline nodes that match a certain path.
 	 *
-	 * @param tre
-	 *                   The tree to outline.
-	 * @param leafMarker
-	 *                   The path to mark nodes with.
+	 * @param <T> The type of the values.
+	 * @param tre The tree to outline.
+	 * @param leafMarker The path to mark nodes with.
 	 * @return The list of marked paths.
 	 */
-	public static <T> IList<IList<T>> outlineTree(ITree<T> tre, Predicate<T> leafMarker) {
-		IList<IList<T>> paths = new FunctionalList<>();
+	public static <T> ListEx<ListEx<T>> outlineTree(Tree<T> tre, Predicate<T> leafMarker) {
+		ListEx<ListEx<T>> paths = new FunctionalList<>();
 
 		LinkedList<T> path = new LinkedList<>();
 		path.add(tre.getHead());
@@ -34,15 +32,13 @@ public class TreeUtils {
 	}
 
 	/* Find a path in a tree. */
-	private static <T> void findPath(ITree<T> subtree, LinkedList<T> path,
-			Predicate<T> leafMarker, IList<IList<T>> paths) {
+	private static <T> void findPath(Tree<T> subtree, LinkedList<T> path,
+			Predicate<T> leafMarker, ListEx<ListEx<T>> paths) {
 		if (subtree.getChildrenCount() == 0 && leafMarker.test(subtree.getHead())) {
 			/* We're at a matching leaf node. Add it. */
-			IList<T> finalPath = new FunctionalList<>();
+			ListEx<T> finalPath = new FunctionalList<>();
 
-			for (T ePath : path) {
-				finalPath.add(ePath);
-			}
+			for (T ePath : path) finalPath.add(ePath);
 
 			finalPath.add(subtree.getHead());
 
@@ -56,4 +52,43 @@ 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> Tree<ContainedType> substitute(
+			Tree<ContainedType> tree,
+			Predicate<ContainedType> marker,
+			Function<ContainedType, Tree<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> Tree<ContainedType> substitute(
+			Tree<ContainedType> tree,
+			MapEx<ContainedType, Tree<ContainedType>> environment) {
+		return substitute(
+				tree,
+				environment::containsKey,
+				(element) -> environment.get(element).get());
+	}
 }