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