1 | /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. |
---|---|
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 | |
16 | #ifndef TENSORFLOW_CORE_GRAPH_EDGESET_H_ |
17 | #define TENSORFLOW_CORE_GRAPH_EDGESET_H_ |
18 | |
19 | #include <stddef.h> |
20 | |
21 | #include "tensorflow/core/lib/gtl/flatset.h" |
22 | #include "tensorflow/core/platform/logging.h" |
23 | #include "tensorflow/core/platform/macros.h" |
24 | #include "tensorflow/core/platform/types.h" |
25 | namespace tensorflow { |
26 | |
27 | class Edge; |
28 | |
29 | // An unordered set of edges. Uses very little memory for small sets. |
30 | // Unlike gtl::FlatSet, EdgeSet does NOT allow mutations during |
31 | // iteration. |
32 | class EdgeSet { |
33 | public: |
34 | EdgeSet(); |
35 | ~EdgeSet(); |
36 | |
37 | typedef const Edge* key_type; |
38 | typedef const Edge* value_type; |
39 | typedef size_t size_type; |
40 | typedef ptrdiff_t difference_type; |
41 | |
42 | class const_iterator; |
43 | typedef const_iterator iterator; |
44 | |
45 | bool empty() const; |
46 | size_type size() const; |
47 | void clear(); |
48 | std::pair<iterator, bool> insert(value_type value); |
49 | size_type erase(key_type key); |
50 | void reserve(size_type new_size) { |
51 | if (new_size > kInline) { |
52 | auto s = new gtl::FlatSet<const Edge*>(new_size); |
53 | s->insert(reinterpret_cast<const Edge**>(std::begin(ptrs_)), |
54 | reinterpret_cast<const Edge**>(&ptrs_[0] + size())); |
55 | ptrs_[0] = this; |
56 | ptrs_[1] = s; |
57 | } |
58 | } |
59 | |
60 | // Caller is not allowed to mutate the EdgeSet while iterating. |
61 | const_iterator begin() const; |
62 | const_iterator end() const; |
63 | |
64 | private: |
65 | // Up to kInline elements are stored directly in ptrs_ (nullptr means none). |
66 | // If ptrs_[0] == this then ptrs_[1] points to a set<const Edge*>. |
67 | // kInline must be >= 2, and is chosen such that ptrs_ fills a 64 byte |
68 | // cacheline. |
69 | static constexpr int kInline = 64 / sizeof(const void*); |
70 | const void* ptrs_[kInline]; |
71 | |
72 | gtl::FlatSet<const Edge*>* get_set() const { |
73 | if (ptrs_[0] == this) { |
74 | return static_cast<gtl::FlatSet<const Edge*>*>( |
75 | const_cast<void*>(ptrs_[1])); |
76 | } else { |
77 | return nullptr; |
78 | } |
79 | } |
80 | |
81 | // To detect mutations while iterating. |
82 | #ifdef NDEBUG |
83 | void RegisterMutation() {} |
84 | #else |
85 | uint32 mutations_ = 0; |
86 | void RegisterMutation() { mutations_++; } |
87 | #endif |
88 | |
89 | TF_DISALLOW_COPY_AND_ASSIGN(EdgeSet); |
90 | }; |
91 | |
92 | class EdgeSet::const_iterator { |
93 | public: |
94 | typedef typename EdgeSet::value_type value_type; |
95 | typedef const typename EdgeSet::value_type& reference; |
96 | typedef const typename EdgeSet::value_type* pointer; |
97 | typedef typename EdgeSet::difference_type difference_type; |
98 | typedef std::forward_iterator_tag iterator_category; |
99 | |
100 | const_iterator() {} |
101 | |
102 | const_iterator& operator++(); |
103 | const_iterator operator++(int /*unused*/); |
104 | const value_type* operator->() const; |
105 | value_type operator*() const; |
106 | bool operator==(const const_iterator& other) const; |
107 | bool operator!=(const const_iterator& other) const { |
108 | return !(*this == other); |
109 | } |
110 | |
111 | private: |
112 | friend class EdgeSet; |
113 | |
114 | void const* const* array_iter_ = nullptr; |
115 | typename gtl::FlatSet<const Edge*>::const_iterator tree_iter_; |
116 | |
117 | #ifdef NDEBUG |
118 | inline void Init(const EdgeSet* e) {} |
119 | inline void CheckNoMutations() const {} |
120 | #else |
121 | inline void Init(const EdgeSet* e) { |
122 | owner_ = e; |
123 | init_mutations_ = e->mutations_; |
124 | } |
125 | inline void CheckNoMutations() const { |
126 | CHECK_EQ(init_mutations_, owner_->mutations_); |
127 | } |
128 | const EdgeSet* owner_ = nullptr; |
129 | uint32 init_mutations_ = 0; |
130 | #endif |
131 | }; |
132 | |
133 | inline EdgeSet::EdgeSet() { |
134 | for (int i = 0; i < kInline; i++) { |
135 | ptrs_[i] = nullptr; |
136 | } |
137 | } |
138 | |
139 | inline EdgeSet::~EdgeSet() { delete get_set(); } |
140 | |
141 | inline bool EdgeSet::empty() const { return size() == 0; } |
142 | |
143 | inline EdgeSet::size_type EdgeSet::size() const { |
144 | auto s = get_set(); |
145 | if (s) { |
146 | return s->size(); |
147 | } else { |
148 | size_t result = 0; |
149 | for (int i = 0; i < kInline; i++) { |
150 | if (ptrs_[i]) result++; |
151 | } |
152 | return result; |
153 | } |
154 | } |
155 | |
156 | inline void EdgeSet::clear() { |
157 | RegisterMutation(); |
158 | delete get_set(); |
159 | for (int i = 0; i < kInline; i++) { |
160 | ptrs_[i] = nullptr; |
161 | } |
162 | } |
163 | |
164 | inline EdgeSet::const_iterator EdgeSet::begin() const { |
165 | const_iterator ci; |
166 | ci.Init(this); |
167 | auto s = get_set(); |
168 | if (s) { |
169 | ci.tree_iter_ = s->begin(); |
170 | } else { |
171 | ci.array_iter_ = &ptrs_[0]; |
172 | } |
173 | return ci; |
174 | } |
175 | |
176 | inline EdgeSet::const_iterator EdgeSet::end() const { |
177 | const_iterator ci; |
178 | ci.Init(this); |
179 | auto s = get_set(); |
180 | if (s) { |
181 | ci.tree_iter_ = s->end(); |
182 | } else { |
183 | ci.array_iter_ = &ptrs_[size()]; |
184 | } |
185 | return ci; |
186 | } |
187 | |
188 | inline EdgeSet::const_iterator& EdgeSet::const_iterator::operator++() { |
189 | CheckNoMutations(); |
190 | if (array_iter_ != nullptr) { |
191 | ++array_iter_; |
192 | } else { |
193 | ++tree_iter_; |
194 | } |
195 | return *this; |
196 | } |
197 | |
198 | inline EdgeSet::const_iterator EdgeSet::const_iterator::operator++( |
199 | int /*unused*/) { |
200 | CheckNoMutations(); |
201 | const_iterator tmp = *this; |
202 | operator++(); |
203 | return tmp; |
204 | } |
205 | |
206 | // gcc's set and multiset always use const_iterator since it will otherwise |
207 | // allow modification of keys. |
208 | inline const EdgeSet::const_iterator::value_type* EdgeSet::const_iterator:: |
209 | operator->() const { |
210 | CheckNoMutations(); |
211 | if (array_iter_ != nullptr) { |
212 | return reinterpret_cast<const value_type*>(array_iter_); |
213 | } else { |
214 | return tree_iter_.operator->(); |
215 | } |
216 | } |
217 | |
218 | // gcc's set and multiset always use const_iterator since it will otherwise |
219 | // allow modification of keys. |
220 | inline EdgeSet::const_iterator::value_type EdgeSet::const_iterator::operator*() |
221 | const { |
222 | CheckNoMutations(); |
223 | if (array_iter_ != nullptr) { |
224 | return static_cast<value_type>(*array_iter_); |
225 | } else { |
226 | return *tree_iter_; |
227 | } |
228 | } |
229 | |
230 | inline bool EdgeSet::const_iterator::operator==( |
231 | const const_iterator& other) const { |
232 | DCHECK((array_iter_ == nullptr) == (other.array_iter_ == nullptr)) |
233 | << "Iterators being compared must be from same set that has not " |
234 | << "been modified since the iterator was constructed"; |
235 | CheckNoMutations(); |
236 | if (array_iter_ != nullptr) { |
237 | return array_iter_ == other.array_iter_; |
238 | } else { |
239 | return other.array_iter_ == nullptr && tree_iter_ == other.tree_iter_; |
240 | } |
241 | } |
242 | |
243 | } // namespace tensorflow |
244 | |
245 | #endif // TENSORFLOW_CORE_GRAPH_EDGESET_H_ |
246 |