Skip to main content

uqa_core/memory/
heap.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Priority queues reuse controlled vector growth and retain the buffer when transferring results.
8
9use super::{BudgetedVec, MemoryBudget, MemoryError};
10
11#[derive(Debug)]
12pub struct BudgetedBinaryHeap<T> {
13    values: BudgetedVec<T>,
14}
15
16impl<T> BudgetedBinaryHeap<T> {
17    pub fn new(budget: &MemoryBudget) -> Self {
18        Self {
19            values: BudgetedVec::new(budget),
20        }
21    }
22
23    pub fn len(&self) -> usize {
24        self.values.len()
25    }
26
27    pub fn is_empty(&self) -> bool {
28        self.values.is_empty()
29    }
30
31    pub fn peek(&self) -> Option<&T> {
32        self.values.first()
33    }
34
35    /// Borrow the heap's values in heap order without transferring their allocation lease.
36    pub fn as_slice(&self) -> &[T] {
37        &self.values
38    }
39
40    pub fn reserve(&mut self, additional: usize) -> Result<(), MemoryError> {
41        self.values.reserve(additional)
42    }
43
44    /// Transfer heap-order values with their allocation lease, without copying or sorting them.
45    pub fn into_vec(self) -> BudgetedVec<T> {
46        self.values
47    }
48}
49
50impl<T: Ord> BudgetedBinaryHeap<T> {
51    pub fn push(&mut self, value: T) -> Result<(), MemoryError> {
52        self.values.push(value)?;
53        let mut child = self.values.len() - 1;
54        while child > 0 {
55            let parent = (child - 1) / 2;
56            if self.values[parent] >= self.values[child] {
57                break;
58            }
59            self.values.swap(parent, child);
60            child = parent;
61        }
62        Ok(())
63    }
64
65    pub fn pop(&mut self) -> Option<T> {
66        let last = self.values.pop()?;
67        if self.values.is_empty() {
68            return Some(last);
69        }
70        let result = std::mem::replace(&mut self.values[0], last);
71        let mut root = 0_usize;
72        while let Some(left) = root
73            .checked_mul(2)
74            .and_then(|position| position.checked_add(1))
75            .filter(|&position| position < self.values.len())
76        {
77            let right = left + 1;
78            let child = if right < self.values.len() && self.values[right] > self.values[left] {
79                right
80            } else {
81                left
82            };
83            if self.values[root] >= self.values[child] {
84                break;
85            }
86            self.values.swap(root, child);
87            root = child;
88        }
89        Some(result)
90    }
91}
92
93#[cfg(test)]
94mod tests;