molframe_spatial/backends/cell/
budget.rs

1//! Reserve construction peak before allocating a cell index.
2
3use super::{CellGrid, CellId, CellList, Layout, choose_layout, finite_bounds, 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        match choose_layout(&bounds, cutoff, options)? {
55            Layout::Dense { cells, .. } => {
56                (cells, bounds.finite_count, cells.min(bounds.finite_count))
57            }
58            // Sorted (cell, target) entries coexist with members, compact
59            // offsets and occupied-cell ids, each bounded by the atom count.
60            Layout::Sparse { .. } => {
61                return bounds
62                    .finite_count
63                    .checked_mul(
64                        size_of::<(CellId, u32)>()
65                            + size_of::<u32>()
66                            + size_of::<usize>()
67                            + size_of::<CellId>(),
68                    )
69                    .and_then(|bytes| bytes.checked_add(size_of::<usize>()))
70                    .ok_or(SpatialError::NumericRangeExceeded);
71            }
72        }
73    } else {
74        return Ok(0);
75    };
76    // Counts and prefix offsets coexist with members and occupied-cell ids.
77    // Each vector reserves its final capacity, so construction never doubles it.
78    cells
79        .checked_mul(2)
80        .and_then(|count| count.checked_add(1))
81        .and_then(|count| count.checked_mul(size_of::<usize>()))
82        .and_then(|bytes| {
83            occupied
84                .checked_mul(size_of::<CellId>())
85                .and_then(|ids| bytes.checked_add(ids))
86        })
87        .and_then(|bytes| {
88            atoms
89                .checked_mul(size_of::<u32>())
90                .and_then(|members| bytes.checked_add(members))
91        })
92        .ok_or(SpatialError::NumericRangeExceeded)
93}
94
95#[cfg(test)]
96#[path = "budget_tests.rs"]
97mod tests;
98
99impl super::Grid {
100    fn retained_bytes(&self) -> usize {
101        self.offsets.capacity() * size_of::<usize>()
102            + self.non_empty_cells.capacity() * size_of::<super::CellId>()
103            + self.members.capacity() * size_of::<u32>()
104    }
105}
molframe_spatial/backends/cell/budget.rs

molframe_spatial/backends/cell/
budget.rs
