1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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"
25namespace tensorflow {
26
27class 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.
32class 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
92class 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
133inline EdgeSet::EdgeSet() {
134 for (int i = 0; i < kInline; i++) {
135 ptrs_[i] = nullptr;
136 }
137}
138
139inline EdgeSet::~EdgeSet() { delete get_set(); }
140
141inline bool EdgeSet::empty() const { return size() == 0; }
142
143inline 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
156inline void EdgeSet::clear() {
157 RegisterMutation();
158 delete get_set();
159 for (int i = 0; i < kInline; i++) {
160 ptrs_[i] = nullptr;
161 }
162}
163
164inline 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
176inline 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
188inline 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
198inline 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.
208inline const EdgeSet::const_iterator::value_type* EdgeSet::const_iterator::
209operator->() 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.
220inline 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
230inline 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