1 /* Copyright 2019 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_COMPILER_XLA_SERVICE_GRAPHCYCLES_ORDERED_SET_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GRAPHCYCLES_ORDERED_SET_H_ 18 19 #include <vector> 20 21 #include "absl/container/flat_hash_map.h" 22 #include "absl/types/span.h" 23 #include "tensorflow/core/platform/logging.h" 24 25 namespace tensorflow { 26 // This is a set data structure that provides a deterministic iteration order. 27 // The iteration order of elements only depends on the sequence of 28 // inserts/deletes, so as long as the inserts/deletes happen in the same 29 // sequence, the set will have the same iteration order. 30 // 31 // Assumes that T can be cheaply copied for simplicity. 32 template <typename T> 33 class OrderedSet { 34 public: 35 // Inserts `value` into the ordered set. Returns true if the value was not 36 // present in the set before the insertion. Insert(T value)37 bool Insert(T value) { 38 bool new_insertion = 39 value_to_index_.insert({value, value_sequence_.size()}).second; 40 if (new_insertion) { 41 value_sequence_.push_back(value); 42 } 43 return new_insertion; 44 } 45 46 // Removes `value` from the set. Assumes `value` is already present in the 47 // set. Erase(T value)48 void Erase(T value) { 49 auto it = value_to_index_.find(value); 50 DCHECK(it != value_to_index_.end()); 51 52 // Since we don't want to move values around in `value_sequence_` we swap 53 // the value in the last position and with value to be deleted and then 54 // pop_back. 55 value_to_index_[value_sequence_.back()] = it->second; 56 std::swap(value_sequence_[it->second], value_sequence_.back()); 57 value_sequence_.pop_back(); 58 value_to_index_.erase(it); 59 } 60 Reserve(size_t new_size)61 void Reserve(size_t new_size) { 62 value_to_index_.reserve(new_size); 63 value_sequence_.reserve(new_size); 64 } 65 Clear()66 void Clear() { 67 value_to_index_.clear(); 68 value_sequence_.clear(); 69 } 70 Contains(T value)71 bool Contains(T value) const { return value_to_index_.contains(value); } Size()72 size_t Size() const { return value_sequence_.size(); } 73 GetSequence()74 absl::Span<T const> GetSequence() const { return value_sequence_; } 75 76 private: 77 // The stable order that we maintain through insertions and deletions. 78 std::vector<T> value_sequence_; 79 80 // Maps values to their indices in `value_sequence_`. 81 absl::flat_hash_map<T, int> value_to_index_; 82 }; 83 } // namespace tensorflow 84 85 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GRAPHCYCLES_ORDERED_SET_H_ 86