Skip to main content

uqa_core/memory/
small_vec.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Fallible vectors whose first elements live inline. Only a heap buffer is charged, and both buffers stay charged while elements move between them.
8
9use smallvec::{Array, SmallVec};
10
11use super::{reconcile_buffer_capacity, replacement, MemoryBudget, MemoryError, MemoryReservation};
12
13pub struct BudgetedSmallVec<A: Array> {
14    values: SmallVec<A>,
15    memory: MemoryReservation,
16}
17
18impl<A: Array> BudgetedSmallVec<A> {
19    pub fn new(budget: &MemoryBudget) -> Self {
20        Self {
21            values: SmallVec::new(),
22            memory: budget.empty_reservation(),
23        }
24    }
25
26    pub fn capacity(&self) -> usize {
27        self.values.capacity()
28    }
29
30    /// Whether the elements moved to a charged heap buffer.
31    pub fn spilled(&self) -> bool {
32        self.values.spilled()
33    }
34
35    pub fn budget(&self) -> &MemoryBudget {
36        self.memory.budget()
37    }
38
39    /// Charge a heap buffer for `additional` more elements before allocating it. A failed reservation or allocation preserves the original elements, buffer and reservation.
40    pub fn reserve(&mut self, additional: usize) -> Result<(), MemoryError> {
41        let required = self
42            .values
43            .len()
44            .checked_add(additional)
45            .ok_or(MemoryError::SizeOverflow)?;
46        if required <= self.values.capacity() {
47            return Ok(());
48        }
49        let heap = if self.values.spilled() {
50            self.values.capacity()
51        } else {
52            0
53        };
54        let (capacity, mut memory) = replacement::<A::Item>(self.memory.budget(), heap, required)?;
55        let mut buffer = Vec::new();
56        buffer.try_reserve_exact(capacity)?;
57        reconcile_buffer_capacity::<A::Item>(&mut memory, buffer.capacity())?;
58        buffer.extend(self.values.drain(..));
59        // The required capacity exceeds the inline capacity, so the buffer stays on the heap. Field order frees the old buffer before releasing its reservation.
60        self.values = SmallVec::from_vec(buffer);
61        self.memory = memory;
62        Ok(())
63    }
64
65    pub fn push(&mut self, value: A::Item) -> Result<(), MemoryError> {
66        self.reserve(1)?;
67        self.values.push(value);
68        Ok(())
69    }
70
71    pub fn pop(&mut self) -> Option<A::Item> {
72        self.values.pop()
73    }
74
75    pub fn clear(&mut self) {
76        self.values.clear();
77    }
78
79    pub fn truncate(&mut self, len: usize) {
80        self.values.truncate(len);
81    }
82}
83
84impl<A: Array> std::ops::Deref for BudgetedSmallVec<A> {
85    type Target = [A::Item];
86
87    fn deref(&self) -> &[A::Item] {
88        &self.values
89    }
90}
91
92impl<A: Array> std::ops::DerefMut for BudgetedSmallVec<A> {
93    fn deref_mut(&mut self) -> &mut [A::Item] {
94        &mut self.values
95    }
96}
97
98impl<A: Array> std::fmt::Debug for BudgetedSmallVec<A>
99where
100    A::Item: std::fmt::Debug,
101{
102    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
103        formatter
104            .debug_struct("BudgetedSmallVec")
105            .field("values", &&*self.values)
106            .field("memory", &self.memory)
107            .finish()
108    }
109}
110
111#[cfg(test)]
112mod tests;