xref: /aosp_15_r20/external/cronet/base/containers/intrusive_heap.h (revision 6777b5387eb2ff775bb5750e3f5d96f37fb7352b)
1 // Copyright 2019 The Chromium Authors
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_CONTAINERS_INTRUSIVE_HEAP_H_
6 #define BASE_CONTAINERS_INTRUSIVE_HEAP_H_
7 
8 // Implements a standard max-heap, but with arbitrary element removal. To
9 // facilitate this, each element has associated with it a HeapHandle (an opaque
10 // wrapper around the index at which the element is stored), which is maintained
11 // by the heap as elements move within it.
12 //
13 // An IntrusiveHeap is implemented as a standard max-heap over a std::vector<T>,
14 // like std::make_heap. Insertion, removal and updating are amortized O(lg size)
15 // (occasional O(size) cost if a new vector allocation is required). Retrieving
16 // an element by handle is O(1). Looking up the top element is O(1). Insertions,
17 // removals and updates invalidate all iterators, but handles remain valid.
18 // Similar to a std::set, all iterators are read-only so as to disallow changing
19 // elements and violating the heap property. That being said, if the type you
20 // are storing is able to have its sort key be changed externally you can
21 // repair the heap by resorting the modified element via a call to "Update".
22 //
23 // Example usage:
24 //
25 //   // Create a heap, wrapping integer elements with WithHeapHandle in order to
26 //   // endow them with heap handles.
27 //   IntrusiveHeap<WithHeapHandle<int>> heap;
28 //
29 //   // WithHeapHandle<T> is for simple or opaque types. In cases where you
30 //   // control the type declaration you can also provide HeapHandle storage by
31 //   // deriving from InternalHeapHandleStorage.
32 //   class Foo : public InternalHeapHandleStorage {
33 //    public:
34 //     explicit Foo(int);
35 //     ...
36 //   };
37 //   IntrusiveHeap<Foo> heap2;
38 //
39 //   // Insert some elements. Like most containers, "insert" returns an iterator
40 //   // to the element in the container.
41 //   heap.insert(3);
42 //   heap.insert(1);
43 //   auto it = heap.insert(4);
44 //
45 //   // By default this is a max heap, so the top element should be 4 at this
46 //   // point.
47 //   EXPECT_EQ(4, heap.top().value());
48 //
49 //   // Iterators are invalidated by further heap operations, but handles are
50 //   // not. Grab a handle to the current top element so we can track it across
51 //   // changes.
52 //   HeapHandle* handle = it->handle();
53 //
54 //   // Insert a new max element. 4 should no longer be the top.
55 //   heap.insert(5);
56 //   EXPECT_EQ(5, heap.top().value());
57 //
58 //   // We can lookup and erase element 4 by its handle, even though it has
59 //   // moved. Note that erasing the element invalidates the handle to it.
60 //   EXPECT_EQ(4, heap.at(*handle).value());
61 //   heap.erase(*handle);
62 //   handle = nullptr;
63 //
64 //   // Popping the current max (5), makes 3 the new max, as we already erased
65 //   // element 4.
66 //   heap.pop();
67 //   EXPECT_EQ(3, heap.top().value());
68 //
69 // Under the hood the HeapHandle is managed by an object implementing the
70 // HeapHandleAccess interface, which is passed as a parameter to the
71 // IntrusiveHeap template:
72 //
73 //   // Gets the heap handle associated with the element. This should return the
74 //   // most recently set handle value, or HeapHandle::Invalid(). This is only
75 //   // called in DCHECK builds.
76 //   HeapHandle GetHeapHandle(const T*);
77 //
78 //   // Changes the result of GetHeapHandle. GetHeapHandle() must return the
79 //   // most recent value provided to SetHeapHandle() or HeapHandle::Invalid().
80 //   // In some implementations, where GetHeapHandle() can independently
81 //   // reproduce the correct value, it is possible that SetHeapHandle() does
82 //   // nothing.
83 //   void SetHeapHandle(T*, HeapHandle);
84 //
85 //   // Clears the heap handle associated with the given element. After calling
86 //   // this GetHeapHandle() must return HeapHandle::Invalid().
87 //   void ClearHeapHandle(T*);
88 //
89 // The default implementation of HeapHandleAccess assumes that your type
90 // provides HeapHandle storage and will simply forward these calls to equivalent
91 // member functions on the type T:
92 //
93 //   void T::SetHeapHandle(HeapHandle)
94 //   void T::ClearHeapHandle()
95 //   HeapHandle T::GetHeapHandle() const
96 //
97 // The WithHeapHandle and InternalHeapHandleStorage classes in turn provide
98 // implementations of that contract.
99 //
100 // In summary, to provide heap handle support for your type, you can do one of
101 // the following (from most manual / least magical, to least manual / most
102 // magical):
103 //
104 // 0. use a custom HeapHandleAccessor, and implement storage however you want;
105 // 1. use the default HeapHandleAccessor, and manually provide storage on your
106 //    your element type and implement the IntrusiveHeap contract;
107 // 2. use the default HeapHandleAccessor, and endow your type with handle
108 //    storage by deriving from a helper class (see InternalHeapHandleStorage);
109 //    or,
110 // 3. use the default HeapHandleAccessor, and wrap your type in a container that
111 //    provides handle storage (see WithHeapHandle<T>).
112 //
113 // Approach 0 is suitable for custom types that already implement something akin
114 // to heap handles, via back pointers or any other mechanism, but where the
115 // storage is external to the objects in the heap. If you already have the
116 // ability to determine where in a container an object lives despite it
117 // being moved, then you don't need the overhead of storing an actual HeapHandle
118 // whose value can be inferred.
119 //
120 // Approach 1 is is suitable in cases like the above, but where the data
121 // allowing you to determine the index of an element in a container is stored
122 // directly in the object itself.
123 //
124 // Approach 2 is suitable for types whose declarations you control, where you
125 // are able to use inheritance.
126 //
127 // Finally, approach 3 is suitable when you are storing PODs, or a type whose
128 // declaration you can not change.
129 //
130 // Most users should be using approach 2 or 3.
131 
132 #include <algorithm>
133 #include <compare>
134 #include <functional>
135 #include <limits>
136 #include <memory>
137 #include <type_traits>
138 #include <utility>
139 #include <vector>
140 
141 #include "base/base_export.h"
142 #include "base/check.h"
143 #include "base/check_op.h"
144 #include "base/memory/ptr_util.h"
145 #include "base/ranges/algorithm.h"
146 #include "third_party/abseil-cpp/absl/container/inlined_vector.h"
147 
148 namespace base {
149 
150 // Intended as a wrapper around an |index_| in the vector storage backing an
151 // IntrusiveHeap. A HeapHandle is associated with each element in an
152 // IntrusiveHeap, and is maintained by the heap as the object moves around
153 // within it. It can be used to subsequently remove the element, or update it
154 // in place.
155 class BASE_EXPORT HeapHandle {
156  public:
157   enum : size_t { kInvalidIndex = std::numeric_limits<size_t>::max() };
158 
159   constexpr HeapHandle() = default;
160   constexpr HeapHandle(const HeapHandle& other) = default;
HeapHandle(HeapHandle && other)161   HeapHandle(HeapHandle&& other) noexcept
162       : index_(std::exchange(other.index_, kInvalidIndex)) {}
163   ~HeapHandle() = default;
164 
165   HeapHandle& operator=(const HeapHandle& other) = default;
166   HeapHandle& operator=(HeapHandle&& other) noexcept {
167     index_ = std::exchange(other.index_, kInvalidIndex);
168     return *this;
169   }
170 
171   static HeapHandle Invalid();
172 
173   // Resets this handle back to an invalid state.
reset()174   void reset() { index_ = kInvalidIndex; }
175 
176   // Accessors.
index()177   size_t index() const { return index_; }
IsValid()178   bool IsValid() const { return index_ != kInvalidIndex; }
179 
180   // Comparison operators.
181   friend bool operator==(const HeapHandle& lhs,
182                          const HeapHandle& rhs) = default;
183   friend std::strong_ordering operator<=>(const HeapHandle& lhs,
184                                           const HeapHandle& rhs) = default;
185 
186  private:
187   template <typename T, typename Compare, typename HeapHandleAccessor>
188   friend class IntrusiveHeap;
189 
190   // Only IntrusiveHeaps can create valid HeapHandles.
HeapHandle(size_t index)191   explicit HeapHandle(size_t index) : index_(index) {}
192 
193   size_t index_ = kInvalidIndex;
194 };
195 
196 // The default HeapHandleAccessor, which simply forwards calls to the underlying
197 // type.
198 template <typename T>
199 struct DefaultHeapHandleAccessor {
SetHeapHandleDefaultHeapHandleAccessor200   void SetHeapHandle(T* element, HeapHandle handle) const {
201     element->SetHeapHandle(handle);
202   }
203 
ClearHeapHandleDefaultHeapHandleAccessor204   void ClearHeapHandle(T* element) const { element->ClearHeapHandle(); }
205 
GetHeapHandleDefaultHeapHandleAccessor206   HeapHandle GetHeapHandle(const T* element) const {
207     return element->GetHeapHandle();
208   }
209 };
210 
211 // Intrusive heap class. This is something like a std::vector (insertion and
212 // removal are similar, objects don't have a fixed address in memory) crossed
213 // with a std::set (elements are considered immutable once they're in the
214 // container).
215 template <typename T,
216           typename Compare = std::less<T>,
217           typename HeapHandleAccessor = DefaultHeapHandleAccessor<T>>
218 class IntrusiveHeap {
219  private:
220   using UnderlyingType = std::vector<T>;
221 
222  public:
223   //////////////////////////////////////////////////////////////////////////////
224   // Types.
225 
226   using value_type = typename UnderlyingType::value_type;
227   using size_type = typename UnderlyingType::size_type;
228   using difference_type = typename UnderlyingType::difference_type;
229   using value_compare = Compare;
230   using heap_handle_accessor = HeapHandleAccessor;
231 
232   using reference = typename UnderlyingType::reference;
233   using const_reference = typename UnderlyingType::const_reference;
234   using pointer = typename UnderlyingType::pointer;
235   using const_pointer = typename UnderlyingType::const_pointer;
236 
237   // Iterators are read-only.
238   using iterator = typename UnderlyingType::const_iterator;
239   using const_iterator = typename UnderlyingType::const_iterator;
240   using reverse_iterator = typename UnderlyingType::const_reverse_iterator;
241   using const_reverse_iterator =
242       typename UnderlyingType::const_reverse_iterator;
243 
244   //////////////////////////////////////////////////////////////////////////////
245   // Lifetime.
246 
247   IntrusiveHeap() = default;
IntrusiveHeap(const value_compare & comp,const heap_handle_accessor & access)248   IntrusiveHeap(const value_compare& comp, const heap_handle_accessor& access)
249       : impl_(comp, access) {}
250 
251   template <class InputIterator>
252   IntrusiveHeap(InputIterator first,
253                 InputIterator last,
254                 const value_compare& comp = value_compare(),
255                 const heap_handle_accessor& access = heap_handle_accessor())
impl_(comp,access)256       : impl_(comp, access) {
257     insert(first, last);
258   }
259 
260   // Moves an intrusive heap. The outstanding handles remain valid and end up
261   // pointing to the new heap.
262   IntrusiveHeap(IntrusiveHeap&& other) = default;
263 
264   // Copy constructor for an intrusive heap.
265   IntrusiveHeap(const IntrusiveHeap&);
266 
267   // Initializer list constructor.
268   template <typename U>
269   IntrusiveHeap(std::initializer_list<U> ilist,
270                 const value_compare& comp = value_compare(),
271                 const heap_handle_accessor& access = heap_handle_accessor())
impl_(comp,access)272       : impl_(comp, access) {
273     insert(std::begin(ilist), std::end(ilist));
274   }
275 
276   ~IntrusiveHeap();
277 
278   //////////////////////////////////////////////////////////////////////////////
279   // Assignment.
280 
281   IntrusiveHeap& operator=(IntrusiveHeap&&) noexcept;
282   IntrusiveHeap& operator=(const IntrusiveHeap&);
283   IntrusiveHeap& operator=(std::initializer_list<value_type> ilist);
284 
285   //////////////////////////////////////////////////////////////////////////////
286   // Element access.
287   //
288   // These provide O(1) const access to the elements in the heap. If you wish to
289   // modify an element in the heap you should first remove it from the heap, and
290   // then reinsert it into the heap, or use the "Replace*" helper functions. In
291   // the rare case where you directly modify an element in the heap you can
292   // subsequently repair the heap with "Update".
293 
at(size_type pos)294   const_reference at(size_type pos) const { return impl_.heap_.at(pos); }
at(HeapHandle pos)295   const_reference at(HeapHandle pos) const {
296     return impl_.heap_.at(pos.index());
297   }
298   const_reference operator[](size_type pos) const { return impl_.heap_[pos]; }
299   const_reference operator[](HeapHandle pos) const {
300     return impl_.heap_[pos.index()];
301   }
front()302   const_reference front() const { return impl_.heap_.front(); }
back()303   const_reference back() const { return impl_.heap_.back(); }
top()304   const_reference top() const { return impl_.heap_.front(); }
305 
306   // May or may not return a null pointer if size() is zero.
data()307   const_pointer data() const { return impl_.heap_.data(); }
308 
309   //////////////////////////////////////////////////////////////////////////////
310   // Memory management.
311 
reserve(size_type new_capacity)312   void reserve(size_type new_capacity) { impl_.heap_.reserve(new_capacity); }
capacity()313   size_type capacity() const { return impl_.heap_.capacity(); }
shrink_to_fit()314   void shrink_to_fit() { impl_.heap_.shrink_to_fit(); }
315 
316   //////////////////////////////////////////////////////////////////////////////
317   // Size management.
318 
319   void clear();
size()320   size_type size() const { return impl_.heap_.size(); }
max_size()321   size_type max_size() const { return impl_.heap_.max_size(); }
empty()322   bool empty() const { return impl_.heap_.empty(); }
323 
324   //////////////////////////////////////////////////////////////////////////////
325   // Iterators.
326   //
327   // Only constant iterators are allowed.
328 
begin()329   const_iterator begin() const { return impl_.heap_.cbegin(); }
cbegin()330   const_iterator cbegin() const { return impl_.heap_.cbegin(); }
331 
end()332   const_iterator end() const { return impl_.heap_.cend(); }
cend()333   const_iterator cend() const { return impl_.heap_.cend(); }
334 
rbegin()335   const_reverse_iterator rbegin() const { return impl_.heap_.crbegin(); }
crbegin()336   const_reverse_iterator crbegin() const { return impl_.heap_.crbegin(); }
337 
rend()338   const_reverse_iterator rend() const { return impl_.heap_.crend(); }
crend()339   const_reverse_iterator crend() const { return impl_.heap_.crend(); }
340 
341   //////////////////////////////////////////////////////////////////////////////
342   // Insertion (these are std::multiset like, with no position hints).
343   //
344   // All insertion operations invalidate iterators, pointers and references.
345   // Handles remain valid. Insertion of one element is amortized O(lg size)
346   // (occasional O(size) cost if a new vector allocation is required).
347 
insert(const value_type & value)348   const_iterator insert(const value_type& value) { return InsertImpl(value); }
insert(value_type && value)349   const_iterator insert(value_type&& value) {
350     return InsertImpl(std::move_if_noexcept(value));
351   }
352 
353   template <class InputIterator>
354   void insert(InputIterator first, InputIterator last);
355 
356   template <typename... Args>
357   const_iterator emplace(Args&&... args);
358 
359   //////////////////////////////////////////////////////////////////////////////
360   // Removing elements.
361   //
362   // Erasing invalidates all outstanding iterators, pointers and references.
363   // Handles remain valid. Removing one element is amortized O(lg size)
364   // (occasional O(size) cost if a new vector allocation is required).
365   //
366   // Note that it is safe for the element being removed to be in an invalid
367   // state (modified such that it may currently violate the heap property)
368   // when this called.
369 
370   // Takes the element from the heap at the given position, erasing that entry
371   // from the heap. This can only be called if |value_type| is movable.
372   value_type take(size_type pos);
373 
374   // Version of take that will accept iterators and handles. This can only be
375   // called if |value_type| is movable.
376   template <typename P>
take(P pos)377   value_type take(P pos) {
378     return take(ToIndex(pos));
379   }
380 
381   // Takes the top element from the heap.
take_top()382   value_type take_top() { return take(0u); }
383 
384   // Erases the element at the given position |pos|.
385   void erase(size_type pos);
386 
387   // Version of erase that will accept iterators and handles.
388   template <typename P>
erase(P pos)389   void erase(P pos) {
390     erase(ToIndex(pos));
391   }
392 
393   // Removes the element at the top of the heap (accessible via "top", or
394   // "front" or "take").
pop()395   void pop() { erase(0u); }
396 
397   // Erases every element that matches the predicate. This is done in-place for
398   // maximum efficiency. Also, to avoid re-entrancy issues, elements are deleted
399   // at the very end.
400   // Note: This function is currently tuned for a use-case where there are
401   // usually 8 or less elements removed at a time. Consider adding a template
402   // parameter if a different tuning is needed.
403   template <typename Functor>
EraseIf(Functor predicate)404   void EraseIf(Functor predicate) {
405     // Stable partition ensures that if no elements are erased, the heap remains
406     // intact.
407     auto erase_start = std::stable_partition(
408         impl_.heap_.begin(), impl_.heap_.end(),
409         [&](const auto& element) { return !predicate(element); });
410 
411     // Clear the heap handle of every element that will be erased.
412     for (size_t i = static_cast<size_t>(erase_start - impl_.heap_.begin());
413          i < impl_.heap_.size(); ++i) {
414       ClearHeapHandle(i);
415     }
416 
417     // Deleting an element can potentially lead to reentrancy, we move all the
418     // elements to be erased into a temporary container before deleting them.
419     // This is to avoid changing the underlying container during the erase()
420     // call.
421     absl::InlinedVector<value_type, 8> elements_to_delete;
422     std::move(erase_start, impl_.heap_.end(),
423               std::back_inserter(elements_to_delete));
424 
425     impl_.heap_.erase(erase_start, impl_.heap_.end());
426 
427     // If no elements were removed, then the heap is still intact.
428     if (elements_to_delete.empty()) {
429       return;
430     }
431 
432     // Repair the heap and ensure handles are pointing to the right index.
433     ranges::make_heap(impl_.heap_, value_comp());
434     for (size_t i = 0; i < size(); ++i)
435       SetHeapHandle(i);
436 
437     // Explicitly delete elements last.
438     elements_to_delete.clear();
439   }
440 
441   //////////////////////////////////////////////////////////////////////////////
442   // Updating.
443   //
444   // Amortized cost of O(lg size).
445 
446   // Replaces the element corresponding to |handle| with a new |element|.
Replace(size_type pos,const T & element)447   const_iterator Replace(size_type pos, const T& element) {
448     return ReplaceImpl(pos, element);
449   }
Replace(size_type pos,T && element)450   const_iterator Replace(size_type pos, T&& element) {
451     return ReplaceImpl(pos, std::move_if_noexcept(element));
452   }
453 
454   // Versions of Replace that will accept handles and iterators.
455   template <typename P>
Replace(P pos,const T & element)456   const_iterator Replace(P pos, const T& element) {
457     return ReplaceImpl(ToIndex(pos), element);
458   }
459   template <typename P>
Replace(P pos,T && element)460   const_iterator Replace(P pos, T&& element) {
461     return ReplaceImpl(ToIndex(pos), std::move_if_noexcept(element));
462   }
463 
464   // Replaces the top element in the heap with the provided element.
ReplaceTop(const T & element)465   const_iterator ReplaceTop(const T& element) {
466     return ReplaceTopImpl(element);
467   }
ReplaceTop(T && element)468   const_iterator ReplaceTop(T&& element) {
469     return ReplaceTopImpl(std::move_if_noexcept(element));
470   }
471 
472   // Causes the object at the given location to be resorted into an appropriate
473   // position in the heap. To be used if the object in the heap was externally
474   // modified, and the heap needs to be repaired. This only works if a single
475   // heap element has been modified, otherwise the behaviour is undefined.
476   const_iterator Update(size_type pos);
477   template <typename P>
Update(P pos)478   const_iterator Update(P pos) {
479     return Update(ToIndex(pos));
480   }
481 
482   // Applies a modification function to the object at the given location, then
483   // repairs the heap. To be used to modify an element in the heap in-place
484   // while keeping the heap intact.
485   template <typename P, typename UnaryOperation>
Modify(P pos,UnaryOperation unary_op)486   const_iterator Modify(P pos, UnaryOperation unary_op) {
487     size_type index = ToIndex(pos);
488     unary_op(impl_.heap_.at(index));
489     return Update(index);
490   }
491 
492   //////////////////////////////////////////////////////////////////////////////
493   // Access to helper functors.
494 
value_comp()495   const value_compare& value_comp() const { return impl_.get_value_compare(); }
496 
heap_handle_access()497   const heap_handle_accessor& heap_handle_access() const {
498     return impl_.get_heap_handle_access();
499   }
500 
501   //////////////////////////////////////////////////////////////////////////////
502   // General operations.
503 
504   void swap(IntrusiveHeap& other) noexcept;
swap(IntrusiveHeap & lhs,IntrusiveHeap & rhs)505   friend void swap(IntrusiveHeap& lhs, IntrusiveHeap& rhs) { lhs.swap(rhs); }
506 
507   // Comparison operators. These check for exact equality. Two heaps that are
508   // semantically equivalent (contain the same elements, but in different
509   // orders) won't compare as equal using these operators.
510   friend bool operator==(const IntrusiveHeap& lhs, const IntrusiveHeap& rhs) {
511     return lhs.impl_.heap_ == rhs.impl_.heap_;
512   }
513 
514   //////////////////////////////////////////////////////////////////////////////
515   // Utility functions.
516 
517   // Converts iterators and handles to indices. Helpers for templated versions
518   // of insert/erase/Replace.
ToIndex(HeapHandle handle)519   size_type ToIndex(HeapHandle handle) { return handle.index(); }
520   size_type ToIndex(const_iterator pos);
521   size_type ToIndex(const_reverse_iterator pos);
522 
523  private:
524   // Templated version of ToIndex that lets insert/erase/Replace work with all
525   // integral types.
526   template <typename I, typename = std::enable_if_t<std::is_integral_v<I>>>
ToIndex(I pos)527   size_type ToIndex(I pos) {
528     return static_cast<size_type>(pos);
529   }
530 
531   // Returns the last valid index in |heap_|.
GetLastIndex()532   size_type GetLastIndex() const { return impl_.heap_.size() - 1; }
533 
534   // Helper functions for setting heap handles.
535   void SetHeapHandle(size_type i);
536   void ClearHeapHandle(size_type i);
537   HeapHandle GetHeapHandle(size_type i);
538 
539   // Helpers for doing comparisons between elements inside and outside of the
540   // heap.
541   bool Less(size_type i, size_type j);
542   bool Less(const T& element, size_type i);
543   bool Less(size_type i, const T& element);
544 
545   // The following function are all related to the basic heap algorithm
546   // underpinning this data structure. They are templated so that they work with
547   // both movable (U = T&&) and non-movable (U = const T&) types.
548 
549   // Primitive helpers for adding removing / elements to the heap. To minimize
550   // moves, the heap is implemented by making a hole where an element used to
551   // be (or where a new element will soon be), and moving the hole around,
552   // before finally filling the hole or deleting the entry corresponding to the
553   // hole.
554   void MakeHole(size_type pos);
555   template <typename U>
556   void FillHole(size_type hole, U element);
557   void MoveHole(size_type new_hole_pos, size_type old_hole_pos);
558 
559   // Moves a hold up the tree and fills it with the provided |element|. Returns
560   // the final index of the element.
561   template <typename U>
562   size_type MoveHoleUpAndFill(size_type hole_pos, U element);
563 
564   // Moves a hole down the tree and fills it with the provided |element|. If
565   // |kFillWithLeaf| is true it will deterministically move the hole all the
566   // way down the tree, avoiding a second comparison per level, before
567   // potentially moving it back up the tree.
568   struct WithLeafElement {
569     static constexpr bool kIsLeafElement = true;
570   };
571   struct WithElement {
572     static constexpr bool kIsLeafElement = false;
573   };
574   template <typename FillElementType, typename U>
575   size_type MoveHoleDownAndFill(size_type hole_pos, U element);
576 
577   // Implementation of Insert and Replace built on top of the MoveHole
578   // primitives.
579   template <typename U>
580   const_iterator InsertImpl(U element);
581   template <typename U>
582   const_iterator ReplaceImpl(size_type pos, U element);
583   template <typename U>
584   const_iterator ReplaceTopImpl(U element);
585 
586   // To support comparators that may not be possible to default-construct, we
587   // have to store an instance of value_compare. Using this to store all
588   // internal state of IntrusiveHeap and using private inheritance to store
589   // compare lets us take advantage of an empty base class optimization to avoid
590   // extra space in the common case when Compare has no state.
591   struct Impl : private value_compare, private heap_handle_accessor {
ImplImpl592     Impl(const value_compare& value_comp,
593          const heap_handle_accessor& heap_handle_access)
594         : value_compare(value_comp), heap_handle_accessor(heap_handle_access) {}
595 
596     Impl() = default;
597     Impl(Impl&&) = default;
598     Impl(const Impl&) = default;
599     Impl& operator=(Impl&& other) = default;
600     Impl& operator=(const Impl& other) = default;
601 
get_value_compareImpl602     const value_compare& get_value_compare() const { return *this; }
get_value_compareImpl603     value_compare& get_value_compare() { return *this; }
604 
get_heap_handle_accessImpl605     const heap_handle_accessor& get_heap_handle_access() const { return *this; }
get_heap_handle_accessImpl606     heap_handle_accessor& get_heap_handle_access() { return *this; }
607 
608     // The items in the heap.
609     UnderlyingType heap_;
610   } impl_;
611 };
612 
613 // Helper class to endow an object with internal HeapHandle storage. By deriving
614 // from this type you endow your class with self-owned storage for a HeapHandle.
615 // This is a move-only type so that the handle follows the element across moves
616 // and resizes of the underlying vector.
617 class BASE_EXPORT InternalHeapHandleStorage {
618  public:
619   InternalHeapHandleStorage();
620   InternalHeapHandleStorage(const InternalHeapHandleStorage&) = delete;
621   InternalHeapHandleStorage(InternalHeapHandleStorage&& other) noexcept;
622   virtual ~InternalHeapHandleStorage();
623 
624   InternalHeapHandleStorage& operator=(const InternalHeapHandleStorage&) =
625       delete;
626   InternalHeapHandleStorage& operator=(
627       InternalHeapHandleStorage&& other) noexcept;
628 
629   // Allows external clients to get a pointer to the heap handle. This allows
630   // them to remove the element from the heap regardless of its location.
handle()631   HeapHandle* handle() const { return handle_.get(); }
632 
633   // Implementation of IntrusiveHeap contract. Inlined to keep heap code as fast
634   // as possible.
SetHeapHandle(HeapHandle handle)635   void SetHeapHandle(HeapHandle handle) {
636     DCHECK(handle.IsValid());
637     if (handle_)
638       *handle_ = handle;
639   }
ClearHeapHandle()640   void ClearHeapHandle() {
641     if (handle_)
642       handle_->reset();
643   }
GetHeapHandle()644   HeapHandle GetHeapHandle() const {
645     if (handle_)
646       return *handle_;
647     return HeapHandle::Invalid();
648   }
649 
650   // Utility functions.
651   void swap(InternalHeapHandleStorage& other) noexcept;
swap(InternalHeapHandleStorage & lhs,InternalHeapHandleStorage & rhs)652   friend void swap(InternalHeapHandleStorage& lhs,
653                    InternalHeapHandleStorage& rhs) {
654     lhs.swap(rhs);
655   }
656 
657  private:
658   std::unique_ptr<HeapHandle> handle_;
659 };
660 
661 // Spiritually akin to a std::pair<T, std::unique_ptr<HeapHandle>>. Can be used
662 // to wrap arbitrary types and provide them with a HeapHandle, making them
663 // appropriate for use in an IntrusiveHeap. This is a move-only type.
664 template <typename T>
665 class WithHeapHandle : public InternalHeapHandleStorage {
666  public:
667   WithHeapHandle() = default;
668   // Allow implicit conversion of any type that T supports for ease of use with
669   // InstrusiveHeap constructors/insert/emplace.
670   template <typename U>
WithHeapHandle(U value)671   WithHeapHandle(U value) : value_(std::move_if_noexcept(value)) {}
WithHeapHandle(T && value)672   WithHeapHandle(T&& value) noexcept : value_(std::move(value)) {}
673   // Constructor that forwards all arguments along to |value_|.
674   template <class... Args>
675   explicit WithHeapHandle(Args&&... args);
676   WithHeapHandle(const WithHeapHandle&) = delete;
677   WithHeapHandle(WithHeapHandle&& other) noexcept = default;
678   ~WithHeapHandle() override = default;
679 
680   WithHeapHandle& operator=(const WithHeapHandle&) = delete;
681   WithHeapHandle& operator=(WithHeapHandle&& other) = default;
682 
value()683   T& value() { return value_; }
value()684   const T& value() const { return value_; }
685 
686   // Utility functions.
687   void swap(WithHeapHandle& other) noexcept;
swap(WithHeapHandle & lhs,WithHeapHandle & rhs)688   friend void swap(WithHeapHandle& lhs, WithHeapHandle& rhs) { lhs.swap(rhs); }
689 
690   // Comparison operators, for compatibility with ordered STL containers.
691   friend bool operator==(const WithHeapHandle& lhs, const WithHeapHandle& rhs) {
692     return lhs.value_ == rhs.value_;
693   }
694   friend auto operator<=>(const WithHeapHandle& lhs,
695                           const WithHeapHandle& rhs) {
696     return lhs.value_ <=> rhs.value_;
697   }
698 
699  private:
700   T value_;
701 };
702 
703 ////////////////////////////////////////////////////////////////////////////////
704 // IMPLEMENTATION DETAILS
705 
706 namespace intrusive_heap {
707 
ParentIndex(size_t i)708 BASE_EXPORT inline size_t ParentIndex(size_t i) {
709   DCHECK_NE(0u, i);
710   return (i - 1) / 2;
711 }
712 
LeftIndex(size_t i)713 BASE_EXPORT inline size_t LeftIndex(size_t i) {
714   return 2 * i + 1;
715 }
716 
717 template <typename HandleType>
IsInvalid(const HandleType & handle)718 bool IsInvalid(const HandleType& handle) {
719   return !handle || !handle->IsValid();
720 }
721 
CheckInvalidOrEqualTo(HeapHandle handle,size_t index)722 BASE_EXPORT inline void CheckInvalidOrEqualTo(HeapHandle handle, size_t index) {
723   if (handle.IsValid())
724     DCHECK_EQ(index, handle.index());
725 }
726 
727 }  // namespace intrusive_heap
728 
729 ////////////////////////////////////////////////////////////////////////////////
730 // IntrusiveHeap
731 
732 template <typename T, typename Compare, typename HeapHandleAccessor>
IntrusiveHeap(const IntrusiveHeap & other)733 IntrusiveHeap<T, Compare, HeapHandleAccessor>::IntrusiveHeap(
734     const IntrusiveHeap& other)
735     : impl_(other.impl_) {
736   for (size_t i = 0; i < size(); ++i) {
737     SetHeapHandle(i);
738   }
739 }
740 
741 template <typename T, typename Compare, typename HeapHandleAccessor>
~IntrusiveHeap()742 IntrusiveHeap<T, Compare, HeapHandleAccessor>::~IntrusiveHeap() {
743   clear();
744 }
745 
746 template <typename T, typename Compare, typename HeapHandleAccessor>
747 IntrusiveHeap<T, Compare, HeapHandleAccessor>&
748 IntrusiveHeap<T, Compare, HeapHandleAccessor>::operator=(
749     IntrusiveHeap&& other) noexcept {
750   clear();
751   impl_ = std::move(other.impl_);
752   return *this;
753 }
754 
755 template <typename T, typename Compare, typename HeapHandleAccessor>
756 IntrusiveHeap<T, Compare, HeapHandleAccessor>&
757 IntrusiveHeap<T, Compare, HeapHandleAccessor>::operator=(
758     const IntrusiveHeap& other) {
759   clear();
760   impl_ = other.impl_;
761   for (size_t i = 0; i < size(); ++i) {
762     SetHeapHandle(i);
763   }
764   return *this;
765 }
766 
767 template <typename T, typename Compare, typename HeapHandleAccessor>
768 IntrusiveHeap<T, Compare, HeapHandleAccessor>&
769 IntrusiveHeap<T, Compare, HeapHandleAccessor>::operator=(
770     std::initializer_list<value_type> ilist) {
771   clear();
772   insert(std::begin(ilist), std::end(ilist));
773 }
774 
775 template <typename T, typename Compare, typename HeapHandleAccessor>
clear()776 void IntrusiveHeap<T, Compare, HeapHandleAccessor>::clear() {
777   // Make all of the handles invalid before cleaning up the heap.
778   for (size_type i = 0; i < size(); ++i) {
779     ClearHeapHandle(i);
780   }
781 
782   // Clear the heap.
783   impl_.heap_.clear();
784 }
785 
786 template <typename T, typename Compare, typename HeapHandleAccessor>
787 template <class InputIterator>
insert(InputIterator first,InputIterator last)788 void IntrusiveHeap<T, Compare, HeapHandleAccessor>::insert(InputIterator first,
789                                                            InputIterator last) {
790   for (auto it = first; it != last; ++it) {
791     insert(value_type(*it));
792   }
793 }
794 
795 template <typename T, typename Compare, typename HeapHandleAccessor>
796 template <typename... Args>
797 typename IntrusiveHeap<T, Compare, HeapHandleAccessor>::const_iterator
emplace(Args &&...args)798 IntrusiveHeap<T, Compare, HeapHandleAccessor>::emplace(Args&&... args) {
799   value_type value(std::forward<Args>(args)...);
800   return InsertImpl(std::move_if_noexcept(value));
801 }
802 
803 template <typename T, typename Compare, typename HeapHandleAccessor>
804 typename IntrusiveHeap<T, Compare, HeapHandleAccessor>::value_type
take(size_type pos)805 IntrusiveHeap<T, Compare, HeapHandleAccessor>::take(size_type pos) {
806   // Make a hole by taking the element out of the heap.
807   MakeHole(pos);
808   value_type val = std::move(impl_.heap_[pos]);
809 
810   // If the element being taken is already the last element then the heap
811   // doesn't need to be repaired.
812   if (pos != GetLastIndex()) {
813     MakeHole(GetLastIndex());
814 
815     // Move the hole down the heap, filling it with the current leaf at the
816     // very end of the heap.
817     MoveHoleDownAndFill<WithLeafElement>(
818         pos, std::move(impl_.heap_[GetLastIndex()]));
819   }
820 
821   impl_.heap_.pop_back();
822 
823   return val;
824 }
825 
826 // This is effectively identical to "take", but it avoids an unnecessary move.
827 template <typename T, typename Compare, typename HeapHandleAccessor>
erase(size_type pos)828 void IntrusiveHeap<T, Compare, HeapHandleAccessor>::erase(size_type pos) {
829   DCHECK_LT(pos, size());
830   // Make a hole by taking the element out of the heap.
831   MakeHole(pos);
832 
833   // If the element being erased is already the last element then the heap
834   // doesn't need to be repaired.
835   if (pos != GetLastIndex()) {
836     MakeHole(GetLastIndex());
837 
838     // Move the hole down the heap, filling it with the current leaf at the
839     // very end of the heap.
840     MoveHoleDownAndFill<WithLeafElement>(
841         pos, std::move_if_noexcept(impl_.heap_[GetLastIndex()]));
842   }
843 
844   impl_.heap_.pop_back();
845 }
846 
847 template <typename T, typename Compare, typename HeapHandleAccessor>
848 typename IntrusiveHeap<T, Compare, HeapHandleAccessor>::const_iterator
Update(size_type pos)849 IntrusiveHeap<T, Compare, HeapHandleAccessor>::Update(size_type pos) {
850   DCHECK_LT(pos, size());
851   MakeHole(pos);
852 
853   // Determine if we're >= parent, in which case we may need to go up.
854   bool child_greater_eq_parent = false;
855   size_type i = 0;
856   if (pos > 0) {
857     i = intrusive_heap::ParentIndex(pos);
858     child_greater_eq_parent = !Less(pos, i);
859   }
860 
861   if (child_greater_eq_parent) {
862     i = MoveHoleUpAndFill(pos, std::move_if_noexcept(impl_.heap_[pos]));
863   } else {
864     i = MoveHoleDownAndFill<WithElement>(
865         pos, std::move_if_noexcept(impl_.heap_[pos]));
866   }
867 
868   return cbegin() + i;
869 }
870 
871 template <typename T, typename Compare, typename HeapHandleAccessor>
swap(IntrusiveHeap & other)872 void IntrusiveHeap<T, Compare, HeapHandleAccessor>::swap(
873     IntrusiveHeap& other) noexcept {
874   std::swap(impl_.get_value_compare(), other.impl_.get_value_compare());
875   std::swap(impl_.get_heap_handle_access(),
876             other.impl_.get_heap_handle_access());
877   std::swap(impl_.heap_, other.impl_.heap_);
878 }
879 
880 template <typename T, typename Compare, typename HeapHandleAccessor>
881 typename IntrusiveHeap<T, Compare, HeapHandleAccessor>::size_type
ToIndex(const_iterator pos)882 IntrusiveHeap<T, Compare, HeapHandleAccessor>::ToIndex(const_iterator pos) {
883   DCHECK(cbegin() <= pos);
884   DCHECK(pos <= cend());
885   if (pos == cend())
886     return HeapHandle::kInvalidIndex;
887   return pos - cbegin();
888 }
889 
890 template <typename T, typename Compare, typename HeapHandleAccessor>
891 typename IntrusiveHeap<T, Compare, HeapHandleAccessor>::size_type
ToIndex(const_reverse_iterator pos)892 IntrusiveHeap<T, Compare, HeapHandleAccessor>::ToIndex(
893     const_reverse_iterator pos) {
894   DCHECK(crbegin() <= pos);
895   DCHECK(pos <= crend());
896   if (pos == crend())
897     return HeapHandle::kInvalidIndex;
898   return (pos.base() - cbegin()) - 1;
899 }
900 
901 template <typename T, typename Compare, typename HeapHandleAccessor>
SetHeapHandle(size_type i)902 void IntrusiveHeap<T, Compare, HeapHandleAccessor>::SetHeapHandle(size_type i) {
903   impl_.get_heap_handle_access().SetHeapHandle(&impl_.heap_[i], HeapHandle(i));
904   intrusive_heap::CheckInvalidOrEqualTo(GetHeapHandle(i), i);
905 }
906 
907 template <typename T, typename Compare, typename HeapHandleAccessor>
ClearHeapHandle(size_type i)908 void IntrusiveHeap<T, Compare, HeapHandleAccessor>::ClearHeapHandle(
909     size_type i) {
910   impl_.get_heap_handle_access().ClearHeapHandle(&impl_.heap_[i]);
911   DCHECK(!GetHeapHandle(i).IsValid());
912 }
913 
914 template <typename T, typename Compare, typename HeapHandleAccessor>
GetHeapHandle(size_type i)915 HeapHandle IntrusiveHeap<T, Compare, HeapHandleAccessor>::GetHeapHandle(
916     size_type i) {
917   return impl_.get_heap_handle_access().GetHeapHandle(&impl_.heap_[i]);
918 }
919 
920 template <typename T, typename Compare, typename HeapHandleAccessor>
Less(size_type i,size_type j)921 bool IntrusiveHeap<T, Compare, HeapHandleAccessor>::Less(size_type i,
922                                                          size_type j) {
923   DCHECK_LT(i, size());
924   DCHECK_LT(j, size());
925   return impl_.get_value_compare()(impl_.heap_[i], impl_.heap_[j]);
926 }
927 
928 template <typename T, typename Compare, typename HeapHandleAccessor>
Less(const T & element,size_type i)929 bool IntrusiveHeap<T, Compare, HeapHandleAccessor>::Less(const T& element,
930                                                          size_type i) {
931   DCHECK_LT(i, size());
932   return impl_.get_value_compare()(element, impl_.heap_[i]);
933 }
934 
935 template <typename T, typename Compare, typename HeapHandleAccessor>
Less(size_type i,const T & element)936 bool IntrusiveHeap<T, Compare, HeapHandleAccessor>::Less(size_type i,
937                                                          const T& element) {
938   DCHECK_LT(i, size());
939   return impl_.get_value_compare()(impl_.heap_[i], element);
940 }
941 
942 template <typename T, typename Compare, typename HeapHandleAccessor>
MakeHole(size_type pos)943 void IntrusiveHeap<T, Compare, HeapHandleAccessor>::MakeHole(size_type pos) {
944   DCHECK_LT(pos, size());
945   ClearHeapHandle(pos);
946 }
947 
948 template <typename T, typename Compare, typename HeapHandleAccessor>
949 template <typename U>
FillHole(size_type hole_pos,U element)950 void IntrusiveHeap<T, Compare, HeapHandleAccessor>::FillHole(size_type hole_pos,
951                                                              U element) {
952   // The hole that we're filling may not yet exist. This can occur when
953   // inserting a new element into the heap.
954   DCHECK_LE(hole_pos, size());
955   if (hole_pos == size()) {
956     impl_.heap_.push_back(std::move_if_noexcept(element));
957   } else {
958     impl_.heap_[hole_pos] = std::move_if_noexcept(element);
959   }
960   SetHeapHandle(hole_pos);
961 }
962 
963 template <typename T, typename Compare, typename HeapHandleAccessor>
MoveHole(size_type new_hole_pos,size_type old_hole_pos)964 void IntrusiveHeap<T, Compare, HeapHandleAccessor>::MoveHole(
965     size_type new_hole_pos,
966     size_type old_hole_pos) {
967   // The old hole position may be one past the end. This occurs when a new
968   // element is being added.
969   DCHECK_NE(new_hole_pos, old_hole_pos);
970   DCHECK_LT(new_hole_pos, size());
971   DCHECK_LE(old_hole_pos, size());
972 
973   if (old_hole_pos == size()) {
974     impl_.heap_.push_back(std::move_if_noexcept(impl_.heap_[new_hole_pos]));
975   } else {
976     impl_.heap_[old_hole_pos] =
977         std::move_if_noexcept(impl_.heap_[new_hole_pos]);
978   }
979   SetHeapHandle(old_hole_pos);
980 }
981 
982 template <typename T, typename Compare, typename HeapHandleAccessor>
983 template <typename U>
984 typename IntrusiveHeap<T, Compare, HeapHandleAccessor>::size_type
MoveHoleUpAndFill(size_type hole_pos,U element)985 IntrusiveHeap<T, Compare, HeapHandleAccessor>::MoveHoleUpAndFill(
986     size_type hole_pos,
987     U element) {
988   // Moving 1 spot beyond the end is fine. This happens when we insert a new
989   // element.
990   DCHECK_LE(hole_pos, size());
991 
992   // Stop when the element is as far up as it can go.
993   while (hole_pos != 0) {
994     // If our parent is >= to us, we can stop.
995     size_type parent = intrusive_heap::ParentIndex(hole_pos);
996     if (!Less(parent, element))
997       break;
998 
999     MoveHole(parent, hole_pos);
1000     hole_pos = parent;
1001   }
1002 
1003   FillHole(hole_pos, std::move_if_noexcept(element));
1004   return hole_pos;
1005 }
1006 
1007 template <typename T, typename Compare, typename HeapHandleAccessor>
1008 template <typename FillElementType, typename U>
1009 typename IntrusiveHeap<T, Compare, HeapHandleAccessor>::size_type
MoveHoleDownAndFill(size_type hole_pos,U element)1010 IntrusiveHeap<T, Compare, HeapHandleAccessor>::MoveHoleDownAndFill(
1011     size_type hole_pos,
1012     U element) {
1013   DCHECK_LT(hole_pos, size());
1014 
1015   // If we're filling with a leaf, then that leaf element is about to be erased.
1016   // We pretend that the space doesn't exist in the heap.
1017   const size_type n = size() - (FillElementType::kIsLeafElement ? 1 : 0);
1018 
1019   DCHECK_LT(hole_pos, n);
1020   DCHECK(!GetHeapHandle(hole_pos).IsValid());
1021 
1022   while (true) {
1023     // If this spot has no children, then we've gone down as far as we can go.
1024     size_type left = intrusive_heap::LeftIndex(hole_pos);
1025     if (left >= n)
1026       break;
1027     size_type right = left + 1;
1028 
1029     // Get the larger of the potentially two child nodes.
1030     size_type largest = left;
1031     if (right < n && Less(left, right))
1032       largest = right;
1033 
1034     // If we're not deterministically moving the element all the way down to
1035     // become a leaf, then stop when it is >= the largest of the children.
1036     if (!FillElementType::kIsLeafElement && !Less(element, largest))
1037       break;
1038 
1039     MoveHole(largest, hole_pos);
1040     hole_pos = largest;
1041   }
1042 
1043   if (FillElementType::kIsLeafElement) {
1044     // If we're filling with a leaf node we may need to bubble the leaf back up
1045     // the tree a bit to repair the heap.
1046     hole_pos = MoveHoleUpAndFill(hole_pos, std::move_if_noexcept(element));
1047   } else {
1048     FillHole(hole_pos, std::move_if_noexcept(element));
1049   }
1050   return hole_pos;
1051 }
1052 
1053 template <typename T, typename Compare, typename HeapHandleAccessor>
1054 template <typename U>
1055 typename IntrusiveHeap<T, Compare, HeapHandleAccessor>::const_iterator
InsertImpl(U element)1056 IntrusiveHeap<T, Compare, HeapHandleAccessor>::InsertImpl(U element) {
1057   // MoveHoleUpAndFill can tolerate the initial hole being in a slot that
1058   // doesn't yet exist. It will be created by MoveHole by copy/move, thus
1059   // removing the need for a default constructor.
1060   size_type i = MoveHoleUpAndFill(size(), std::move_if_noexcept(element));
1061   return cbegin() + static_cast<difference_type>(i);
1062 }
1063 
1064 template <typename T, typename Compare, typename HeapHandleAccessor>
1065 template <typename U>
1066 typename IntrusiveHeap<T, Compare, HeapHandleAccessor>::const_iterator
ReplaceImpl(size_type pos,U element)1067 IntrusiveHeap<T, Compare, HeapHandleAccessor>::ReplaceImpl(size_type pos,
1068                                                            U element) {
1069   // If we're greater than our parent we need to go up, otherwise we may need
1070   // to go down.
1071   MakeHole(pos);
1072   size_type i = 0;
1073   if (!Less(element, pos)) {
1074     i = MoveHoleUpAndFill(pos, std::move_if_noexcept(element));
1075   } else {
1076     i = MoveHoleDownAndFill<WithElement>(pos, std::move_if_noexcept(element));
1077   }
1078   return cbegin() + static_cast<difference_type>(i);
1079 }
1080 
1081 template <typename T, typename Compare, typename HeapHandleAccessor>
1082 template <typename U>
1083 typename IntrusiveHeap<T, Compare, HeapHandleAccessor>::const_iterator
ReplaceTopImpl(U element)1084 IntrusiveHeap<T, Compare, HeapHandleAccessor>::ReplaceTopImpl(U element) {
1085   MakeHole(0u);
1086   size_type i =
1087       MoveHoleDownAndFill<WithElement>(0u, std::move_if_noexcept(element));
1088   return cbegin() + static_cast<difference_type>(i);
1089 }
1090 
1091 ////////////////////////////////////////////////////////////////////////////////
1092 // WithHeapHandle
1093 
1094 template <typename T>
1095 template <class... Args>
WithHeapHandle(Args &&...args)1096 WithHeapHandle<T>::WithHeapHandle(Args&&... args)
1097     : value_(std::forward<Args>(args)...) {}
1098 
1099 template <typename T>
swap(WithHeapHandle & other)1100 void WithHeapHandle<T>::swap(WithHeapHandle& other) noexcept {
1101   InternalHeapHandleStorage::swap(other);
1102   std::swap(value_, other.value_);
1103 }
1104 
1105 }  // namespace base
1106 
1107 #endif  // BASE_CONTAINERS_INTRUSIVE_HEAP_H_
1108