js/src/ds/PriorityQueue.h

   1 /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*-
   2  * vim: set ts=8 sts=2 et sw=2 tw=80:
   3  * This Source Code Form is subject to the terms of the Mozilla Public
   4  * License, v. 2.0. If a copy of the MPL was not distributed with this
   5  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
   6
   7 #ifndef ds_PriorityQueue_h
   8 #define ds_PriorityQueue_h
   9
  10 #include "js/Vector.h"
  11
  12 namespace js {
  13
  14 /*
  15  * Class which represents a heap based priority queue using a vector.
  16  * Inserting elements and removing the highest priority one are both O(log n).
  17  *
  18  * Template parameters are the same as for Vector, with the addition that P
  19  * must have a static priority(const T&) method which returns higher numbers
  20  * for higher priority elements.
  21  */
  22 template <class T, class P, size_t MinInlineCapacity = 0,
  23           class AllocPolicy = TempAllocPolicy>
  24 class PriorityQueue {
  25   Vector<T, MinInlineCapacity, AllocPolicy> heap;
  26
  27   PriorityQueue(const PriorityQueue&) = delete;
  28   PriorityQueue& operator=(const PriorityQueue&) = delete;
  29
  30  public:
  31   explicit PriorityQueue(AllocPolicy ap = AllocPolicy())
  32       : heap(std::move(ap)) {}
  33
  34   [[nodiscard]] bool reserve(size_t capacity) { return heap.reserve(capacity); }
  35
  36   size_t length() const { return heap.length(); }
  37
  38   bool empty() const { return heap.empty(); }
  39
  40   T removeHighest() {
  41     T highest = heap[0];
  42     T last = heap.popCopy();
  43     if (!heap.empty()) {
  44       heap[0] = last;
  45       siftDown(0);
  46     }
  47     return highest;
  48   }
  49
  50   [[nodiscard]] bool insert(const T& v) {
  51     if (!heap.append(v)) {
  52       return false;
  53     }
  54     siftUp(heap.length() - 1);
  55     return true;
  56   }
  57
  58   void infallibleInsert(const T& v) {
  59     heap.infallibleAppend(v);
  60     siftUp(heap.length() - 1);
  61   }
  62
  63  private:
  64   /*
  65    * Elements of the vector encode a binary tree:
  66    *
  67    *      0
  68    *    1   2
  69    *   3 4 5 6
  70    *
  71    * The children of element N are (2N + 1) and (2N + 2).
  72    * The parent of element N is (N - 1) / 2.
  73    *
  74    * Each element has higher priority than its children.
  75    */
  76
  77   void siftDown(size_t n) {
  78     while (true) {
  79       size_t left = n * 2 + 1;
  80       size_t right = n * 2 + 2;
  81
  82       if (left < heap.length()) {
  83         if (right < heap.length()) {
  84           if (P::priority(heap[n]) < P::priority(heap[right]) &&
  85               P::priority(heap[left]) < P::priority(heap[right])) {
  86             swap(n, right);
  87             n = right;
  88             continue;
  89           }
  90         }
  91
  92         if (P::priority(heap[n]) < P::priority(heap[left])) {
  93           swap(n, left);
  94           n = left;
  95           continue;
  96         }
  97       }
  98
  99       break;
 100     }
 101   }
 102
 103   void siftUp(size_t n) {
 104     while (n > 0) {
 105       size_t parent = (n - 1) / 2;
 106
 107       if (P::priority(heap[parent]) > P::priority(heap[n])) {
 108         break;
 109       }
 110
 111       swap(n, parent);
 112       n = parent;
 113     }
 114   }
 115
 116   void swap(size_t a, size_t b) {
 117     T tmp = heap[a];
 118     heap[a] = heap[b];
 119     heap[b] = tmp;
 120   }
 121 };
 122
 123 } /* namespace js */
 124
 125 #endif /* ds_PriorityQueue_h */