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 "base/callback.h"
9 #include "base/check.h"
10 #include "base/containers/stack.h"
11 #include "base/macros.h"
12 
13 namespace 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 // }
22 template <class NodeType>
23 class TreeNodeIterator {
24  public:
25   typedef base::RepeatingCallback<bool(NodeType*)> PruneCallback;
26 
27   // This constructor accepts an optional filter function |prune| which could be
28   // used to prune complete branches of the tree. The filter function will be
29   // evaluated on each tree node and if it evaluates to true the node and all
30   // its descendants will be skipped by the iterator.
TreeNodeIterator(NodeType * node,const PruneCallback & prune)31   TreeNodeIterator(NodeType* node, const PruneCallback& prune)
32       : prune_(prune) {
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     const auto i =
38         std::find_if(node->children().cbegin(), node->children().cend(),
39                      [prune](const auto& child) {
40                        return prune.is_null() || !prune.Run(child.get());
41                      });
42     if (i != node->children().cend())
43       positions_.emplace(node, i - node->children().cbegin());
44   }
45 
TreeNodeIterator(NodeType * node)46   explicit TreeNodeIterator(NodeType* node) {
47     if (!node->children().empty())
48       positions_.emplace(node, 0);
49   }
50 
51   // Returns true if there are more descendants.
has_next()52   bool has_next() const { return !positions_.empty(); }
53 
54   // Returns the next descendant.
Next()55   NodeType* Next() {
56     DCHECK(has_next());
57 
58     // There must always be a valid node in the current Position index.
59     NodeType* result =
60         positions_.top().node->children()[positions_.top().index].get();
61 
62     // Make sure we don't attempt to visit result again.
63     ++positions_.top().index;
64 
65     // Iterate over result's children.
66     positions_.emplace(result, 0);
67 
68     // Advance to next valid node by skipping over the pruned nodes and the
69     // empty Positions. At the end of this loop two cases are possible:
70     // - the current index of the top() Position points to a valid node
71     // - the _position list is empty, the iterator has_next() will return false.
72     while (!positions_.empty()) {
73       auto& top = positions_.top();
74       if (top.index >= top.node->children().size())
75         positions_.pop(); // This Position is all processed, move to the next.
76       else if (!prune_.is_null() &&
77                prune_.Run(top.node->children()[top.index].get()))
78         ++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 {
PositionPosition89     Position(PositionNodeType* node, size_t index) : node(node), index(index) {}
90 
91     PositionNodeType* node;
92     size_t index;
93   };
94 
95   base::stack<Position<NodeType>> positions_;
96   PruneCallback prune_;
97 
98   DISALLOW_COPY_AND_ASSIGN(TreeNodeIterator);
99 };
100 
101 }  // namespace ui
102 
103 #endif  // UI_BASE_MODELS_TREE_NODE_ITERATOR_H_
104