001/*
002 * Copyright (C) 2012 The Guava Authors
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 * http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package com.google.common.collect;
018
019import static com.google.common.base.Preconditions.checkNotNull;
020
021import com.google.common.annotations.Beta;
022import com.google.common.annotations.GwtCompatible;
023import com.google.common.base.Function;
024import java.util.ArrayDeque;
025import java.util.Deque;
026import java.util.Iterator;
027import java.util.Queue;
028import java.util.function.Consumer;
029
030/**
031 * Views elements of a type {@code T} as nodes in a tree, and provides methods to traverse the trees
032 * induced by this traverser.
033 *
034 * <p>For example, the tree
035 *
036 * <pre>{@code
037 *        h
038 *      / | \
039 *     /  e  \
040 *    d       g
041 *   /|\      |
042 *  / | \     f
043 * a  b  c
044 * }</pre>
045 *
046 * <p>can be iterated over in preorder (hdabcegf), postorder (abcdefgh), or breadth-first order
047 * (hdegabcf).
048 *
049 * <p>Null nodes are strictly forbidden.
050 *
051 * <p><b>For Java 8 users:</b> Because this is an abstract class, not an interface, you can't use a
052 * lambda expression to extend it:
053 *
054 * <pre>{@code
055 * // won't work
056 * TreeTraverser<NodeType> traverser = node -> node.getChildNodes();
057 * }</pre>
058 *
059 * Instead, you can pass a lambda expression to the {@code using} factory method:
060 *
061 * <pre>{@code
062 * TreeTraverser<NodeType> traverser = TreeTraverser.using(node -> node.getChildNodes());
063 * }</pre>
064 *
065 * @author Louis Wasserman
066 * @since 15.0
067 */
068@Beta
069@GwtCompatible
070public abstract class TreeTraverser<T> {
071
072  /**
073   * Returns a tree traverser that uses the given function to navigate from a node to its children.
074   * This is useful if the function instance already exists, or so that you can supply a lambda
075   * expressions. If those circumstances don't apply, you probably don't need to use this; subclass
076   * {@code TreeTraverser} and implement its {@link #children} method directly.
077   *
078   * @since 20.0
079   */
080  public static <T> TreeTraverser<T> using(
081      final Function<T, ? extends Iterable<T>> nodeToChildrenFunction) {
082    checkNotNull(nodeToChildrenFunction);
083    return new TreeTraverser<T>() {
084      @Override
085      public Iterable<T> children(T root) {
086        return nodeToChildrenFunction.apply(root);
087      }
088    };
089  }
090
091  /**
092   * Returns the children of the specified node.  Must not contain null.
093   */
094  public abstract Iterable<T> children(T root);
095
096  /**
097   * Returns an unmodifiable iterable over the nodes in a tree structure, using pre-order
098   * traversal. That is, each node's subtrees are traversed after the node itself is returned.
099   *
100   * <p>No guarantees are made about the behavior of the traversal when nodes change while
101   * iteration is in progress or when the iterators generated by {@link #children} are advanced.
102   */
103  public final FluentIterable<T> preOrderTraversal(final T root) {
104    checkNotNull(root);
105    return new FluentIterable<T>() {
106      @Override
107      public UnmodifiableIterator<T> iterator() {
108        return preOrderIterator(root);
109      }
110
111      @Override
112      public void forEach(Consumer<? super T> action) {
113        checkNotNull(action);
114        new Consumer<T>() {
115          @Override
116          public void accept(T t) {
117            action.accept(t);
118            children(t).forEach(this);
119          }
120        }.accept(root);
121      }
122    };
123  }
124
125  // overridden in BinaryTreeTraverser
126  UnmodifiableIterator<T> preOrderIterator(T root) {
127    return new PreOrderIterator(root);
128  }
129
130  private final class PreOrderIterator extends UnmodifiableIterator<T> {
131    private final Deque<Iterator<T>> stack;
132
133    PreOrderIterator(T root) {
134      this.stack = new ArrayDeque<Iterator<T>>();
135      stack.addLast(Iterators.singletonIterator(checkNotNull(root)));
136    }
137
138    @Override
139    public boolean hasNext() {
140      return !stack.isEmpty();
141    }
142
143    @Override
144    public T next() {
145      Iterator<T> itr = stack.getLast(); // throws NSEE if empty
146      T result = checkNotNull(itr.next());
147      if (!itr.hasNext()) {
148        stack.removeLast();
149      }
150      Iterator<T> childItr = children(result).iterator();
151      if (childItr.hasNext()) {
152        stack.addLast(childItr);
153      }
154      return result;
155    }
156  }
157
158  /**
159   * Returns an unmodifiable iterable over the nodes in a tree structure, using post-order
160   * traversal. That is, each node's subtrees are traversed before the node itself is returned.
161   *
162   * <p>No guarantees are made about the behavior of the traversal when nodes change while
163   * iteration is in progress or when the iterators generated by {@link #children} are advanced.
164   */
165  public final FluentIterable<T> postOrderTraversal(final T root) {
166    checkNotNull(root);
167    return new FluentIterable<T>() {
168      @Override
169      public UnmodifiableIterator<T> iterator() {
170        return postOrderIterator(root);
171      }
172
173      @Override
174      public void forEach(Consumer<? super T> action) {
175        checkNotNull(action);
176        new Consumer<T>() {
177          @Override
178          public void accept(T t) {
179            children(t).forEach(this);
180            action.accept(t);
181          }
182        }.accept(root);
183      }
184    };
185  }
186
187  // overridden in BinaryTreeTraverser
188  UnmodifiableIterator<T> postOrderIterator(T root) {
189    return new PostOrderIterator(root);
190  }
191
192  private static final class PostOrderNode<T> {
193    final T root;
194    final Iterator<T> childIterator;
195
196    PostOrderNode(T root, Iterator<T> childIterator) {
197      this.root = checkNotNull(root);
198      this.childIterator = checkNotNull(childIterator);
199    }
200  }
201
202  private final class PostOrderIterator extends AbstractIterator<T> {
203    private final ArrayDeque<PostOrderNode<T>> stack;
204
205    PostOrderIterator(T root) {
206      this.stack = new ArrayDeque<PostOrderNode<T>>();
207      stack.addLast(expand(root));
208    }
209
210    @Override
211    protected T computeNext() {
212      while (!stack.isEmpty()) {
213        PostOrderNode<T> top = stack.getLast();
214        if (top.childIterator.hasNext()) {
215          T child = top.childIterator.next();
216          stack.addLast(expand(child));
217        } else {
218          stack.removeLast();
219          return top.root;
220        }
221      }
222      return endOfData();
223    }
224
225    private PostOrderNode<T> expand(T t) {
226      return new PostOrderNode<T>(t, children(t).iterator());
227    }
228  }
229
230  /**
231   * Returns an unmodifiable iterable over the nodes in a tree structure, using breadth-first
232   * traversal. That is, all the nodes of depth 0 are returned, then depth 1, then 2, and so on.
233   *
234   * <p>No guarantees are made about the behavior of the traversal when nodes change while
235   * iteration is in progress or when the iterators generated by {@link #children} are advanced.
236   */
237  public final FluentIterable<T> breadthFirstTraversal(final T root) {
238    checkNotNull(root);
239    return new FluentIterable<T>() {
240      @Override
241      public UnmodifiableIterator<T> iterator() {
242        return new BreadthFirstIterator(root);
243      }
244    };
245  }
246
247  private final class BreadthFirstIterator extends UnmodifiableIterator<T>
248      implements PeekingIterator<T> {
249    private final Queue<T> queue;
250
251    BreadthFirstIterator(T root) {
252      this.queue = new ArrayDeque<T>();
253      queue.add(root);
254    }
255
256    @Override
257    public boolean hasNext() {
258      return !queue.isEmpty();
259    }
260
261    @Override
262    public T peek() {
263      return queue.element();
264    }
265
266    @Override
267    public T next() {
268      T result = queue.remove();
269      Iterables.addAll(queue, children(result));
270      return result;
271    }
272  }
273}