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