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;
028
029/**
030 * Views elements of a type {@code T} as nodes in a tree, and provides methods to traverse the trees
031 * induced by this traverser.
032 *
033 * <p>For example, the tree
034 *
035 * <pre>{@code
036 *        h
037 *      / | \
038 *     /  e  \
039 *    d       g
040 *   /|\      |
041 *  / | \     f
042 * a  b  c
043 * }</pre>
044 *
045 * <p>can be iterated over in preorder (hdabcegf), postorder (abcdefgh), or breadth-first order
046 * (hdegabcf).
047 *
048 * <p>Null nodes are strictly forbidden.
049 *
050 * <p><b>For Java 8 users:</b> Because this is an abstract class, not an interface, you can't use a
051 * lambda expression to extend it:
052 *
053 * <pre>{@code
054 * // won't work
055 * TreeTraverser<NodeType> traverser = node -> node.getChildNodes();
056 * }</pre>
057 *
058 * Instead, you can pass a lambda expression to the {@code using} factory method:
059 *
060 * <pre>{@code
061 * TreeTraverser<NodeType> traverser = TreeTraverser.using(node -> node.getChildNodes());
062 * }</pre>
063 *
064 * @author Louis Wasserman
065 * @since 15.0
066 */
067@Beta
068@GwtCompatible
069public abstract class TreeTraverser<T> {
070
071  /**
072   * Returns a tree traverser that uses the given function to navigate from a node to its children.
073   * This is useful if the function instance already exists, or so that you can supply a lambda
074   * expressions. If those circumstances don't apply, you probably don't need to use this; subclass
075   * {@code TreeTraverser} and implement its {@link #children} method directly.
076   *
077   * @since 20.0
078   */
079  public static <T> TreeTraverser<T> using(
080      final Function<T, ? extends Iterable<T>> nodeToChildrenFunction) {
081    checkNotNull(nodeToChildrenFunction);
082    return new TreeTraverser<T>() {
083      @Override
084      public Iterable<T> children(T root) {
085        return nodeToChildrenFunction.apply(root);
086      }
087    };
088  }
089
090  /**
091   * Returns the children of the specified node.  Must not contain null.
092   */
093  public abstract Iterable<T> children(T root);
094
095  /**
096   * Returns an unmodifiable iterable over the nodes in a tree structure, using pre-order
097   * traversal. That is, each node's subtrees are traversed after the node itself is returned.
098   *
099   * <p>No guarantees are made about the behavior of the traversal when nodes change while
100   * iteration is in progress or when the iterators generated by {@link #children} are advanced.
101   */
102  public final FluentIterable<T> preOrderTraversal(final T root) {
103    checkNotNull(root);
104    return new FluentIterable<T>() {
105      @Override
106      public UnmodifiableIterator<T> iterator() {
107        return preOrderIterator(root);
108      }
109    };
110  }
111
112  // overridden in BinaryTreeTraverser
113  UnmodifiableIterator<T> preOrderIterator(T root) {
114    return new PreOrderIterator(root);
115  }
116
117  private final class PreOrderIterator extends UnmodifiableIterator<T> {
118    private final Deque<Iterator<T>> stack;
119
120    PreOrderIterator(T root) {
121      this.stack = new ArrayDeque<>();
122      stack.addLast(Iterators.singletonIterator(checkNotNull(root)));
123    }
124
125    @Override
126    public boolean hasNext() {
127      return !stack.isEmpty();
128    }
129
130    @Override
131    public T next() {
132      Iterator<T> itr = stack.getLast(); // throws NSEE if empty
133      T result = checkNotNull(itr.next());
134      if (!itr.hasNext()) {
135        stack.removeLast();
136      }
137      Iterator<T> childItr = children(result).iterator();
138      if (childItr.hasNext()) {
139        stack.addLast(childItr);
140      }
141      return result;
142    }
143  }
144
145  /**
146   * Returns an unmodifiable iterable over the nodes in a tree structure, using post-order
147   * traversal. That is, each node's subtrees are traversed before the node itself is returned.
148   *
149   * <p>No guarantees are made about the behavior of the traversal when nodes change while
150   * iteration is in progress or when the iterators generated by {@link #children} are advanced.
151   */
152  public final FluentIterable<T> postOrderTraversal(final T root) {
153    checkNotNull(root);
154    return new FluentIterable<T>() {
155      @Override
156      public UnmodifiableIterator<T> iterator() {
157        return postOrderIterator(root);
158      }
159    };
160  }
161
162  // overridden in BinaryTreeTraverser
163  UnmodifiableIterator<T> postOrderIterator(T root) {
164    return new PostOrderIterator(root);
165  }
166
167  private static final class PostOrderNode<T> {
168    final T root;
169    final Iterator<T> childIterator;
170
171    PostOrderNode(T root, Iterator<T> childIterator) {
172      this.root = checkNotNull(root);
173      this.childIterator = checkNotNull(childIterator);
174    }
175  }
176
177  private final class PostOrderIterator extends AbstractIterator<T> {
178    private final ArrayDeque<PostOrderNode<T>> stack;
179
180    PostOrderIterator(T root) {
181      this.stack = new ArrayDeque<>();
182      stack.addLast(expand(root));
183    }
184
185    @Override
186    protected T computeNext() {
187      while (!stack.isEmpty()) {
188        PostOrderNode<T> top = stack.getLast();
189        if (top.childIterator.hasNext()) {
190          T child = top.childIterator.next();
191          stack.addLast(expand(child));
192        } else {
193          stack.removeLast();
194          return top.root;
195        }
196      }
197      return endOfData();
198    }
199
200    private PostOrderNode<T> expand(T t) {
201      return new PostOrderNode<T>(t, children(t).iterator());
202    }
203  }
204
205  /**
206   * Returns an unmodifiable iterable over the nodes in a tree structure, using breadth-first
207   * traversal. That is, all the nodes of depth 0 are returned, then depth 1, then 2, and so on.
208   *
209   * <p>No guarantees are made about the behavior of the traversal when nodes change while
210   * iteration is in progress or when the iterators generated by {@link #children} are advanced.
211   */
212  public final FluentIterable<T> breadthFirstTraversal(final T root) {
213    checkNotNull(root);
214    return new FluentIterable<T>() {
215      @Override
216      public UnmodifiableIterator<T> iterator() {
217        return new BreadthFirstIterator(root);
218      }
219    };
220  }
221
222  private final class BreadthFirstIterator extends UnmodifiableIterator<T>
223      implements PeekingIterator<T> {
224    private final Queue<T> queue;
225
226    BreadthFirstIterator(T root) {
227      this.queue = new ArrayDeque<T>();
228      queue.add(root);
229    }
230
231    @Override
232    public boolean hasNext() {
233      return !queue.isEmpty();
234    }
235
236    @Override
237    public T peek() {
238      return queue.element();
239    }
240
241    @Override
242    public T next() {
243      T result = queue.remove();
244      Iterables.addAll(queue, children(result));
245      return result;
246    }
247  }
248}