1 files changed, 87 insertions, 54 deletions

diff --git a/base/src/main/java/bjc/utils/data/TopDownTransformIterator.java b/base/src/main/java/bjc/utils/data/TopDownTransformIterator.java
index d4a676c..1e69071 100644
--- a/base/src/main/java/bjc/utils/data/TopDownTransformIterator.java
+++ b/base/src/main/java/bjc/utils/data/TopDownTransformIterator.java
@@ -16,6 +16,14 @@ import java.util.function.Function;
  * 	Figure out what is broken with this, and fix it so that step-wise
  * 	iteration works correctly.
  */
+/**
+ * An iterative top-down transform of a tree.
+ * 
+ * @author EVE
+ *
+ * @param <ContainedType>
+ *        The type of the nodes in the tree.
+ */
 public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<ContainedType>> {
 	private final Function<ContainedType, TopDownTransformResult>							picker;
 	private final BiFunction<ITree<ContainedType>, Consumer<Iterator<ITree<ContainedType>>>, ITree<ContainedType>>	transform;
@@ -34,6 +42,16 @@ public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<C
 	private final Deque<Iterator<ITree<ContainedType>>>	toYield;
 	private Iterator<ITree<ContainedType>>			curYield;
 
+	/**
+	 * Create a new tree iterator.
+	 * 
+	 * @param pickr
+	 *        The function to use to pick how to process nodes.
+	 * @param transfrm
+	 *        The transform to apply to the nodes.
+	 * @param tree
+	 *        The tree to transform.
+	 */
 	public TopDownTransformIterator(final Function<ContainedType, TopDownTransformResult> pickr,
 			final BiFunction<ITree<ContainedType>, Consumer<Iterator<ITree<ContainedType>>>, ITree<ContainedType>> transfrm,
 			final ITree<ContainedType> tree) {
@@ -50,6 +68,12 @@ public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<C
 		initial = true;
 	}
 
+	/**
+	 * Add a set of nodes to yield.
+	 * 
+	 * @param src
+	 *        The nodes to yield.
+	 */
 	public void addYield(final Iterator<ITree<ContainedType>> src) {
 		if(curYield != null) {
 			toYield.push(curYield);
@@ -63,25 +87,34 @@ public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<C
 		return !done;
 	}
 
+	/**
+	 * Get the next yielded value.
+	 * 
+	 * @param val
+	 *        The sentinel value to yield.
+	 * @return The next yielded value.
+	 */
 	public ITree<ContainedType> flushYields(final ITree<ContainedType> val) {
 		if(curYield != null) {
 			toYield.add(new SingleIterator<>(val));
 
-			if(curYield.hasNext())
+			if(curYield.hasNext()) {
 				return curYield.next();
-			else {
-				while(toYield.size() != 0 && !curYield.hasNext()) {
-					curYield = toYield.pop();
-				}
+			}
 
-				if(toYield.size() == 0 && !curYield.hasNext()) {
-					curYield = null;
-					return val;
-				} else
-					return curYield.next();
+			while(toYield.size() != 0 && !curYield.hasNext()) {
+				curYield = toYield.pop();
 			}
-		} else
-			return val;
+
+			if(toYield.size() == 0 && !curYield.hasNext()) {
+				curYield = null;
+				return val;
+			}
+
+			return curYield.next();
+		}
+
+		return val;
 	}
 
 	@Override
@@ -89,17 +122,18 @@ public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<C
 		if(done) throw new NoSuchElementException();
 
 		if(curYield != null) {
-			if(curYield.hasNext())
+			if(curYield.hasNext()) {
 				return curYield.next();
-			else {
-				while(toYield.size() != 0 && !curYield.hasNext()) {
-					curYield = toYield.pop();
-				}
+			}
+
+			while(toYield.size() != 0 && !curYield.hasNext()) {
+				curYield = toYield.pop();
+			}
 
-				if(toYield.size() == 0 && !curYield.hasNext()) {
-					curYield = null;
-				} else
-					return curYield.next();
+			if(toYield.size() == 0 && !curYield.hasNext()) {
+				curYield = null;
+			} else {
+				return curYield.next();
 			}
 		}
 
@@ -117,10 +151,10 @@ public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<C
 
 					// Return whatever the first child is
 					break;
-				} else {
-					done = true;
-					return flushYields(postParent);
 				}
+
+				done = true;
+				return flushYields(postParent);
 			case SKIP:
 				done = true;
 				return flushYields(preParent);
@@ -138,11 +172,10 @@ public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<C
 
 					// Return whatever the first child is
 					break;
-				} else {
-					done = true;
-					return flushYields(transform.apply(new Tree<>(preParent.getHead()),
-							this::addYield));
 				}
+
+				done = true;
+				return flushYields(transform.apply(new Tree<>(preParent.getHead()), this::addYield));
 			case PULLUP:
 				final ITree<ContainedType> intRes = transform.apply(preParent, this::addYield);
 
@@ -155,10 +188,10 @@ public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<C
 
 					// Return whatever the first child is
 					break;
-				} else {
-					done = true;
-					return flushYields(postParent);
 				}
+
+				done = true;
+				return flushYields(postParent);
 			default:
 				throw new IllegalArgumentException("Unknown result type " + res);
 			}
@@ -177,38 +210,38 @@ public class TopDownTransformIterator<ContainedType> implements Iterator<ITree<C
 				postChildren.add(res);
 
 				return flushYields(res);
-			} else {
-				ITree<ContainedType> res = null;
-
-				if(postParent == null) {
-					res = new Tree<>(preParent.getHead());
+			}
 
-					System.out.println("\t\tTRACE: adding nodes " + postChildren + " to " + res);
+			ITree<ContainedType> res = null;
 
-					for(final ITree<ContainedType> child : postChildren) {
-						res.addChild(child);
-					}
+			if(postParent == null) {
+				res = new Tree<>(preParent.getHead());
 
-					// res = transform.apply(res,
-					// this::addYield);
-				} else {
-					res = postParent;
+				System.out.println("\t\tTRACE: adding nodes " + postChildren + " to " + res);
 
-					System.out.println("\t\tTRACE: adding nodes " + postChildren + " to " + res);
-					for(final ITree<ContainedType> child : postChildren) {
-						res.addChild(child);
-					}
+				for(final ITree<ContainedType> child : postChildren) {
+					res.addChild(child);
 				}
 
-				done = true;
-				return flushYields(res);
+				// res = transform.apply(res,
+				// this::addYield);
+			} else {
+				res = postParent;
+
+				System.out.println("\t\tTRACE: adding nodes " + postChildren + " to " + res);
+				for(final ITree<ContainedType> child : postChildren) {
+					res.addChild(child);
+				}
 			}
-		} else {
-			final ITree<ContainedType> res = curChild.next();
-			System.out.println("\t\tTRACE: adding node " + res + " to children");
-			postChildren.add(res);
 
+			done = true;
 			return flushYields(res);
 		}
+
+		final ITree<ContainedType> res = curChild.next();
+		System.out.println("\t\tTRACE: adding node " + res + " to children");
+		postChildren.add(res);
+
+		return flushYields(res);
 	}
 }