|
21 | 21 | package com.apple.foundationdb.record.query.plan.cascades; |
22 | 22 |
|
23 | 23 | import com.google.common.base.Verify; |
| 24 | +import com.google.common.collect.ImmutableList; |
24 | 25 | import com.google.common.collect.Iterables; |
25 | 26 | import org.apache.commons.lang3.tuple.MutablePair; |
26 | 27 | import org.apache.commons.lang3.tuple.Pair; |
|
30 | 31 | import java.util.ArrayDeque; |
31 | 32 | import java.util.Deque; |
32 | 33 | import java.util.Iterator; |
33 | | -import java.util.List; |
34 | 34 | import java.util.NoSuchElementException; |
| 35 | +import java.util.function.Predicate; |
35 | 36 |
|
36 | 37 | /** |
37 | 38 | * An iterator that accesses all elements of a {@link TreeLike} object in pre-order fashion. |
|
44 | 45 | * @param <T> The type of the iterator element. |
45 | 46 | */ |
46 | 47 | @NotThreadSafe |
47 | | -public final class PreOrderIterator<T extends TreeLike<T>> implements Iterator<T> { |
| 48 | +public final class PreOrderPruningIterator<T extends TreeLike<T>> implements Iterator<T> { |
48 | 49 |
|
49 | 50 | @Nonnull |
50 | 51 | private final Deque<Pair<Iterable<? extends T>, Integer>> stack; |
51 | 52 |
|
| 53 | + @Nonnull |
| 54 | + private final Predicate<T> descendIntoChildrenPredicate; |
| 55 | + |
52 | 56 | private static final int INITIAL_POSITION = -1; |
53 | 57 |
|
54 | | - private PreOrderIterator(@Nonnull final T traversable) { |
| 58 | + private PreOrderPruningIterator(@Nonnull final T treeLike) { |
| 59 | + this(treeLike, ignored -> true); |
| 60 | + } |
| 61 | + |
| 62 | + private PreOrderPruningIterator(@Nonnull final T treeLike, @Nonnull final Predicate<T> descendIntoChildrenPredicate) { |
55 | 63 | // initialize the stack with the {@link TreeLike}'s depth as capacity to avoid resizing. |
56 | | - stack = new ArrayDeque<>(traversable.height()); |
| 64 | + stack = new ArrayDeque<>(treeLike.height()); |
57 | 65 | // this is the only list allocation done to put the root in the stack. |
58 | 66 | // all the remaining lists added to the stack are references to children |
59 | 67 | // lists (copy by reference). |
60 | | - stack.push(MutablePair.of(List.of(traversable.getThis()), INITIAL_POSITION)); |
| 68 | + stack.push(MutablePair.of(ImmutableList.of(treeLike.getThis()), INITIAL_POSITION)); |
| 69 | + this.descendIntoChildrenPredicate = descendIntoChildrenPredicate; |
61 | 70 | } |
62 | 71 |
|
63 | 72 | @Override |
@@ -97,17 +106,42 @@ public T next() { |
97 | 106 | // that incrementing it would lead to finding that next element correctly. |
98 | 107 | final var result = Iterables.get(currentLevelItems, nextItemIndex); |
99 | 108 | top.setValue(nextItemIndex); |
100 | | - final var resultChildren = result.getChildren(); |
101 | 109 |
|
102 | | - // descend immediately to the children (if any) so to conform to pre-order DFS semantics. |
103 | | - if (!Iterables.isEmpty(resultChildren)) { |
104 | | - stack.add(MutablePair.of(resultChildren, INITIAL_POSITION)); |
| 110 | + // test whether we should descend into the children of the current node or prune them |
| 111 | + // and continue to the next unexplored node in pre-order. |
| 112 | + if (descendIntoChildrenPredicate.test(result)) { |
| 113 | + final var resultChildren = result.getChildren(); |
| 114 | + |
| 115 | + // descend immediately to the children (if any) so to conform to pre-order DFS semantics. |
| 116 | + if (!Iterables.isEmpty(resultChildren)) { |
| 117 | + stack.add(MutablePair.of(resultChildren, INITIAL_POSITION)); |
| 118 | + } |
105 | 119 | } |
106 | 120 | return result; |
107 | 121 | } |
108 | 122 |
|
| 123 | + /** |
| 124 | + * Retuns an iterator that traverses {@code treeLike} in pre-order. |
| 125 | + * @param treeLike a {@link TreeLike} treeLike. |
| 126 | + * @param <T> The type of {@code treeLike} items. |
| 127 | + * @return an iterator that traverses the items in pre-order. |
| 128 | + */ |
| 129 | + @Nonnull |
| 130 | + public static <T extends TreeLike<T>> PreOrderPruningIterator<T> over(@Nonnull final T treeLike) { |
| 131 | + return new PreOrderPruningIterator<>(treeLike); |
| 132 | + } |
| 133 | + |
| 134 | + /** |
| 135 | + * Retuns an iterator that traverses {@code treeLike} in pre-order with a children-pruning condition. |
| 136 | + * @param treeLike a {@link TreeLike} treeLike. |
| 137 | + * @param descendIntoChildrenPredicate a condition that determines whether, for the currently visited node, the iterator |
| 138 | + * should descend to the children or not. |
| 139 | + * @param <T> The type of {@code treeLike} items. |
| 140 | + * @return an iterator that traverses the items in pre-order. |
| 141 | + */ |
109 | 142 | @Nonnull |
110 | | - public static <T extends TreeLike<T>> PreOrderIterator<T> over(@Nonnull final T traversable) { |
111 | | - return new PreOrderIterator<>(traversable); |
| 143 | + public static <T extends TreeLike<T>> PreOrderPruningIterator<T> overWithPruningPredicate(@Nonnull final T treeLike, |
| 144 | + @Nonnull Predicate<T> descendIntoChildrenPredicate) { |
| 145 | + return new PreOrderPruningIterator<>(treeLike, descendIntoChildrenPredicate); |
112 | 146 | } |
113 | 147 | } |
0 commit comments