VecGeom/base/PriorityQueue.h

0001 #ifndef VECGEOM_BASE_PRIOQUEUE_H_
0002 #define VECGEOM_BASE_PRIOQUEUE_H_
0003
0004 #include "VecGeom/base/Assert.h"
0005
0006 /**
0007  * @brief A simple implementation of a small priority queue, compilable under CPU + GPU
0008  * Can be used in BVH traversal etc
0009  */
0010 namespace vecgeom {
0011
0012 template <typename T, typename Priority, int MaxN = 1000>
0013 struct PriorityQueue {
0014   struct Entry {
0015     T value;
0016     Priority priority;
0017   };
0018   // make sure to never initialize
0019   Entry data[MaxN];
0020   int size = 0;
0021   VECCORE_ATT_HOST_DEVICE void push(T val, Priority p)
0022   {
0023     int i = size++;
0024     VECGEOM_ASSERT(i < MaxN);
0025     data[i] = {val, p};
0026     // Sift up (min-heap on priority)
0027     while (i > 0) {
0028       int parent = (i - 1) >> 1;
0029       if (data[parent].priority <= data[i].priority) {
0030         break;
0031       }
0032       Entry tmp    = data[parent];
0033       data[parent] = data[i];
0034       data[i]      = tmp;
0035       i            = parent;
0036     }
0037   }
0038   VECCORE_ATT_HOST_DEVICE
0039   T pop()
0040   {
0041     T top = data[0].value;
0042     VECGEOM_ASSERT(size > 0);
0043     data[0] = data[--size];
0044     // Sift down
0045     int i = 0;
0046     while (true) {
0047       int l = 2 * i + 1, r = 2 * i + 2, smallest = i;
0048       if (l < size && data[l].priority < data[smallest].priority) {
0049         smallest = l;
0050       }
0051       if (r < size && data[r].priority < data[smallest].priority) {
0052         smallest = r;
0053       }
0054       if (smallest == i) {
0055         break;
0056       }
0057       Entry tmp      = data[smallest];
0058       data[smallest] = data[i];
0059       data[i]        = tmp;
0060       i              = smallest;
0061     }
0062     return top;
0063   }
0064   VECCORE_ATT_HOST_DEVICE T peek() const { return data[0].value; }
0065   VECCORE_ATT_HOST_DEVICE bool empty() const { return size == 0; }
0066   VECCORE_ATT_HOST_DEVICE void clear() { size = 0; }
0067   VECCORE_ATT_HOST_DEVICE Priority peek_priority() const { return data[0].priority; }
0068   VECCORE_ATT_HOST_DEVICE void trim(Priority threshold)
0069   {
0070     while (size > 0 && data[size - 1].priority > threshold)
0071       --size;
0072   }
0073 };
0074
0075 } // namespace vecgeom
0076
0077 #endif