Simplify Tree.collapse

2 files changed, 7 insertions, 6 deletions

diff --git a/base/src/main/java/bjc/utils/data/ITree.java b/base/src/main/java/bjc/utils/data/ITree.java
index 5a4d645..3e3d5ce 100644
--- a/base/src/main/java/bjc/utils/data/ITree.java
+++ b/base/src/main/java/bjc/utils/data/ITree.java
@@ -1,10 +1,12 @@
 package bjc.utils.data;
 
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.UnaryOperator;
 
+import bjc.utils.funcdata.IList;
 import bjc.utils.funcdata.bst.TreeLinearizationMethod;
 import bjc.utils.functypes.ListFlattener;
 
@@ -57,7 +59,7 @@ public interface ITree<ContainedType> {
 	 * @return The final transformed state.
 	 */
 	<NewType, ReturnedType> ReturnedType collapse(Function<ContainedType, NewType> leafTransform,
-			Function<ContainedType, ListFlattener<NewType>> nodeCollapser,
+			BiFunction<ContainedType, IList<NewType>, NewType> nodeCollapser,
 			Function<NewType, ReturnedType> resultTransformer);
 
 	/**
diff --git a/base/src/main/java/bjc/utils/data/Tree.java b/base/src/main/java/bjc/utils/data/Tree.java
index f7ff90a..8fcbead 100644
--- a/base/src/main/java/bjc/utils/data/Tree.java
+++ b/base/src/main/java/bjc/utils/data/Tree.java
@@ -1,5 +1,6 @@
 package bjc.utils.data;
 
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
@@ -187,7 +188,7 @@ public class Tree<ContainedType> implements ITree<ContainedType> {
 
 	@Override
 	public <NewType, ReturnedType> ReturnedType collapse(final Function<ContainedType, NewType> leafTransform,
-			final Function<ContainedType, ListFlattener<NewType>> nodeCollapser,
+			final BiFunction<ContainedType, IList<NewType>, NewType> nodeCollapser,
 			final Function<NewType, ReturnedType> resultTransformer) {
 		return resultTransformer.apply(internalCollapse(leafTransform, nodeCollapser));
 	}
@@ -216,10 +217,8 @@ public class Tree<ContainedType> implements ITree<ContainedType> {
 	 */
 
 	protected <NewType> NewType internalCollapse(final Function<ContainedType, NewType> leafTransform,
-			final Function<ContainedType, ListFlattener<NewType>> nodeCollapser) {
+			final BiFunction<ContainedType, IList<NewType>, NewType> nodeCollapser) {
 		if(hasChildren) {
-			final Function<IList<NewType>, NewType> nodeTransformer = nodeCollapser.apply(data);
-
 			final IList<NewType> collapsedChildren = children.map(child -> {
 				final NewType collapsed = child.collapse(leafTransform, nodeCollapser,
 						subTreeVal -> subTreeVal);
@@ -227,7 +226,7 @@ public class Tree<ContainedType> implements ITree<ContainedType> {
 				return collapsed;
 			});
 
-			return nodeTransformer.apply(collapsedChildren);
+			return nodeCollapser.apply(data, collapsedChildren);
 		}
 
 		return leafTransform.apply(data);