uqa_core/

memory.rs

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Shared byte allowances with reservations that follow allocation ownership.
//!
//! Owners reserve requested buffer layouts before allocating and charge any additional reported capacity before moving values. Allocator bookkeeping, borrowed data, and separately owned immutable resources are outside this allowance.

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

use triomphe::Arc as StrongArc;

mod deque;
mod hash_set;
mod heap;
mod map;
mod production;
mod small_vec;
mod string;
mod vec;

pub use deque::BudgetedDeque;
pub use hash_set::BudgetedHashSet;
pub use heap::BudgetedBinaryHeap;
pub use map::{
    BudgetedMap, BudgetedMapIter, BudgetedSharedMap, BudgetedSharedMapIter,
    BudgetedSharedMapSnapshot, OwnedMap, OwnedMapIntoIter, OwnedSet, OwnedSetIntoIter,
    OwnedSetIter, PreparedMapEntry,
};
pub use production::{Produced, ProductionControl, ProductionString, ProductionVec};
pub use small_vec::BudgetedSmallVec;
pub use string::BudgetedString;
pub use vec::BudgetedVec;

#[derive(Debug, thiserror::Error)]
pub enum MemoryError {
    #[error("memory requires {required} bytes, exceeding limit {limit}")]
    Limit { required: usize, limit: usize },
    #[error("memory allocation size overflow")]
    SizeOverflow,
    #[error("memory allocation failed: {0}")]
    Allocation(#[from] std::collections::TryReserveError),
}

#[derive(Debug)]
struct Allowance {
    limit: usize,
    used: AtomicUsize,
    peak: AtomicUsize,
    parent: Option<MemoryBudget>,
}

impl Allowance {
    fn claim(&self, additional: usize) -> Result<(), MemoryError> {
        let mut used = self.used.load(Ordering::Relaxed);
        loop {
            let required = used
                .checked_add(additional)
                .ok_or(MemoryError::SizeOverflow)?;
            if required > self.limit {
                return Err(MemoryError::Limit {
                    required,
                    limit: self.limit,
                });
            }
            match self.used.compare_exchange_weak(
                used,
                required,
                Ordering::Relaxed,
                Ordering::Relaxed,
            ) {
                Ok(_) => {
                    self.peak.fetch_max(required, Ordering::Relaxed);
                    return Ok(());
                }
                Err(current) => used = current,
            }
        }
    }
}

impl Drop for Allowance {
    fn drop(&mut self) {
        // Destroy an unshared parent chain iteratively, including deeply nested application budgets.
        let mut parent = self.parent.take();
        while let Some(budget) = parent {
            match StrongArc::try_unwrap(budget.0) {
                Ok(mut allowance) => parent = allowance.parent.take(),
                Err(_) => break,
            }
        }
    }
}

/// Clones share one allowance, including reservations retained by completed producers.
#[derive(Debug, Clone)]
pub struct MemoryBudget(StrongArc<Allowance>);

impl MemoryBudget {
    pub fn new(limit: usize) -> Self {
        Self(StrongArc::new(Allowance {
            limit,
            used: AtomicUsize::new(0),
            peak: AtomicUsize::new(0),
            parent: None,
        }))
    }

    /// Add a component limit without enlarging this allowance. Every descendant reservation also charges each ancestor until its final owner releases it.
    pub fn child(&self, limit: usize) -> Self {
        Self(StrongArc::new(Allowance {
            limit,
            used: AtomicUsize::new(0),
            peak: AtomicUsize::new(0),
            parent: Some(self.clone()),
        }))
    }

    fn allowances(&self) -> impl Iterator<Item = &Allowance> {
        std::iter::successors(Some(self), |budget| budget.0.parent.as_ref())
            .map(|budget| &*budget.0)
    }

    pub fn limit(&self) -> usize {
        self.0.limit
    }

    pub fn used(&self) -> usize {
        self.0.used.load(Ordering::Relaxed)
    }

    /// Currently unreserved bytes under this limit and every ancestor limit. This is a sizing hint; concurrent owners can consume it before a reservation succeeds.
    pub fn available(&self) -> usize {
        self.allowances()
            .map(|allowance| {
                allowance
                    .limit
                    .saturating_sub(allowance.used.load(Ordering::Relaxed))
            })
            .min()
            .expect("a budget has its own allowance")
    }

    pub fn shares_allowance(&self, other: &Self) -> bool {
        StrongArc::ptr_eq(&self.0, &other.0)
    }

    /// Largest simultaneous reservation, including old and replacement buffers.
    pub fn peak(&self) -> usize {
        self.0.peak.load(Ordering::Relaxed)
    }

    pub fn reserve(&self, bytes: usize) -> Result<MemoryReservation, MemoryError> {
        let mut reservation = self.empty_reservation();
        reservation.grow(bytes)?;
        Ok(reservation)
    }

    pub fn empty_reservation(&self) -> MemoryReservation {
        MemoryReservation {
            budget: self.clone(),
            bytes: 0,
        }
    }
}

/// A unique lease: release it only after the associated allocation has been freed.
#[derive(Debug)]
pub struct MemoryReservation {
    budget: MemoryBudget,
    bytes: usize,
}

impl MemoryReservation {
    pub fn bytes(&self) -> usize {
        self.bytes
    }

    pub fn budget(&self) -> &MemoryBudget {
        &self.budget
    }

    /// Transfer part of an existing allocation lease without releasing or reserving bytes.
    ///
    /// Panics if `bytes` exceeds this lease. The original lease is unchanged on failure.
    pub fn split(&mut self, bytes: usize) -> Self {
        self.bytes = self
            .bytes
            .checked_sub(bytes)
            .expect("insufficient reserved bytes");
        Self {
            budget: self.budget.clone(),
            bytes,
        }
    }

    /// Reserve additional bytes atomically; a failed request leaves this lease unchanged.
    pub fn grow(&mut self, additional: usize) -> Result<(), MemoryError> {
        if additional == 0 {
            return Ok(());
        }
        for (claimed, allowance) in self.budget.allowances().enumerate() {
            if let Err(error) = allowance.claim(additional) {
                for previous in self.budget.allowances().take(claimed) {
                    previous.used.fetch_sub(additional, Ordering::Relaxed);
                }
                return Err(error);
            }
        }
        self.bytes += additional;
        Ok(())
    }

    /// Transfer ownership without releasing and reacquiring the shared allowance.
    ///
    /// Panics if the leases belong to different allowances.
    pub fn absorb(&mut self, mut other: Self) {
        assert!(
            self.budget.shares_allowance(&other.budget),
            "different memory allowances"
        );
        self.bytes += other.bytes;
        other.bytes = 0;
    }
}

impl Drop for MemoryReservation {
    fn drop(&mut self) {
        for allowance in self.budget.allowances() {
            allowance.used.fetch_sub(self.bytes, Ordering::Relaxed);
        }
    }
}

/// An immutable result and its owned allocations. Cloning the value is a separate allocation.
#[derive(Debug)]
pub struct Budgeted<T> {
    // Field order ensures that the value is destroyed before its allowance is returned.
    value: T,
    memory: MemoryReservation,
}

impl<T> Budgeted<T> {
    /// The caller must associate all owned allocations with this reservation.
    pub fn new(value: T, memory: MemoryReservation) -> Self {
        Self { value, memory }
    }

    pub fn reserved_bytes(&self) -> usize {
        self.memory.bytes()
    }

    pub fn budget(&self) -> &MemoryBudget {
        self.memory.budget()
    }

    /// Share the immutable value and its lease, reserving the shared payload before allocation. Reference-count bookkeeping is outside the payload allowance.
    pub fn into_shared(self) -> Result<Arc<Self>, MemoryError> {
        let mut memory = self.memory.budget().reserve(std::mem::size_of::<Self>())?;
        let (value, allocations) = self.into_parts();
        memory.absorb(allocations);
        Ok(Arc::new(Self::new(value, memory)))
    }

    /// Move the value and lease together to another allocation owner.
    pub fn into_parts(self) -> (T, MemoryReservation) {
        (self.value, self.memory)
    }
}

impl<T> std::ops::Deref for Budgeted<T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.value
    }
}

impl<T: PartialEq> PartialEq for Budgeted<T> {
    fn eq(&self, other: &Self) -> bool {
        self.value == other.value
    }
}

impl<T: Eq> Eq for Budgeted<T> {}

impl<T: AsRef<U>, U: ?Sized> AsRef<U> for Budgeted<T> {
    fn as_ref(&self) -> &U {
        self.value.as_ref()
    }
}

fn buffer_bytes<T>(capacity: usize) -> Result<usize, MemoryError> {
    capacity
        .checked_mul(std::mem::size_of::<T>())
        .filter(|bytes| isize::try_from(*bytes).is_ok())
        .ok_or(MemoryError::SizeOverflow)
}

fn reconcile_buffer_capacity<T>(
    memory: &mut MemoryReservation,
    capacity: usize,
) -> Result<(), MemoryError> {
    memory.grow(buffer_bytes::<T>(capacity)?.saturating_sub(memory.bytes()))
}

fn replacement<T>(
    budget: &MemoryBudget,
    capacity: usize,
    required: usize,
) -> Result<(usize, MemoryReservation), MemoryError> {
    let preferred = capacity.saturating_mul(2).max(required);
    if let Ok(bytes) = buffer_bytes::<T>(preferred) {
        if let Ok(memory) = budget.reserve(bytes) {
            return Ok((preferred, memory));
        }
    }
    Ok((required, budget.reserve(buffer_bytes::<T>(required)?)?))
}

#[cfg(test)]
mod tests;