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