Add upwards and direct-children-only traversal

2 files changed, 65 insertions, 1 deletions

diff --git a/javaparser-core/src/main/java/com/github/javaparser/ast/Node.java b/javaparser-core/src/main/java/com/github/javaparser/ast/Node.java
index 4168fccb8..9996cccf9 100644
--- a/javaparser-core/src/main/java/com/github/javaparser/ast/Node.java
+++ b/javaparser-core/src/main/java/com/github/javaparser/ast/Node.java
@@ -686,7 +686,7 @@ public abstract class Node implements Cloneable, HasParentNode<Node>, Visitable,
 
 
     public enum TreeTraversal {
-        PREORDER, BREADTHFIRST, POSTORDER
+        PREORDER, BREADTHFIRST, POSTORDER, PARENTS, DIRECT_CHILDREN
     }
 
     public Iterator<Node> treeIterator(TreeTraversal traversal) {
@@ -697,6 +697,10 @@ public abstract class Node implements Cloneable, HasParentNode<Node>, Visitable,
                 return new TreeVisitor.PostOrderIterator(this);
             case PREORDER:
                 return new TreeVisitor.PreOrderIterator(this);
+            case DIRECT_CHILDREN:
+                return new TreeVisitor.DirectChildrenIterator(this);
+            case PARENTS:
+                return new TreeVisitor.ParentsVisitor(this);
             default:
                 throw new IllegalArgumentException("Unknown traversal choice.");
         }
diff --git a/javaparser-core/src/main/java/com/github/javaparser/ast/visitor/TreeVisitor.java b/javaparser-core/src/main/java/com/github/javaparser/ast/visitor/TreeVisitor.java
index 9d1222126..89d4ce843 100644
--- a/javaparser-core/src/main/java/com/github/javaparser/ast/visitor/TreeVisitor.java
+++ b/javaparser-core/src/main/java/com/github/javaparser/ast/visitor/TreeVisitor.java
@@ -65,6 +65,47 @@ public abstract class TreeVisitor {
         }
     }
 
+    public static class DirectChildrenIterator implements Iterator<Node> {
+        private final Iterator<Node> childrenIterator;
+
+        public DirectChildrenIterator(Node node) {
+            childrenIterator = new ArrayList<>(node.getChildNodes()).iterator();
+        }
+
+        @Override
+        public boolean hasNext() {
+            return childrenIterator.hasNext();
+        }
+
+        @Override
+        public Node next() {
+            return childrenIterator.next();
+        }
+    }
+
+    /**
+     * Iterates over the parent of the node, then the parent's parent, then the parent's parent's parent, until running
+     * out of parents.
+     */
+    public static class ParentsVisitor implements Iterator<Node> {
+        private Node node;
+
+        public ParentsVisitor(Node node) {
+            this.node = node;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return node.getParentNode().isPresent();
+        }
+
+        @Override
+        public Node next() {
+            node = node.getParentNode().orElse(null);
+            return node;
+        }
+    }
+
     public static class PreOrderIterator implements Iterator<Node> {
         private final Stack<Node> stack = new Stack<>();
 
@@ -178,6 +219,25 @@ public abstract class TreeVisitor {
     }
 
     /**
+     * Calls "process" for each direct child node of "node"
+     *
+     * @param node the parent node.
+     */
+    public void visitDirectChildren(Node node) {
+        new DirectChildrenIterator(node).forEachRemaining(this::process);
+    }
+
+    /**
+     * Calls "process" for the parent of "node", then the parent's parent, then the parent's parent's parent, until
+     * running out of parents. <br/>Note that "node" itself is not processed.
+     *
+     * @param node the start node.
+     */
+    public void visitParents(Node node) {
+        new ParentsVisitor(node).forEachRemaining(this::process);
+    }
+
+    /**
      * Process the given node.
      *
      * @param node The current node to process.