1// Copyright (c) 2017 Google Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#ifndef SOURCE_OPT_TREE_ITERATOR_H_
16#define SOURCE_OPT_TREE_ITERATOR_H_
17
18#include <stack>
19#include <type_traits>
20#include <utility>
21
22namespace spvtools {
23namespace opt {
24
25// Helper class to iterate over a tree in a depth first order.
26// The class assumes the data structure is a tree, tree node type implements a
27// forward iterator.
28// At each step, the iterator holds the pointer to the current node and state of
29// the walk.
30// The state is recorded by stacking the iteration position of the node
31// children. To move to the next node, the iterator:
32//  - Looks at the top of the stack;
33//  - Sets the node behind the iterator as the current node;
34//  - Increments the iterator if it has more children to visit, pops otherwise;
35//  - If the current node has children, the children iterator is pushed into the
36//    stack.
37template <typename NodeTy>
38class TreeDFIterator {
39  static_assert(!std::is_pointer<NodeTy>::value &&
40                    !std::is_reference<NodeTy>::value,
41                "NodeTy should be a class");
42  // Type alias to keep track of the const qualifier.
43  using NodeIterator =
44      typename std::conditional<std::is_const<NodeTy>::value,
45                                typename NodeTy::const_iterator,
46                                typename NodeTy::iterator>::type;
47
48  // Type alias to keep track of the const qualifier.
49  using NodePtr = NodeTy*;
50
51 public:
52  // Standard iterator interface.
53  using reference = NodeTy&;
54  using value_type = NodeTy;
55
56  explicit inline TreeDFIterator(NodePtr top_node) : current_(top_node) {
57    if (current_ && current_->begin() != current_->end())
58      parent_iterators_.emplace(make_pair(current_, current_->begin()));
59  }
60
61  // end() iterator.
62  inline TreeDFIterator() : TreeDFIterator(nullptr) {}
63
64  bool operator==(const TreeDFIterator& x) const {
65    return current_ == x.current_;
66  }
67
68  bool operator!=(const TreeDFIterator& x) const { return !(*this == x); }
69
70  reference operator*() const { return *current_; }
71
72  NodePtr operator->() const { return current_; }
73
74  TreeDFIterator& operator++() {
75    MoveToNextNode();
76    return *this;
77  }
78
79  TreeDFIterator operator++(int) {
80    TreeDFIterator tmp = *this;
81    ++*this;
82    return tmp;
83  }
84
85 private:
86  // Moves the iterator to the next node in the tree.
87  // If we are at the end, do nothing, otherwise
88  // if our current node has children, use the children iterator and push the
89  // current node into the stack.
90  // If we reach the end of the local iterator, pop it.
91  inline void MoveToNextNode() {
92    if (!current_) return;
93    if (parent_iterators_.empty()) {
94      current_ = nullptr;
95      return;
96    }
97    std::pair<NodePtr, NodeIterator>& next_it = parent_iterators_.top();
98    // Set the new node.
99    current_ = *next_it.second;
100    // Update the iterator for the next child.
101    ++next_it.second;
102    // If we finished with node, pop it.
103    if (next_it.first->end() == next_it.second) parent_iterators_.pop();
104    // If our current node is not a leaf, store the iteration state for later.
105    if (current_->begin() != current_->end())
106      parent_iterators_.emplace(make_pair(current_, current_->begin()));
107  }
108
109  // The current node of the tree.
110  NodePtr current_;
111  // State of the tree walk: each pair contains the parent node (which has been
112  // already visited) and the iterator of the next children to visit.
113  // When all the children has been visited, we pop the entry, get the next
114  // child and push back the pair if the children iterator is not end().
115  std::stack<std::pair<NodePtr, NodeIterator>> parent_iterators_;
116};
117
118// Helper class to iterate over a tree in a depth first post-order.
119// The class assumes the data structure is a tree, tree node type implements a
120// forward iterator.
121// At each step, the iterator holds the pointer to the current node and state of
122// the walk.
123// The state is recorded by stacking the iteration position of the node
124// children. To move to the next node, the iterator:
125//  - Looks at the top of the stack;
126//  - If the children iterator has reach the end, then the node become the
127//    current one and we pop the stack;
128//  - Otherwise, we save the child and increment the iterator;
129//  - We walk the child sub-tree until we find a leaf, stacking all non-leaves
130//    states (pair of node pointer and child iterator) as we walk it.
131template <typename NodeTy>
132class PostOrderTreeDFIterator {
133  static_assert(!std::is_pointer<NodeTy>::value &&
134                    !std::is_reference<NodeTy>::value,
135                "NodeTy should be a class");
136  // Type alias to keep track of the const qualifier.
137  using NodeIterator =
138      typename std::conditional<std::is_const<NodeTy>::value,
139                                typename NodeTy::const_iterator,
140                                typename NodeTy::iterator>::type;
141
142  // Type alias to keep track of the const qualifier.
143  using NodePtr = NodeTy*;
144
145 public:
146  // Standard iterator interface.
147  using reference = NodeTy&;
148  using value_type = NodeTy;
149
150  static inline PostOrderTreeDFIterator begin(NodePtr top_node) {
151    return PostOrderTreeDFIterator(top_node);
152  }
153
154  static inline PostOrderTreeDFIterator end(NodePtr sentinel_node) {
155    return PostOrderTreeDFIterator(sentinel_node, false);
156  }
157
158  bool operator==(const PostOrderTreeDFIterator& x) const {
159    return current_ == x.current_;
160  }
161
162  bool operator!=(const PostOrderTreeDFIterator& x) const {
163    return !(*this == x);
164  }
165
166  reference operator*() const { return *current_; }
167
168  NodePtr operator->() const { return current_; }
169
170  PostOrderTreeDFIterator& operator++() {
171    MoveToNextNode();
172    return *this;
173  }
174
175  PostOrderTreeDFIterator operator++(int) {
176    PostOrderTreeDFIterator tmp = *this;
177    ++*this;
178    return tmp;
179  }
180
181 private:
182  explicit inline PostOrderTreeDFIterator(NodePtr top_node)
183      : current_(top_node) {
184    if (current_) WalkToLeaf();
185  }
186
187  // Constructor for the "end()" iterator.
188  // |end_sentinel| is the value that acts as end value (can be null). The bool
189  // parameters is to distinguish from the start() Ctor.
190  inline PostOrderTreeDFIterator(NodePtr sentinel_node, bool)
191      : current_(sentinel_node) {}
192
193  // Moves the iterator to the next node in the tree.
194  // If we are at the end, do nothing, otherwise
195  // if our current node has children, use the children iterator and push the
196  // current node into the stack.
197  // If we reach the end of the local iterator, pop it.
198  inline void MoveToNextNode() {
199    if (!current_) return;
200    if (parent_iterators_.empty()) {
201      current_ = nullptr;
202      return;
203    }
204    std::pair<NodePtr, NodeIterator>& next_it = parent_iterators_.top();
205    // If we visited all children, the current node is the top of the stack.
206    if (next_it.second == next_it.first->end()) {
207      // Set the new node.
208      current_ = next_it.first;
209      parent_iterators_.pop();
210      return;
211    }
212    // We have more children to visit, set the current node to the first child
213    // and dive to leaf.
214    current_ = *next_it.second;
215    // Update the iterator for the next child (avoid unneeded pop).
216    ++next_it.second;
217    WalkToLeaf();
218  }
219
220  // Moves the iterator to the next node in the tree.
221  // If we are at the end, do nothing, otherwise
222  // if our current node has children, use the children iterator and push the
223  // current node into the stack.
224  // If we reach the end of the local iterator, pop it.
225  inline void WalkToLeaf() {
226    while (current_->begin() != current_->end()) {
227      NodeIterator next = ++current_->begin();
228      parent_iterators_.emplace(make_pair(current_, next));
229      // Set the first child as the new node.
230      current_ = *current_->begin();
231    }
232  }
233
234  // The current node of the tree.
235  NodePtr current_;
236  // State of the tree walk: each pair contains the parent node and the iterator
237  // of the next children to visit.
238  // When all the children has been visited, we pop the first entry and the
239  // parent node become the current node.
240  std::stack<std::pair<NodePtr, NodeIterator>> parent_iterators_;
241};
242
243}  // namespace opt
244}  // namespace spvtools
245
246#endif  // SOURCE_OPT_TREE_ITERATOR_H_
247