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 
23 using namespace android;
24 
25 namespace android {
26 namespace os {
27 namespace statsd {
28 
29 /** Defines a hash function for sp<const AA>, returning the hash of the underlying pointer. */
30 template <class AA>
31 struct SpHash {
operatorSpHash32     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  */
43 template <class AA, class Comparator>
44 class indexed_priority_queue {
45 public:
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. */
size()63     size_t size() const {
64         return pq.size() - 1;
65     }  // pq is 1-indexed
66     /** Returns true iff priority queue is empty. */
empty()67     bool empty() const {
68         return size() < 1;
69     }
70 
71 private:
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 
87 template <class AA, class Comparator>
indexed_priority_queue()88 indexed_priority_queue<AA, Comparator>::indexed_priority_queue() {
89     init();
90 }
91 
92 template <class AA, class Comparator>
push(sp<const AA> a)93 void 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 
102 template <class AA, class Comparator>
remove(sp<const AA> a)103 bool 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()).
130 template <class AA, class Comparator>
pop()131 void 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 
151 template <class AA, class Comparator>
clear()152 void indexed_priority_queue<AA, Comparator>::clear() {
153     pq.clear();
154     indices.clear();
155     init();
156 }
157 
158 template <class AA, class Comparator>
top()159 sp<const AA> indexed_priority_queue<AA, Comparator>::top() const {
160     if (empty()) return nullptr;
161     return pq[1];
162 }
163 
164 template <class AA, class Comparator>
init()165 void 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 
170 template <class AA, class Comparator>
sift_up(size_t idx)171 void 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 
182 template <class AA, class Comparator>
sift_down(size_t idx)183 void 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 
195 template <class AA, class Comparator>
higher(size_t idx1,size_t idx2)196 bool 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 
203 template <class AA, class Comparator>
contains(sp<const AA> a)204 bool 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 
209 template <class AA, class Comparator>
swap_indices(size_t i,size_t j)210 void 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