1 // Copyright 2018 The Chromium Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 #ifndef BASE_TASK_SEQUENCE_MANAGER_INTRUSIVE_HEAP_H_ 6 #define BASE_TASK_SEQUENCE_MANAGER_INTRUSIVE_HEAP_H_ 7 8 #include <algorithm> 9 #include <vector> 10 11 #include "base/logging.h" 12 13 namespace base { 14 namespace sequence_manager { 15 namespace internal { 16 17 template <typename T> 18 class IntrusiveHeap; 19 20 // Intended as an opaque wrapper around |index_|. 21 class HeapHandle { 22 public: HeapHandle()23 HeapHandle() : index_(0u) {} 24 IsValid()25 bool IsValid() const { return index_ != 0u; } 26 27 private: 28 template <typename T> 29 friend class IntrusiveHeap; 30 HeapHandle(size_t index)31 HeapHandle(size_t index) : index_(index) {} 32 33 size_t index_; 34 }; 35 36 // A standard min-heap with the following assumptions: 37 // 1. T has operator <= 38 // 2. T has method void SetHeapHandle(HeapHandle handle) 39 // 3. T has method void ClearHeapHandle() 40 // 4. T is moveable 41 // 5. T is default constructible 42 // 6. The heap size never gets terribly big so reclaiming memory on pop/erase 43 // isn't a priority. 44 // 45 // The reason IntrusiveHeap exists is to provide similar performance to 46 // std::priority_queue while allowing removal of arbitrary elements. 47 template <typename T> 48 class IntrusiveHeap { 49 public: IntrusiveHeap()50 IntrusiveHeap() : nodes_(kMinimumHeapSize), size_(0) {} 51 ~IntrusiveHeap()52 ~IntrusiveHeap() { 53 for (size_t i = 1; i <= size_; i++) { 54 MakeHole(i); 55 } 56 } 57 empty()58 bool empty() const { return size_ == 0; } 59 size()60 size_t size() const { return size_; } 61 Clear()62 void Clear() { 63 for (size_t i = 1; i <= size_; i++) { 64 MakeHole(i); 65 } 66 nodes_.resize(kMinimumHeapSize); 67 size_ = 0; 68 } 69 Min()70 const T& Min() const { 71 DCHECK_GE(size_, 1u); 72 return nodes_[1]; 73 } 74 Pop()75 void Pop() { 76 DCHECK_GE(size_, 1u); 77 MakeHole(1u); 78 size_t top_index = size_--; 79 if (!empty()) 80 MoveHoleDownAndFillWithLeafElement(1u, std::move(nodes_[top_index])); 81 } 82 insert(T && element)83 void insert(T&& element) { 84 size_++; 85 if (size_ >= nodes_.size()) 86 nodes_.resize(nodes_.size() * 2); 87 // Notionally we have a hole in the tree at index |size_|, move this up 88 // to find the right insertion point. 89 MoveHoleUpAndFillWithElement(size_, std::move(element)); 90 } 91 erase(HeapHandle handle)92 void erase(HeapHandle handle) { 93 DCHECK_GT(handle.index_, 0u); 94 DCHECK_LE(handle.index_, size_); 95 MakeHole(handle.index_); 96 size_t top_index = size_--; 97 if (empty() || top_index == handle.index_) 98 return; 99 if (nodes_[handle.index_] <= nodes_[top_index]) { 100 MoveHoleDownAndFillWithLeafElement(handle.index_, 101 std::move(nodes_[top_index])); 102 } else { 103 MoveHoleUpAndFillWithElement(handle.index_, std::move(nodes_[top_index])); 104 } 105 } 106 ReplaceMin(T && element)107 void ReplaceMin(T&& element) { 108 // Note |element| might not be a leaf node so we can't use 109 // MoveHoleDownAndFillWithLeafElement. 110 MoveHoleDownAndFillWithElement(1u, std::move(element)); 111 } 112 ChangeKey(HeapHandle handle,T && element)113 void ChangeKey(HeapHandle handle, T&& element) { 114 if (nodes_[handle.index_] <= element) { 115 MoveHoleDownAndFillWithLeafElement(handle.index_, std::move(element)); 116 } else { 117 MoveHoleUpAndFillWithElement(handle.index_, std::move(element)); 118 } 119 } 120 121 // Caution mutating the heap invalidates the iterators. begin()122 const T* begin() const { return &nodes_[1u]; } end()123 const T* end() const { return begin() + size_; } 124 125 private: 126 enum { 127 // The majority of sets in the scheduler have 0-3 items in them (a few will 128 // have perhaps up to 100), so this means we usually only have to allocate 129 // memory once. 130 kMinimumHeapSize = 4u 131 }; 132 133 friend class IntrusiveHeapTest; 134 MoveHole(size_t new_hole_pos,size_t old_hole_pos)135 size_t MoveHole(size_t new_hole_pos, size_t old_hole_pos) { 136 DCHECK_GT(new_hole_pos, 0u); 137 DCHECK_LE(new_hole_pos, size_); 138 DCHECK_GT(new_hole_pos, 0u); 139 DCHECK_LE(new_hole_pos, size_); 140 DCHECK_NE(old_hole_pos, new_hole_pos); 141 nodes_[old_hole_pos] = std::move(nodes_[new_hole_pos]); 142 nodes_[old_hole_pos].SetHeapHandle(HeapHandle(old_hole_pos)); 143 return new_hole_pos; 144 } 145 146 // Notionally creates a hole in the tree at |index|. MakeHole(size_t index)147 void MakeHole(size_t index) { 148 DCHECK_GT(index, 0u); 149 DCHECK_LE(index, size_); 150 nodes_[index].ClearHeapHandle(); 151 } 152 FillHole(size_t hole,T && element)153 void FillHole(size_t hole, T&& element) { 154 DCHECK_GT(hole, 0u); 155 DCHECK_LE(hole, size_); 156 nodes_[hole] = std::move(element); 157 nodes_[hole].SetHeapHandle(HeapHandle(hole)); 158 DCHECK(std::is_heap(begin(), end(), CompareNodes)); 159 } 160 161 // is_heap requires a strict comparator. CompareNodes(const T & a,const T & b)162 static bool CompareNodes(const T& a, const T& b) { return !(a <= b); } 163 164 // Moves the |hole| up the tree and when the right position has been found 165 // |element| is moved in. MoveHoleUpAndFillWithElement(size_t hole,T && element)166 void MoveHoleUpAndFillWithElement(size_t hole, T&& element) { 167 DCHECK_GT(hole, 0u); 168 DCHECK_LE(hole, size_); 169 while (hole >= 2u) { 170 size_t parent_pos = hole / 2; 171 if (nodes_[parent_pos] <= element) 172 break; 173 174 hole = MoveHole(parent_pos, hole); 175 } 176 FillHole(hole, std::move(element)); 177 } 178 179 // Moves the |hole| down the tree and when the right position has been found 180 // |element| is moved in. MoveHoleDownAndFillWithElement(size_t hole,T && element)181 void MoveHoleDownAndFillWithElement(size_t hole, T&& element) { 182 DCHECK_GT(hole, 0u); 183 DCHECK_LE(hole, size_); 184 size_t child_pos = hole * 2; 185 while (child_pos < size_) { 186 if (nodes_[child_pos + 1] <= nodes_[child_pos]) 187 child_pos++; 188 189 if (element <= nodes_[child_pos]) 190 break; 191 192 hole = MoveHole(child_pos, hole); 193 child_pos *= 2; 194 } 195 if (child_pos == size_ && !(element <= nodes_[child_pos])) 196 hole = MoveHole(child_pos, hole); 197 FillHole(hole, std::move(element)); 198 } 199 200 // Moves the |hole| down the tree and when the right position has been found 201 // |leaf_element| is moved in. Faster than MoveHoleDownAndFillWithElement 202 // (it does one key comparison per level instead of two) but only valid for 203 // leaf elements (i.e. one of the max values). MoveHoleDownAndFillWithLeafElement(size_t hole,T && leaf_element)204 void MoveHoleDownAndFillWithLeafElement(size_t hole, T&& leaf_element) { 205 DCHECK_GT(hole, 0u); 206 DCHECK_LE(hole, size_); 207 size_t child_pos = hole * 2; 208 while (child_pos < size_) { 209 size_t second_child = child_pos + 1; 210 if (nodes_[second_child] <= nodes_[child_pos]) 211 child_pos = second_child; 212 213 hole = MoveHole(child_pos, hole); 214 child_pos *= 2; 215 } 216 if (child_pos == size_) 217 hole = MoveHole(child_pos, hole); 218 MoveHoleUpAndFillWithElement(hole, std::move(leaf_element)); 219 } 220 221 std::vector<T> nodes_; // NOTE we use 1-based indexing 222 size_t size_; 223 }; 224 225 } // namespace internal 226 } // namespace sequence_manager 227 } // namespace base 228 229 #endif // BASE_TASK_SEQUENCE_MANAGER_INTRUSIVE_HEAP_H_ 230