feature: add CtElement#descendantIterator and CtElement#asIterable (#…

…1986)

src/main/java/spoon/reflect/declaration/CtElement.java

@@ -32,6 +32,7 @@
 
 import java.lang.annotation.Annotation;
 import java.util.Collection;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Set;
 
@@ -368,4 +369,16 @@ <E extends CtElement> List<E> getAnnotatedChildren(
 	 */
 	CtPath getPath();
 
+	/**
+	 * Returns an iterator over this CtElement's descendants.
+	 * @return An iterator over this CtElement's descendants.
+	 */
+	Iterator<CtElement> descendantIterator();
+
+	/**
+	 * Returns an Iterable instance of this CtElement, allowing for dfs traversal of its descendants.
+	 * @return an Iterable object that allows iterating through this CtElement's descendants.
+	 */
+	Iterable<CtElement> asIterable();
+
 }

src/main/java/spoon/reflect/visitor/CtIterator.java

@@ -9,7 +9,7 @@
 /**
  * A class to be able to iterate over the children elements in the tree of a given node, in depth-first order.
  */
-public class CtIterator extends CtScanner implements Iterator {
+public class CtIterator extends CtScanner implements Iterator<CtElement> {
 	/**
 	 * A deque containing the elements the iterator has seen but not expanded
 	 */
@@ -67,7 +67,7 @@ public boolean hasNext() {
 	 * @return CtElement the next element in DFS order without going down the tree
 	 */
 	@Override
-	public Object next() {
+	public CtElement next() {
 		CtElement next = deque.pollFirst(); // get the element to expand from the deque
 		current_children.clear(); // clear for this scan
 		next.accept(this); // call @scan for each direct child of the node

src/main/java/spoon/support/reflect/declaration/CtElementImpl.java

@@ -52,6 +52,7 @@
 import spoon.reflect.visitor.chain.CtFunction;
 import spoon.reflect.visitor.chain.CtQuery;
 import spoon.reflect.visitor.filter.AnnotationFilter;
+import spoon.reflect.visitor.CtIterator;
 import spoon.support.DefaultCoreFactory;
 import spoon.support.DerivedProperty;
 import spoon.support.StandardEnvironment;
@@ -72,6 +73,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.Iterator;
 
 import static spoon.reflect.ModelElementContainerDefaultCapacities.ANNOTATIONS_CONTAINER_DEFAULT_CAPACITY;
 import static spoon.reflect.ModelElementContainerDefaultCapacities.COMMENT_CONTAINER_DEFAULT_CAPACITY;
@@ -560,4 +562,14 @@ public CtPath getPath() {
 			throw new SpoonException(e);
 		}
 	}
+
+	@Override
+	public Iterator<CtElement> descendantIterator() {
+		return new CtIterator(this);
+	}
+
+	@Override
+	public Iterable<CtElement> asIterable() {
+		return this::descendantIterator;
+	}
 }

src/main/java/spoon/support/visitor/replace/ReplacementVisitor.java

@@ -1304,10 +1304,10 @@ private <K, V extends spoon.reflect.declaration.CtElement> void replaceInMapIfEx
 				if (val != null) {
 					map.put(key, val);
 					val.setParent(shouldBeDeleted.getParent());
-				}else {
+				} else {
 					map.remove(key);
 				}
-			}else {
+			} else {
 				map.remove(key);
 			}
 			listener.set(map);

src/test/java/spoon/reflect/visitor/CtIteratorTest.java

@@ -5,12 +5,14 @@
 import spoon.reflect.declaration.CtElement;
 
 import java.util.ArrayDeque;
+import java.util.Deque;
 
 import static org.junit.Assert.assertEquals;
 
 public class CtIteratorTest {
+
     @Test
-    public void testMethodsInIterator() throws Exception {
+    public void testCtElementIteration() throws Exception {
         // contract: CtIterator must go over all nodes in dfs order
         final Launcher launcher = new Launcher();
         launcher.setArgs(new String[] {"--output-type", "nooutput"});
@@ -22,25 +24,49 @@ public void testMethodsInIterator() throws Exception {
         // get the first Type
         CtElement root = launcher.getFactory().getModel().getAllTypes().iterator().next();
 
-        // use custom CtScanner to assert the proper behaviour
-        CtScannerList counter = new CtScannerList();
+        testCtIterator(root);
+        testAsIterable(root);
+        testDescendantIterator(root);
+    }
 
-        // scan the root to get the elements in DFS order
-        root.accept(counter);
+    public void testCtIterator(CtElement root) {
+        Deque<CtElement> ctElements = getDescendantsInDFS(root);
 
-        // test the iterator by testing that it matches the DFS order as expected
         CtIterator iterator = new CtIterator(root);
         while (iterator.hasNext()) {
-            assertEquals(counter.nodes.pollFirst(), iterator.next());
+            assertEquals(ctElements.pollFirst(), iterator.next());
+        }
+    }
+
+    public void testAsIterable(CtElement root) {
+        Deque<CtElement> ctElements = getDescendantsInDFS(root);
+
+        for (CtElement elem : root.asIterable()) {
+            assertEquals(elem, ctElements.pollFirst());
         }
     }
 
+    public void testDescendantIterator(CtElement root) {
+        Deque<CtElement> ctElements = getDescendantsInDFS(root);
+
+        root.descendantIterator().forEachRemaining((CtElement elem) ->
+                assertEquals(ctElements.pollFirst(), elem)
+        );
+    }
+
+    Deque<CtElement> getDescendantsInDFS(CtElement root) {
+        CtScannerList counter = new CtScannerList();
+        root.accept(counter);
+
+        return counter.nodes;
+    }
+
     /**
      * Class that saves a deque with all the nodes the {@link CtScanner} visits,
      * in DFS order, for the {@link CtIterator} test
      */
     class CtScannerList extends CtScanner {
-        public ArrayDeque<CtElement> nodes = new ArrayDeque<>();
+        public Deque<CtElement> nodes = new ArrayDeque<>();
 
         @Override
         protected void enter(CtElement e) {