Home | History | Annotate | Download | only in sequence_manager 
      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:
     23   HeapHandle() : index_(0u) {}
     24 
     25   bool IsValid() const { return index_ != 0u; }
     26 
     27  private:
     28   template <typename T>
     29   friend class IntrusiveHeap;
     30 
     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:
     50   IntrusiveHeap() : nodes_(kMinimumHeapSize), size_(0) {}
     51 
     52   ~IntrusiveHeap() {
     53     for (size_t i = 1; i <= size_; i++) {
     54       MakeHole(i);
     55     }
     56   }
     57 
     58   bool empty() const { return size_ == 0; }
     59 
     60   size_t size() const { return size_; }
     61 
     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 
     70   const T& Min() const {
     71     DCHECK_GE(size_, 1u);
     72     return nodes_[1];
     73   }
     74 
     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 
     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 
     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 
    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 
    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.
    122   const T* begin() const { return &nodes_[1u]; }
    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 
    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|.
    147   void MakeHole(size_t index) {
    148     DCHECK_GT(index, 0u);
    149     DCHECK_LE(index, size_);
    150     nodes_[index].ClearHeapHandle();
    151   }
    152 
    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.
    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.
    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.
    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).
    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