1 files changed, 103 insertions, 3 deletions
diff --git a/plugins/infra/core/org.eclipse.papyrus.infra.tools/src/org/eclipse/papyrus/infra/tools/util/Iterators2.java b/plugins/infra/core/org.eclipse.papyrus.infra.tools/src/org/eclipse/papyrus/infra/tools/util/Iterators2.java
index ef0a1257ff9..c83ded14c59 100644
--- a/plugins/infra/core/org.eclipse.papyrus.infra.tools/src/org/eclipse/papyrus/infra/tools/util/Iterators2.java
+++ b/plugins/infra/core/org.eclipse.papyrus.infra.tools/src/org/eclipse/papyrus/infra/tools/util/Iterators2.java
@@ -1,5 +1,5 @@
 /*****************************************************************************
- * Copyright (c) 2014, 2020 Christian W. Damus, CEA LIST, and others.
+ * Copyright (c) 2014, 2021 Christian W. Damus, CEA LIST, and others.
  *
  * All rights reserved. This program and the accompanying materials
  * are made available under the terms of the Eclipse Public License 2.0
@@ -18,6 +18,7 @@ package org.eclipse.papyrus.infra.tools.util;
 import java.util.Iterator;
 import java.util.Spliterator;
 import java.util.Spliterators;
+import java.util.function.Predicate;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
@@ -65,6 +66,49 @@ public class Iterators2 {
 	}
 
 	/**
+	 * Filters an EMF tree iterator for elements of a particular {@code type}.
+	 *
+	 * @param treeIterator
+	 *            the tree iterator to filter
+	 * @param type
+	 *            the type of elements to include in the filtered tree iterator
+	 * @param shouldPrune
+	 *            a predicate that tests {@code true} on elements whose sub-trees should be {@linkplain TreeIterator#prune() pruned}.
+	 *            A {@link null} predicate will never prune
+	 * @return the self-pruning tree iterator
+	 *
+	 * @since 4.2
+	 */
+	public static <T> Iterator<T> autoPrune(final TreeIterator<T> treeIterator, Predicate<? super T> shouldPrune) {
+		class SelfPruningIterator extends AbstractIterator<T> {
+			@Override
+			protected T computeNext() {
+				T result;
+
+				if (!treeIterator.hasNext()) {
+					result = endOfData();
+				} else {
+					result = treeIterator.next();
+					if (shouldPrune.test(result)) {
+						// Prune, now
+						treeIterator.prune();
+					}
+				}
+
+				return result;
+			}
+		}
+
+		Iterator<T> result = treeIterator;
+
+		if (shouldPrune != null) {
+			result = new SelfPruningIterator();
+		}
+
+		return result;
+	}
+
+	/**
 	 * Obtain a spliterator over an EMF tree iterator. The spliterator will have characteristics implied by an
 	 * EMF content tree, namely:
 	 * <ul>
@@ -78,9 +122,40 @@ public class Iterators2 {
 	 * @param treeIterator
 	 *            a tree iterator
 	 * @return a spliterator over the tree
+	 *
+	 * @see #spliterator(TreeIterator, Predicate)
 	 */
 	public static <T> Spliterator<T> spliterator(final TreeIterator<T> treeIterator) {
-		return Spliterators.spliteratorUnknownSize(treeIterator, Spliterator.ORDERED | Spliterator.DISTINCT | Spliterator.NONNULL);
+		return spliterator(treeIterator, (Predicate<T>) null);
+	}
+
+	/**
+	 * Obtain a spliterator over an EMF tree iterator that prunes itself automatically. The spliterator will have characteristics implied by an
+	 * EMF content tree, namely:
+	 * <ul>
+	 * <li>{@link Spliterator#ORDERED}</li>
+	 * <li>{@link Spliterator#DISTINCT}</li>
+	 * <li>{@link Spliterator#NONNULL}</li>
+	 * </ul>
+	 *
+	 * @param <T>
+	 *            the tree element type
+	 * @param treeIterator
+	 *            a tree iterator
+	 * @param shouldPrune
+	 *            a predicate that tests {@code true} on elements whose sub-trees should be {@linkplain TreeIterator#prune() pruned}.
+	 *            A {@link null} predicate will never prune
+	 * @return a spliterator over the tree
+	 *
+	 * @since 4.2
+	 *
+	 * @see #autoPrune(TreeIterator, Predicate)
+	 * @see #stream(TreeIterator, Predicate)
+	 * @see #spliterator(TreeIterator)
+	 */
+	public static <T> Spliterator<T> spliterator(final TreeIterator<T> treeIterator, Predicate<? super T> shouldPrune) {
+		Iterator<T> delegate = autoPrune(treeIterator, shouldPrune);
+		return Spliterators.spliteratorUnknownSize(delegate, Spliterator.ORDERED | Spliterator.DISTINCT | Spliterator.NONNULL);
 	}
 
 	/**
@@ -91,9 +166,34 @@ public class Iterators2 {
 	 * @param treeIterator
 	 *            a tree iterator
 	 * @return a stream over the tree
+	 *
+	 * @see #stream(TreeIterator, Predicate)r
 	 */
 	public static <T> Stream<T> stream(final TreeIterator<T> treeIterator) {
-		return StreamSupport.stream(spliterator(treeIterator), false);
+		return stream(treeIterator, (Predicate<T>) null);
+	}
+
+	/**
+	 * Obtain a stream over an EMF tree iterator that prunes itself automatically.
+	 *
+	 * @param <T>
+	 *            the tree element type
+	 * @param treeIterator
+	 *            a tree iterator
+	 * @param shouldPrune
+	 *            a predicate that tests {@code true} on elements whose sub-trees should be {@linkplain TreeIterator#prune() pruned}.
+	 *            A {@link null} predicate will never prune
+	 *
+	 * @return a self-pruning stream over the tree
+	 *
+	 * @since 4.2
+	 *
+	 * @see #autoPrune(TreeIterator, Predicate)
+	 * @see #spliterator(TreeIterator, Predicate)
+	 * @see #stream(TreeIterator)
+	 */
+	public static <T> Stream<T> stream(final TreeIterator<T> treeIterator, Predicate<? super T> shouldPrune) {
+		return StreamSupport.stream(spliterator(treeIterator, shouldPrune), false);
 	}
 
 }