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    pub fn reserve(&mut self, additional: usize) -> Result<(), MemoryError> {
36        self.values.reserve(additional)
37    }
38
39    /// Transfer heap-order values with their allocation lease, without copying or sorting them.
40    pub fn into_vec(self) -> BudgetedVec<T> {
41        self.values
42    }
43}
44
45impl<T: Ord> BudgetedBinaryHeap<T> {
46    pub fn push(&mut self, value: T) -> Result<(), MemoryError> {
47        self.values.push(value)?;
48        let mut child = self.values.len() - 1;
49        while child > 0 {
50            let parent = (child - 1) / 2;
51            if self.values[parent] >= self.values[child] {
52                break;
53            }
54            self.values.swap(parent, child);
55            child = parent;
56        }
57        Ok(())
58    }
59
60    pub fn pop(&mut self) -> Option<T> {
61        let last = self.values.pop()?;
62        if self.values.is_empty() {
63            return Some(last);
64        }
65        let result = std::mem::replace(&mut self.values[0], last);
66        let mut root = 0_usize;
67        while let Some(left) = root
68            .checked_mul(2)
69            .and_then(|position| position.checked_add(1))
70            .filter(|&position| position < self.values.len())
71        {
72            let right = left + 1;
73            let child = if right < self.values.len() && self.values[right] > self.values[left] {
74                right
75            } else {
76                left
77            };
78            if self.values[root] >= self.values[child] {
79                break;
80            }
81            self.values.swap(root, child);
82            root = child;
83        }
84        Some(result)
85    }
86}
87
88#[cfg(test)]
89mod tests;