deadlock/unsync/
slotheap.rs

1use std::{
2    cmp::Ordering,
3    ops::{Deref, DerefMut},
4};
5
6use crate::{unsync::SlotMap, util};
7
8/// Single-threaded min-heap with stable, reusable IDs.
9///
10/// Elements are ordered by key `K` via [`Ord`]; each inserted `(key, value)`
11/// receives a stable `usize` ID that remains valid until the element is
12/// removed. Min access, key/value updates by ID, and heap repair after key
13/// changes are provided via guard types ([`PeekMut`], [`PeekKeyMut`],
14/// [`RefMut`], [`RefKeyMut`]) that re-heapify on drop when the key (or the
15/// pair) was mutated. All operations have the time complexities documented
16/// on the methods below.
17///
18/// # Examples
19///
20/// ```rust
21/// use deadlock::unsync::SlotHeap;
22///
23/// let mut heap = SlotHeap::new();
24/// let a = heap.insert(3, "three");
25/// let b = heap.insert(1, "one");
26/// let c = heap.insert(2, "two");
27///
28/// assert_eq!(heap.pop(), Some((1, "one")));
29/// assert_eq!(heap.pop(), Some((2, "two")));
30///
31/// heap.remove(a);
32/// assert_eq!(heap.len(), 0);
33/// ```
34pub struct SlotHeap<K, V> {
35    entries: Vec<Entry<K, V>>,
36    indices: SlotMap<usize>,
37}
38
39struct Entry<K, V> {
40    item: (K, V),
41    id: usize,
42}
43
44impl<K, V> SlotHeap<K, V>
45where
46    K: Ord,
47{
48    /// Creates an empty `SlotHeap`. Time: O(1).
49    pub fn new() -> Self {
50        Self {
51            entries: Vec::new(),
52            indices: SlotMap::new(),
53        }
54    }
55
56    /// Removes all elements and resets internal state. Time: O(n).
57    pub fn clear(&mut self) {
58        self.entries.clear();
59        self.indices.clear();
60    }
61
62    /// Returns the number of live elements. Time: O(1).
63    pub fn len(&self) -> usize {
64        self.entries.len()
65    }
66
67    /// Returns `true` if there are no live elements. Time: O(1).
68    pub fn is_empty(&self) -> bool {
69        self.entries.is_empty()
70    }
71
72    /// Returns `true` if `id` refers to a live element. Time: O(1).
73    pub fn contains(&self, id: usize) -> bool {
74        self.indices.contains(id)
75    }
76
77    /// Inserts `(key, value)` and returns the new element's stable ID.
78    /// Time: O(log n).
79    pub fn insert(&mut self, key: K, value: V) -> usize {
80        let id = self.indices.insert(self.entries.len());
81        self.entries.push(Entry {
82            item: (key, value),
83            id,
84        });
85        unsafe { self.heapify_up(self.entries.len() - 1) };
86        id
87    }
88
89    /// Removes and returns the minimum-key element, or `None` if the heap is
90    /// empty. Time: O(log n).
91    pub fn pop(&mut self) -> Option<(K, V)> {
92        if self.entries.is_empty() {
93            return None;
94        }
95
96        if self.entries.len() == 1 {
97            self.indices.clear();
98            return Some(unsafe { self.entries.pop().unwrap_unchecked().item });
99        }
100
101        let item = unsafe {
102            let entry = util::swap_remove_unchecked(&mut self.entries, 0);
103            self.indices.remove_unchecked(entry.id);
104            *self
105                .indices
106                .get_unchecked_mut(self.entries.get_unchecked(0).id) = 0;
107            entry.item
108        };
109
110        unsafe { self.heapify_down(0) };
111        Some(item)
112    }
113
114    /// Removes the element with the given `id` and returns `(key, value)`, or
115    /// `None` if `id` is not a live element. Time: O(log n).
116    pub fn remove(&mut self, id: usize) -> Option<(K, V)> {
117        let index = self.indices.remove(id)?;
118        let len = self.entries.len();
119
120        if index == len - 1 {
121            return Some(unsafe { self.entries.pop().unwrap_unchecked().item });
122        }
123
124        let item = unsafe {
125            let entry = util::swap_remove_unchecked(&mut self.entries, index);
126            *self
127                .indices
128                .get_unchecked_mut(self.entries.get_unchecked(index).id) = index;
129            entry.item
130        };
131
132        unsafe { self.heapify(index) };
133        Some(item)
134    }
135
136    /// Returns a shared reference to the minimum element's `(key, value)` pair.
137    /// Time: O(1).
138    pub fn peek(&self) -> Option<&(K, V)> {
139        self.entries.first().map(|item| &item.item)
140    }
141
142    /// Returns a shared reference to the minimum element's key. Time: O(1).
143    pub fn peek_key(&self) -> Option<&K> {
144        self.entries.first().map(|item| &item.item.0)
145    }
146
147    /// Returns a shared reference to the minimum element's value. Time: O(1).
148    pub fn peek_value(&self) -> Option<&V> {
149        self.entries.first().map(|item| &item.item.1)
150    }
151
152    /// Returns an exclusive guard over the minimum element's `(key, value)` pair.
153    /// If the guard is mutated, the heap invariant is restored on drop.
154    /// Time: O(1) to create; drop may run O(log n) heapify.
155    pub fn peek_mut(&mut self) -> Option<PeekMut<'_, K, V>> {
156        if self.entries.is_empty() {
157            return None;
158        }
159
160        Some(PeekMut {
161            dirty: false,
162            from: self,
163        })
164    }
165
166    /// Returns an exclusive guard over the minimum element's key. If the key
167    /// is modified, the heap invariant is restored on drop.
168    /// Time: O(1) to create; drop may run O(log n) heapify.
169    pub fn peek_key_mut(&mut self) -> Option<PeekKeyMut<'_, K, V>> {
170        util::ensure!(!self.entries.is_empty());
171
172        Some(PeekKeyMut {
173            dirty: false,
174            from: self,
175        })
176    }
177
178    /// Returns an exclusive reference to the minimum element's value.
179    /// Modifying only the value does not affect the heap ordering. Time: O(1).
180    pub fn peek_value_mut(&mut self) -> Option<&mut V> {
181        self.entries.first_mut().map(|entry| &mut entry.item.1)
182    }
183
184    /// Returns a shared reference to the `(key, value)` pair of the element
185    /// with `id`, or `None` if `id` is not live. Time: O(1).
186    pub fn get(&self, id: usize) -> Option<&(K, V)> {
187        let index = *self.indices.get(id)?;
188        Some(unsafe { &self.entries.get_unchecked(index).item })
189    }
190
191    /// Returns a shared reference to the key of the element with `id`. Time: O(1).
192    pub fn get_key(&self, id: usize) -> Option<&K> {
193        let index = *self.indices.get(id)?;
194        Some(unsafe { &self.entries.get_unchecked(index).item.0 })
195    }
196
197    /// Returns a shared reference to the value of the element with `id`. Time: O(1).
198    pub fn get_value(&self, id: usize) -> Option<&V> {
199        let index = *self.indices.get(id)?;
200        Some(unsafe { &self.entries.get_unchecked(index).item.1 })
201    }
202
203    /// Returns an exclusive guard over the `(key, value)` pair of the element
204    /// with `id`. If the guard is mutated, the heap invariant is restored on drop.
205    /// Time: O(1) to create; drop may run O(log n) heapify.
206    pub fn get_mut(&mut self, id: usize) -> Option<RefMut<'_, K, V>> {
207        let index = *self.indices.get(id)?;
208        Some(RefMut {
209            dirty: false,
210            from: self,
211            index,
212        })
213    }
214
215    /// Returns an exclusive guard over the key of the element with `id`. If the
216    /// key is modified, the heap invariant is restored on drop.
217    /// Time: O(1) to create; drop may run O(log n) heapify.
218    pub fn get_key_mut(&mut self, id: usize) -> Option<RefKeyMut<'_, K, V>> {
219        let index = *self.indices.get(id)?;
220        Some(RefKeyMut {
221            dirty: false,
222            from: self,
223            index,
224        })
225    }
226
227    /// Returns an exclusive reference to the value of the element with `id`.
228    /// Modifying only the value does not affect the heap ordering. Time: O(1).
229    pub fn get_value_mut(&mut self, id: usize) -> Option<&mut V> {
230        let index = *self.indices.get(id)?;
231        Some(unsafe { &mut self.entries.get_unchecked_mut(index).item.1 })
232    }
233
234    pub(crate) unsafe fn heapify(&mut self, now: usize) {
235        let now = unsafe { self.heapify_up(now) };
236        unsafe { self.heapify_down(now) };
237    }
238
239    unsafe fn heapify_up(&mut self, mut now: usize) -> usize {
240        unsafe {
241            while let Some(up) = self.next_up(now) {
242                self.swap_entries(now, up);
243                now = up;
244            }
245        }
246
247        now
248    }
249
250    pub(crate) unsafe fn heapify_down(&mut self, mut now: usize) {
251        unsafe {
252            while let Some(down) = self.next_down(now) {
253                self.swap_entries(now, down);
254                now = down;
255            }
256        }
257    }
258
259    unsafe fn next_up(&self, now: usize) -> Option<usize> {
260        now.checked_sub(1).map(|x| x / 2).filter(|up| unsafe {
261            self.entries.get_unchecked(now) < self.entries.get_unchecked(*up)
262        })
263    }
264
265    unsafe fn next_down(&self, now: usize) -> Option<usize> {
266        let now_item = unsafe { self.entries.get_unchecked(now) };
267        let (left, right) = (now * 2 + 1, now * 2 + 2);
268
269        if right < self.entries.len() {
270            let left_item = unsafe { self.entries.get_unchecked(left) };
271            let right_item = unsafe { self.entries.get_unchecked(right) };
272
273            if left_item < right_item {
274                (left_item < now_item).then_some(left)
275            } else {
276                (right_item < now_item).then_some(right)
277            }
278        } else {
279            let left_item = self.entries.get(left)?;
280            (left_item < now_item).then_some(left)
281        }
282    }
283
284    unsafe fn swap_entries(&mut self, index0: usize, index1: usize) {
285        unsafe {
286            util::swap_unchecked(&mut self.entries, index0, index1);
287            let id0 = self.entries.get_unchecked(index0).id;
288            let id1 = self.entries.get_unchecked(index1).id;
289            *self.indices.get_unchecked_mut(id0) = index0;
290            *self.indices.get_unchecked_mut(id1) = index1;
291        }
292    }
293
294    pub(crate) fn get_index(&self, id: usize) -> Option<usize> {
295        self.indices.get(id).copied()
296    }
297
298    pub(crate) unsafe fn by_index(&self, index: usize) -> &(K, V) {
299        unsafe { &self.entries.get_unchecked(index).item }
300    }
301
302    pub(crate) unsafe fn by_index_mut(&mut self, index: usize) -> &mut (K, V) {
303        unsafe { &mut self.entries.get_unchecked_mut(index).item }
304    }
305}
306
307impl<K, V> Default for SlotHeap<K, V>
308where
309    K: Ord,
310{
311    fn default() -> Self {
312        Self::new()
313    }
314}
315
316impl<K, V> PartialEq for Entry<K, V>
317where
318    K: Ord,
319{
320    fn eq(&self, other: &Self) -> bool {
321        self.id == other.id
322    }
323}
324
325impl<K, V> Eq for Entry<K, V> where K: Ord {}
326
327impl<K, V> PartialOrd for Entry<K, V>
328where
329    K: Ord,
330{
331    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
332        Some(self.cmp(other))
333    }
334}
335
336impl<K, V> Ord for Entry<K, V>
337where
338    K: Ord,
339{
340    fn cmp(&self, other: &Self) -> Ordering {
341        self.item
342            .0
343            .cmp(&other.item.0)
344            .then_with(|| self.id.cmp(&other.id))
345    }
346}
347
348/// Exclusive guard over the minimum element's `(key, value)` pair. Mutating
349/// the guard marks the heap as dirty; on drop the heap is re-heapified.
350pub struct PeekMut<'a, K, V>
351where
352    K: Ord,
353{
354    dirty: bool,
355    from: &'a mut SlotHeap<K, V>,
356}
357
358impl<'a, K, V> Deref for PeekMut<'a, K, V>
359where
360    K: Ord,
361{
362    type Target = (K, V);
363
364    fn deref(&self) -> &Self::Target {
365        unsafe { &self.from.entries.get_unchecked(0).item }
366    }
367}
368
369impl<'a, K, V> DerefMut for PeekMut<'a, K, V>
370where
371    K: Ord,
372{
373    fn deref_mut(&mut self) -> &mut Self::Target {
374        self.dirty = true;
375        unsafe { &mut self.from.entries.get_unchecked_mut(0).item }
376    }
377}
378
379impl<'a, K, V> AsRef<(K, V)> for PeekMut<'a, K, V>
380where
381    K: Ord,
382{
383    fn as_ref(&self) -> &(K, V) {
384        self.deref()
385    }
386}
387
388impl<'a, K, V> AsMut<(K, V)> for PeekMut<'a, K, V>
389where
390    K: Ord,
391{
392    fn as_mut(&mut self) -> &mut (K, V) {
393        self.deref_mut()
394    }
395}
396
397impl<'a, K, V> Drop for PeekMut<'a, K, V>
398where
399    K: Ord,
400{
401    fn drop(&mut self) {
402        if self.dirty {
403            unsafe { self.from.heapify_down(0) }
404        }
405    }
406}
407
408/// Exclusive guard over the minimum element's key. Mutating the guard marks
409/// the heap as dirty; on drop the heap is re-heapified.
410pub struct PeekKeyMut<'a, K, V>
411where
412    K: Ord,
413{
414    dirty: bool,
415    from: &'a mut SlotHeap<K, V>,
416}
417
418impl<'a, K, V> Deref for PeekKeyMut<'a, K, V>
419where
420    K: Ord,
421{
422    type Target = K;
423
424    fn deref(&self) -> &Self::Target {
425        unsafe { &self.from.entries.get_unchecked(0).item.0 }
426    }
427}
428
429impl<'a, K, V> DerefMut for PeekKeyMut<'a, K, V>
430where
431    K: Ord,
432{
433    fn deref_mut(&mut self) -> &mut Self::Target {
434        self.dirty = true;
435        unsafe { &mut self.from.entries.get_unchecked_mut(0).item.0 }
436    }
437}
438
439impl<'a, K, V> AsRef<K> for PeekKeyMut<'a, K, V>
440where
441    K: Ord,
442{
443    fn as_ref(&self) -> &K {
444        self.deref()
445    }
446}
447
448impl<'a, K, V> AsMut<K> for PeekKeyMut<'a, K, V>
449where
450    K: Ord,
451{
452    fn as_mut(&mut self) -> &mut K {
453        self.deref_mut()
454    }
455}
456
457impl<'a, K, V> Drop for PeekKeyMut<'a, K, V>
458where
459    K: Ord,
460{
461    fn drop(&mut self) {
462        if self.dirty {
463            unsafe { self.from.heapify_down(0) }
464        }
465    }
466}
467
468/// Exclusive guard over the `(key, value)` pair of an element identified by
469/// index. Mutating the guard marks the heap as dirty; on drop the heap is
470/// re-heapified.
471pub struct RefMut<'a, K, V>
472where
473    K: Ord,
474{
475    dirty: bool,
476    from: &'a mut SlotHeap<K, V>,
477    index: usize,
478}
479
480impl<'a, K, V> Deref for RefMut<'a, K, V>
481where
482    K: Ord,
483{
484    type Target = (K, V);
485
486    fn deref(&self) -> &Self::Target {
487        unsafe { &self.from.entries.get_unchecked(self.index).item }
488    }
489}
490
491impl<'a, K, V> DerefMut for RefMut<'a, K, V>
492where
493    K: Ord,
494{
495    fn deref_mut(&mut self) -> &mut Self::Target {
496        self.dirty = true;
497        unsafe { &mut self.from.entries.get_unchecked_mut(self.index).item }
498    }
499}
500
501impl<'a, K, V> AsRef<(K, V)> for RefMut<'a, K, V>
502where
503    K: Ord,
504{
505    fn as_ref(&self) -> &(K, V) {
506        self.deref()
507    }
508}
509
510impl<'a, K, V> AsMut<(K, V)> for RefMut<'a, K, V>
511where
512    K: Ord,
513{
514    fn as_mut(&mut self) -> &mut (K, V) {
515        self.deref_mut()
516    }
517}
518
519impl<'a, K, V> Drop for RefMut<'a, K, V>
520where
521    K: Ord,
522{
523    fn drop(&mut self) {
524        if self.dirty {
525            unsafe { self.from.heapify(self.index) }
526        }
527    }
528}
529
530/// Exclusive guard over the key of an element identified by index. Mutating
531/// the guard marks the heap as dirty; on drop the heap is re-heapified.
532pub struct RefKeyMut<'a, K, V>
533where
534    K: Ord,
535{
536    dirty: bool,
537    from: &'a mut SlotHeap<K, V>,
538    index: usize,
539}
540
541impl<'a, K, V> Deref for RefKeyMut<'a, K, V>
542where
543    K: Ord,
544{
545    type Target = K;
546
547    fn deref(&self) -> &Self::Target {
548        unsafe { &self.from.entries.get_unchecked(self.index).item.0 }
549    }
550}
551
552impl<'a, K, V> DerefMut for RefKeyMut<'a, K, V>
553where
554    K: Ord,
555{
556    fn deref_mut(&mut self) -> &mut Self::Target {
557        self.dirty = true;
558        unsafe { &mut self.from.entries.get_unchecked_mut(self.index).item.0 }
559    }
560}
561
562impl<'a, K, V> AsRef<K> for RefKeyMut<'a, K, V>
563where
564    K: Ord,
565{
566    fn as_ref(&self) -> &K {
567        self.deref()
568    }
569}
570
571impl<'a, K, V> AsMut<K> for RefKeyMut<'a, K, V>
572where
573    K: Ord,
574{
575    fn as_mut(&mut self) -> &mut K {
576        self.deref_mut()
577    }
578}
579
580impl<'a, K, V> Drop for RefKeyMut<'a, K, V>
581where
582    K: Ord,
583{
584    fn drop(&mut self) {
585        if self.dirty {
586            unsafe { self.from.heapify(self.index) }
587        }
588    }
589}
deadlock/unsync/slotheap.rs

deadlock/unsync/
slotheap.rs
