1/*
2 * Copyright (C) 2017 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#pragma once
18
19#include <utils/RefBase.h>
20#include <unordered_map>
21#include <vector>
22
23using namespace android;
24
25namespace android {
26namespace os {
27namespace statsd {
28
29/** Defines a hash function for sp<const AA>, returning the hash of the underlying pointer. */
30template <class AA>
31struct SpHash {
32    size_t operator()(const sp<const AA>& k) const {
33        return std::hash<const AA*>()(k.get());
34    }
35};
36
37/**
38 * Min priority queue for generic type AA.
39 * Unlike a regular priority queue, this class is also capable of removing interior elements.
40 * @tparam Comparator must implement [bool operator()(sp<const AA> a, sp<const AA> b)], returning
41 *    whether a should be closer to the top of the queue than b.
42 */
43template <class AA, class Comparator>
44class indexed_priority_queue {
45public:
46    indexed_priority_queue();
47    /** Adds a into the priority queue. If already present or a==nullptr, does nothing. */
48    void push(sp<const AA> a);
49    /*
50     * Removes a from the priority queue. If not present or a==nullptr, does nothing.
51     * Returns true if a had been present (and is now removed), else false.
52     */
53    bool remove(sp<const AA> a);
54    /** Removes the top element, if there is one. */
55    void pop();
56    /** Removes all elements. */
57    void clear();
58    /** Returns whether priority queue contains a (not just a copy of a, but a itself). */
59    bool contains(sp<const AA> a) const;
60    /** Returns min element. Returns nullptr iff empty(). */
61    sp<const AA> top() const;
62    /** Returns number of elements in priority queue. */
63    size_t size() const {
64        return pq.size() - 1;
65    }  // pq is 1-indexed
66    /** Returns true iff priority queue is empty. */
67    bool empty() const {
68        return size() < 1;
69    }
70
71private:
72    /** Vector representing a min-heap (1-indexed, with nullptr at 0). */
73    std::vector<sp<const AA>> pq;
74    /** Mapping of each element in pq to its index in pq (i.e. the inverse of a=pq[i]). */
75    std::unordered_map<sp<const AA>, size_t, SpHash<AA>> indices;
76
77    void init();
78    void sift_up(size_t idx);
79    void sift_down(size_t idx);
80    /** Returns whether pq[idx1] is considered higher than pq[idx2], according to Comparator. */
81    bool higher(size_t idx1, size_t idx2) const;
82    void swap_indices(size_t i, size_t j);
83};
84
85// Implementation must be done in this file due to use of template.
86
87template <class AA, class Comparator>
88indexed_priority_queue<AA, Comparator>::indexed_priority_queue() {
89    init();
90}
91
92template <class AA, class Comparator>
93void indexed_priority_queue<AA, Comparator>::push(sp<const AA> a) {
94    if (a == nullptr) return;
95    if (contains(a)) return;
96    pq.push_back(a);
97    size_t idx = size();  // index of last element since 1-indexed
98    indices.insert({a, idx});
99    sift_up(idx);  // get the pq back in order
100}
101
102template <class AA, class Comparator>
103bool indexed_priority_queue<AA, Comparator>::remove(sp<const AA> a) {
104    if (a == nullptr) return false;
105    if (!contains(a)) return false;
106    size_t idx = indices[a];
107    if (idx >= pq.size()) {
108        return false;
109    }
110    if (idx == size()) {  // if a is the last element, i.e. at index idx == size() == (pq.size()-1)
111        pq.pop_back();
112        indices.erase(a);
113        return true;
114    }
115    // move last element (guaranteed not to be at idx) to idx, then delete a
116    sp<const AA> last_a = pq.back();
117    pq[idx] = last_a;
118    pq.pop_back();
119    indices[last_a] = idx;
120    indices.erase(a);
121
122    // get the heap back in order (since the element at idx is not in order)
123    sift_up(idx);
124    sift_down(idx);
125
126    return true;
127}
128
129// The same as, but slightly more efficient than, remove(top()).
130template <class AA, class Comparator>
131void indexed_priority_queue<AA, Comparator>::pop() {
132    sp<const AA> a = top();
133    if (a == nullptr) return;
134    const size_t idx = 1;
135    if (idx == size()) {  // if a is the last element
136        pq.pop_back();
137        indices.erase(a);
138        return;
139    }
140    // move last element (guaranteed not to be at idx) to idx, then delete a
141    sp<const AA> last_a = pq.back();
142    pq[idx] = last_a;
143    pq.pop_back();
144    indices[last_a] = idx;
145    indices.erase(a);
146
147    // get the heap back in order (since the element at idx is not in order)
148    sift_down(idx);
149}
150
151template <class AA, class Comparator>
152void indexed_priority_queue<AA, Comparator>::clear() {
153    pq.clear();
154    indices.clear();
155    init();
156}
157
158template <class AA, class Comparator>
159sp<const AA> indexed_priority_queue<AA, Comparator>::top() const {
160    if (empty()) return nullptr;
161    return pq[1];
162}
163
164template <class AA, class Comparator>
165void indexed_priority_queue<AA, Comparator>::init() {
166    pq.push_back(nullptr);         // so that pq is 1-indexed.
167    indices.insert({nullptr, 0});  // just to be consistent with pq.
168}
169
170template <class AA, class Comparator>
171void indexed_priority_queue<AA, Comparator>::sift_up(size_t idx) {
172    while (idx > 1) {
173        size_t parent = idx / 2;
174        if (higher(idx, parent))
175            swap_indices(idx, parent);
176        else
177            break;
178        idx = parent;
179    }
180}
181
182template <class AA, class Comparator>
183void indexed_priority_queue<AA, Comparator>::sift_down(size_t idx) {
184    while (2 * idx <= size()) {
185        size_t child = 2 * idx;
186        if (child < size() && higher(child + 1, child)) child++;
187        if (higher(child, idx))
188            swap_indices(child, idx);
189        else
190            break;
191        idx = child;
192    }
193}
194
195template <class AA, class Comparator>
196bool indexed_priority_queue<AA, Comparator>::higher(size_t idx1, size_t idx2) const {
197    if (!(0u < idx1 && idx1 < pq.size() && 0u < idx2 && idx2 < pq.size())) {
198        return false;  // got to do something.
199    }
200    return Comparator()(pq[idx1], pq[idx2]);
201}
202
203template <class AA, class Comparator>
204bool indexed_priority_queue<AA, Comparator>::contains(sp<const AA> a) const {
205    if (a == nullptr) return false;  // publicly, we pretend that nullptr is not actually in pq.
206    return indices.count(a) > 0;
207}
208
209template <class AA, class Comparator>
210void indexed_priority_queue<AA, Comparator>::swap_indices(size_t i, size_t j) {
211    if (!(0u < i && i < pq.size() && 0u < j && j < pq.size())) {
212        return;
213    }
214    sp<const AA> val_i = pq[i];
215    sp<const AA> val_j = pq[j];
216    pq[i] = val_j;
217    pq[j] = val_i;
218    indices[val_i] = j;
219    indices[val_j] = i;
220}
221
222}  // namespace statsd
223}  // namespace os
224}  // namespace android
225