molframe_spatial/backends/cell/
budget.rs

1//! Reserve construction peak before allocating a cell index.
2
3use super::{CellGrid, CellList, finite_bounds, grid_geometry, periodic_grid};
4use crate::{CellGridOptions, PeriodicBox, SpatialError};
5use molframe_core::ExecutionContext;
6
7impl<'a> CellList<'a> {
8    /// Builds an index whose construction peak and retained buffers share the
9    /// execution account. Dropping the index releases its retained charge.
10    ///
11    /// # Errors
12    ///
13    /// Returns invalid geometry, cancellation or memory admission errors before
14    /// allocating the grid buffers.
15    pub fn build_in(
16        positions: &'a [[f32; 3]],
17        targets: &[u32],
18        cutoff: f32,
19        periodic: Option<PeriodicBox>,
20        options: CellGridOptions,
21        context: &ExecutionContext,
22    ) -> Result<Self, SpatialError> {
23        if context.cancellation().is_cancelled() {
24            return Err(SpatialError::Cancelled);
25        }
26        super::validate_cutoff(cutoff)?;
27        super::validate_indices(targets, positions.len())?;
28        options.validate()?;
29        let bytes = workspace_bytes(positions, targets, cutoff, periodic.as_ref(), options)?;
30        let mut reservation = context.try_reserve(bytes)?;
31        let mut index = Self::build_with_options(positions, targets, cutoff, periodic, options)?;
32        let retained = match &index.grid {
33            Some(CellGrid::Cartesian(grid)) => grid.retained_bytes(),
34            Some(CellGrid::Periodic(grid)) => grid.retained_bytes(),
35            None => 0,
36        };
37        reservation.shrink_to(retained);
38        index.reservation = Some(reservation);
39        Ok(index)
40    }
41}
42
43fn workspace_bytes(
44    positions: &[[f32; 3]],
45    targets: &[u32],
46    cutoff: f32,
47    periodic: Option<&PeriodicBox>,
48    options: CellGridOptions,
49) -> Result<usize, SpatialError> {
50    let (cells, atoms, occupied) = if let Some(periodic) = periodic {
51        let dims = periodic_grid::periodic_dimensions(periodic, cutoff, options)?;
52        (periodic_grid::cell_count(dims)?, targets.len(), 0)
53    } else if let Some(bounds) = finite_bounds(positions, targets) {
54        let cells = if cutoff == 0.0 {
55            1
56        } else {
57            grid_geometry(bounds.min, bounds.max, cutoff, options)?.2
58        };
59        (cells, bounds.finite_count, cells.min(bounds.finite_count))
60    } else {
61        return Ok(0);
62    };
63    // Counts and prefix offsets coexist with members and occupied-cell indices.
64    // Each vector reserves its final capacity, so construction never doubles it.
65    cells
66        .checked_mul(2)
67        .and_then(|count| count.checked_add(1))
68        .and_then(|count| count.checked_add(occupied))
69        .and_then(|count| count.checked_mul(size_of::<usize>()))
70        .and_then(|bytes| {
71            atoms
72                .checked_mul(size_of::<u32>())
73                .and_then(|members| bytes.checked_add(members))
74        })
75        .ok_or(SpatialError::NumericRangeExceeded)
76}
77
78#[cfg(test)]
79#[path = "budget_tests.rs"]
80mod tests;
81
82impl super::Grid {
83    fn retained_bytes(&self) -> usize {
84        (self.offsets.capacity() + self.non_empty_cells.capacity()) * size_of::<usize>()
85            + self.members.capacity() * size_of::<u32>()
86    }
87}
molframe_spatial/backends/cell/budget.rs

molframe_spatial/backends/cell/
budget.rs
