uqa_core/memory/
vec.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Fallible vectors charge both buffers while moving to a larger allocation.
8
9use super::{replacement, MemoryBudget, MemoryError, MemoryReservation};
10
11#[derive(Debug)]
12pub struct BudgetedVec<T> {
13    values: Vec<T>,
14    memory: MemoryReservation,
15}
16
17impl<T> BudgetedVec<T> {
18    pub fn new(budget: &MemoryBudget) -> Self {
19        Self {
20            values: Vec::new(),
21            memory: budget.empty_reservation(),
22        }
23    }
24
25    pub fn capacity(&self) -> usize {
26        self.values.capacity()
27    }
28
29    pub fn budget(&self) -> &MemoryBudget {
30        self.memory.budget()
31    }
32
33    pub fn reserve(&mut self, additional: usize) -> Result<(), MemoryError> {
34        let required = self
35            .values
36            .len()
37            .checked_add(additional)
38            .ok_or(MemoryError::SizeOverflow)?;
39        if required <= self.values.capacity() {
40            return Ok(());
41        }
42        let (capacity, memory) = replacement::<T>(self.memory.budget(), self.capacity(), required)?;
43        let mut values = Vec::new();
44        values.try_reserve_exact(capacity)?;
45        values.append(&mut self.values);
46        // Free the old buffer before releasing its reservation.
47        self.values = values;
48        self.memory = memory;
49        Ok(())
50    }
51
52    pub fn push(&mut self, value: T) -> Result<(), MemoryError> {
53        self.reserve(1)?;
54        self.values.push(value);
55        Ok(())
56    }
57
58    pub fn clear(&mut self) {
59        self.values.clear();
60    }
61
62    pub fn truncate(&mut self, len: usize) {
63        self.values.truncate(len);
64    }
65
66    pub fn pop(&mut self) -> Option<T> {
67        self.values.pop()
68    }
69
70    /// Transfer the buffer and its reservation together; element allocations have separate owners.
71    pub fn into_parts(self) -> (Vec<T>, MemoryReservation) {
72        (self.values, self.memory)
73    }
74}
75
76impl<T> std::ops::Deref for BudgetedVec<T> {
77    type Target = [T];
78
79    fn deref(&self) -> &[T] {
80        &self.values
81    }
82}
83
84impl<T: Copy> BudgetedVec<T> {
85    /// Reserve the complete destination before copying a slice into this buffer.
86    pub fn extend_from_slice(&mut self, values: &[T]) -> Result<(), MemoryError> {
87        self.reserve(values.len())?;
88        self.values.extend_from_slice(values);
89        Ok(())
90    }
91}
92
93impl<T> std::ops::DerefMut for BudgetedVec<T> {
94    fn deref_mut(&mut self) -> &mut [T] {
95        &mut self.values
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn extending_a_slice_reserves_before_mutating_and_preserves_failed_inputs() {
105        let budget = MemoryBudget::new(16);
106        let mut values = BudgetedVec::new(&budget);
107        values.extend_from_slice(&[1_u8, 2, 3, 4]).unwrap();
108        let retained = budget.used();
109        assert!(matches!(
110            values.extend_from_slice(&[9; 32]),
111            Err(MemoryError::Limit { .. })
112        ));
113        assert_eq!(&*values, &[1, 2, 3, 4]);
114        assert_eq!(budget.used(), retained);
115        values.extend_from_slice(&[5, 6]).unwrap();
116        assert_eq!(&*values, &[1, 2, 3, 4, 5, 6]);
117        assert_eq!(budget.used(), values.capacity());
118        drop(values);
119        assert_eq!(budget.used(), 0);
120    }
121}
uqa_core/memory/vec.rs

uqa_core/memory/
vec.rs
