1// Copyright (c) 2011 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#ifndef UI_BASE_MODELS_TREE_NODE_ITERATOR_H_
6#define UI_BASE_MODELS_TREE_NODE_ITERATOR_H_
7
8#include <stack>
9
10#include "base/basictypes.h"
11#include "base/callback.h"
12#include "base/logging.h"
13
14namespace ui {
15
16// Iterator that iterates over the descendants of a node. The iteration does
17// not include the node itself, only the descendants. The following illustrates
18// typical usage:
19// while (iterator.has_next()) {
20//   Node* node = iterator.Next();
21//   // do something with node.
22// }
23template <class NodeType>
24class TreeNodeIterator {
25 public:
26  typedef base::Callback<bool(NodeType*)> PruneCallback;
27
28  // This contructor accepts an optional filter function |prune| which could be
29  // used to prune complete branches of the tree. The filter function will be
30  // evaluated on each tree node and if it evaluates to true the node and all
31  // its descendants will be skipped by the iterator.
32  TreeNodeIterator(NodeType* node, const PruneCallback& prune)
33      : prune_(prune) {
34    int index = 0;
35
36    // Move forward through the children list until the first non prunable node.
37    // This is to satisfy the iterator invariant that the current index in the
38    // Position at the top of the _positions list must point to a node the
39    // iterator will be returning.
40    for (; index < node->child_count(); ++index)
41      if (prune.is_null() || !prune.Run(node->GetChild(index)))
42        break;
43
44    if (index < node->child_count())
45      positions_.push(Position<NodeType>(node, index));
46  }
47
48  explicit TreeNodeIterator(NodeType* node) {
49    if (!node->empty())
50      positions_.push(Position<NodeType>(node, 0));
51  }
52
53  // Returns true if there are more descendants.
54  bool has_next() const { return !positions_.empty(); }
55
56  // Returns the next descendant.
57  NodeType* Next() {
58    if (!has_next()) {
59      NOTREACHED();
60      return NULL;
61    }
62
63    // There must always be a valid node in the current Position index.
64    NodeType* result = positions_.top().node->GetChild(positions_.top().index);
65
66    // Make sure we don't attempt to visit result again.
67    positions_.top().index++;
68
69    // Iterate over result's children.
70    positions_.push(Position<NodeType>(result, 0));
71
72    // Advance to next valid node by skipping over the pruned nodes and the
73    // empty Positions. At the end of this loop two cases are possible:
74    // - the current index of the top() Position points to a valid node
75    // - the _position list is empty, the iterator has_next() will return false.
76    while (!positions_.empty()) {
77      if (positions_.top().index >= positions_.top().node->child_count())
78        positions_.pop(); // This Position is all processed, move to the next.
79      else if (!prune_.is_null() &&
80          prune_.Run(positions_.top().node->GetChild(positions_.top().index)))
81        positions_.top().index++;  // Prune the branch.
82      else
83        break;  // Now positioned at the next node to be returned.
84    }
85
86    return result;
87  }
88
89 private:
90  template <class PositionNodeType>
91  struct Position {
92    Position(PositionNodeType* node, int index) : node(node), index(index) {}
93    Position() : node(NULL), index(-1) {}
94
95    PositionNodeType* node;
96    int index;
97  };
98
99  std::stack<Position<NodeType> > positions_;
100  PruneCallback prune_;
101
102  DISALLOW_COPY_AND_ASSIGN(TreeNodeIterator);
103};
104
105}  // namespace ui
106
107#endif  // UI_BASE_MODELS_TREE_NODE_ITERATOR_H_
108