ui/base/models/tree_node_iterator.h

   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
  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   // 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_