molframe_core/io/
structure_batch.rs

1//! Bounded columnar rows shared by structural format readers.
2
3use crate::column::Presence;
4use crate::execution::Batch;
5use crate::optional::OptionalSymbol;
6use crate::provider::{ChunkDescriptor, ProviderError};
7use crate::symbol::{AltId, DictionaryFull, Interner, SymbolId};
8use crate::{Diagnostic, Element};
9use num_traits::ToPrimitive;
10use std::fmt;
11
12/// Failure shared by bounded structural batch readers and collectors.
13#[derive(Debug)]
14pub enum StructureBatchError {
15    /// A source, syntax or semantic diagnostic.
16    Diagnostic(Diagnostic),
17    /// A single record cannot fit the caller's row or byte demand.
18    DemandTooSmall {
19        /// Minimum retained bytes required.
20        required: usize,
21        /// Retained bytes offered by the demand.
22        available: usize,
23    },
24    /// A single indivisible record cannot fit the execution budget.
25    RecordExceedsBudget {
26        /// Minimum bytes required for the indivisible record.
27        required: usize,
28        /// Bytes available to the source stage.
29        available: usize,
30    },
31    /// Stable dataset, chunk or logical-row arithmetic overflowed.
32    Identity(ProviderError),
33    /// Execution memory was exhausted after source advancement was prevented.
34    Memory(crate::MemoryBudgetError),
35    /// A batch-local identifier dictionary was exhausted.
36    DictionaryFull,
37}
38
39impl fmt::Display for StructureBatchError {
40    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
41        match self {
42            Self::Diagnostic(error) => error.fmt(formatter),
43            Self::DemandTooSmall {
44                required,
45                available,
46            } => write!(
47                formatter,
48                "record requires {required} bytes but demand permits {available}"
49            ),
50            Self::RecordExceedsBudget {
51                required,
52                available,
53            } => write!(
54                formatter,
55                "record requires {required} bytes but budget permits {available}"
56            ),
57            Self::Identity(error) => error.fmt(formatter),
58            Self::Memory(error) => error.fmt(formatter),
59            Self::DictionaryFull => formatter.write_str("batch identifier dictionary is full"),
60        }
61    }
62}
63
64impl std::error::Error for StructureBatchError {}
65
66impl From<Diagnostic> for StructureBatchError {
67    fn from(error: Diagnostic) -> Self {
68        Self::Diagnostic(error)
69    }
70}
71
72impl From<ProviderError> for StructureBatchError {
73    fn from(error: ProviderError) -> Self {
74        Self::Identity(error)
75    }
76}
77
78impl From<DictionaryFull> for StructureBatchError {
79    fn from(_: DictionaryFull) -> Self {
80        Self::DictionaryFull
81    }
82}
83
84impl From<crate::MemoryBudgetError> for StructureBatchError {
85    fn from(error: crate::MemoryBudgetError) -> Self {
86        Self::Memory(error)
87    }
88}
89
90/// Deepest hierarchy level that continues across one batch boundary.
91#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
92pub enum ContinuityLevel {
93    /// No hierarchy row crosses the boundary.
94    #[default]
95    None,
96    /// The model continues but its chain does not.
97    Model,
98    /// The model and chain continue but the residue does not.
99    Chain,
100    /// The model, chain and residue all continue.
101    Residue,
102}
103
104/// Whether hierarchy rows continue across adjacent batches.
105#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
106pub struct BatchContinuity {
107    /// Deepest hierarchy row begun in the preceding batch.
108    pub before: ContinuityLevel,
109    /// Deepest hierarchy row completed in the following batch.
110    pub after: ContinuityLevel,
111}
112
113/// One atom row accepted by [`StructureBatchBuilder`](crate::io::StructureBatchBuilder).
114#[derive(Clone, Copy, Debug)]
115pub struct StructureAtomRecord<'a> {
116    /// Deposited model number.
117    pub model: i32,
118    /// Normalised chain identifier.
119    pub chain: &'a str,
120    /// Component identifier.
121    pub component: &'a str,
122    /// Deposited residue sequence number, or `i32::MIN` when absent.
123    pub sequence: i32,
124    /// Insertion code, or an empty string when absent.
125    pub insertion: &'a str,
126    /// Atom identifier.
127    pub atom: &'a str,
128    /// Alternate-location identifier, or an empty string when absent.
129    pub alternate: &'a str,
130    /// Chemical element.
131    pub element: Element,
132    /// Cartesian position, or `None` when unrecorded.
133    pub position: Option<[f32; 3]>,
134    /// Occupancy and validity.
135    pub occupancy: (f32, Presence),
136    /// B factor and validity.
137    pub b_factor: (f32, Presence),
138    /// Formal charge and validity.
139    pub formal_charge: (i8, Presence),
140    /// Deposited atom-site identifier, with zero reserved for absence.
141    pub atom_site_id: u32,
142    /// Whether the source marked the residue as a heterogen.
143    pub heterogen: bool,
144}
145
146impl StructureAtomRecord<'_> {
147    /// The element a symbol names, or [`Element::UNKNOWN`].
148    #[must_use]
149    pub fn element_of(symbol: &str) -> Element {
150        match Element::from_symbol(symbol) {
151            Some(element) => element,
152            None => Element::UNKNOWN,
153        }
154    }
155
156    /// Whether a row with this element symbol and model survives the read
157    /// options' hydrogen and first-model filters.
158    #[must_use]
159    pub fn admitted(options: &super::ReadOptions, symbol: &str, model: i32) -> bool {
160        let hydrogen = Element::from_symbol(symbol).is_some_and(Element::is_hydrogen);
161        !(options.discard_hydrogens && hydrogen || options.only_first_model && model != 1)
162    }
163
164    /// A position from three recorded coordinates, or `None` when any is
165    /// missing or does not fit single precision.
166    #[must_use]
167    pub fn position_of(x: Option<f64>, y: Option<f64>, z: Option<f64>) -> Option<[f32; 3]> {
168        Some([x?.to_f32()?, y?.to_f32()?, z?.to_f32()?])
169    }
170}
171
172/// Bounded structure rows in a format-neutral layout.
173#[derive(Clone, Debug)]
174pub struct StructureBatch {
175    pub(super) descriptor: ChunkDescriptor,
176    pub(super) continuity: BatchContinuity,
177    pub(super) dictionary: Interner,
178    pub(super) model: Vec<i32>,
179    pub(super) chain: Vec<SymbolId>,
180    pub(super) component: Vec<SymbolId>,
181    pub(super) sequence: Vec<i32>,
182    pub(super) insertion: Vec<OptionalSymbol>,
183    pub(super) atom: Vec<SymbolId>,
184    pub(super) alternate: Vec<AltId>,
185    pub(super) element: Vec<u8>,
186    pub(super) position: Vec<[f32; 3]>,
187    pub(super) coordinate_presence: Vec<Presence>,
188    pub(super) occupancy: Vec<f32>,
189    pub(super) occupancy_presence: Vec<Presence>,
190    pub(super) b_factor: Vec<f32>,
191    pub(super) b_factor_presence: Vec<Presence>,
192    pub(super) formal_charge: Vec<i8>,
193    pub(super) formal_charge_presence: Vec<Presence>,
194    pub(super) atom_site_id: Vec<u32>,
195    pub(super) heterogen: Vec<u8>,
196    pub(super) diagnostics: Vec<Diagnostic>,
197}
198
199impl StructureBatch {
200    /// Reports whether this allocation can serve another batch without
201    /// growing any hot row column or exceeding the next byte demand.
202    #[must_use]
203    pub fn can_reuse(&self, rows: u32, max_retained_bytes: usize) -> bool {
204        let rows = rows as usize;
205        self.model.capacity() >= rows && self.retained_bytes() <= max_retained_bytes
206    }
207
208    /// Stable global identity and local row extent.
209    #[must_use]
210    pub const fn descriptor(&self) -> ChunkDescriptor {
211        self.descriptor
212    }
213
214    /// Hierarchy continuation flags.
215    #[must_use]
216    pub const fn continuity(&self) -> BatchContinuity {
217        self.continuity
218    }
219
220    /// Batch-local identifier dictionary.
221    #[must_use]
222    pub const fn dictionary(&self) -> &Interner {
223        &self.dictionary
224    }
225
226    /// Deposited model numbers.
227    #[must_use]
228    pub fn models(&self) -> &[i32] {
229        &self.model
230    }
231
232    /// Chain identifiers.
233    #[must_use]
234    pub fn chains(&self) -> &[SymbolId] {
235        &self.chain
236    }
237
238    /// Component identifiers.
239    #[must_use]
240    pub fn components(&self) -> &[SymbolId] {
241        &self.component
242    }
243
244    /// Deposited sequence numbers, with `i32::MIN` reserved for absence.
245    #[must_use]
246    pub fn sequences(&self) -> &[i32] {
247        &self.sequence
248    }
249
250    /// Insertion codes.
251    #[must_use]
252    pub fn insertions(&self) -> &[OptionalSymbol] {
253        &self.insertion
254    }
255
256    /// Atom identifiers.
257    #[must_use]
258    pub fn atoms(&self) -> &[SymbolId] {
259        &self.atom
260    }
261
262    /// Alternate-location identifiers.
263    #[must_use]
264    pub fn alternates(&self) -> &[AltId] {
265        &self.alternate
266    }
267
268    /// Atomic numbers.
269    #[must_use]
270    pub fn elements(&self) -> &[u8] {
271        &self.element
272    }
273
274    /// Cartesian positions; missing rows contain `NaN`.
275    #[must_use]
276    pub fn positions(&self) -> &[[f32; 3]] {
277        &self.position
278    }
279
280    /// Position validity.
281    #[must_use]
282    pub fn coordinate_presence(&self) -> &[Presence] {
283        &self.coordinate_presence
284    }
285
286    /// Occupancy values and validity.
287    #[must_use]
288    pub fn occupancies(&self) -> (&[f32], &[Presence]) {
289        (&self.occupancy, &self.occupancy_presence)
290    }
291
292    /// B factors and validity.
293    #[must_use]
294    pub fn b_factors(&self) -> (&[f32], &[Presence]) {
295        (&self.b_factor, &self.b_factor_presence)
296    }
297
298    /// Formal charges and validity.
299    #[must_use]
300    pub fn formal_charges(&self) -> (&[i8], &[Presence]) {
301        (&self.formal_charge, &self.formal_charge_presence)
302    }
303
304    /// Deposited atom-site identifiers.
305    #[must_use]
306    pub fn atom_site_ids(&self) -> &[u32] {
307        &self.atom_site_id
308    }
309
310    /// Per-row heterogen sentinels.
311    #[must_use]
312    pub fn heterogens(&self) -> &[u8] {
313        &self.heterogen
314    }
315
316    /// Diagnostics emitted while producing these logical rows.
317    #[must_use]
318    pub fn diagnostics(&self) -> &[Diagnostic] {
319        &self.diagnostics
320    }
321}
322
323impl Batch for StructureBatch {
324    fn rows(&self) -> usize {
325        self.model.len()
326    }
327
328    fn retained_bytes(&self) -> usize {
329        let mut bytes = self.dictionary.retained_bytes();
330        bytes = bytes.saturating_add(self.model.capacity_bytes());
331        bytes = bytes.saturating_add(self.chain.capacity_bytes());
332        bytes = bytes.saturating_add(self.component.capacity_bytes());
333        bytes = bytes.saturating_add(self.sequence.capacity_bytes());
334        bytes = bytes.saturating_add(self.insertion.capacity_bytes());
335        bytes = bytes.saturating_add(self.atom.capacity_bytes());
336        bytes = bytes.saturating_add(self.alternate.capacity_bytes());
337        bytes = bytes.saturating_add(self.element.capacity_bytes());
338        bytes = bytes.saturating_add(self.position.capacity_bytes());
339        bytes = bytes.saturating_add(self.coordinate_presence.capacity_bytes());
340        bytes = bytes.saturating_add(self.occupancy.capacity_bytes());
341        bytes = bytes.saturating_add(self.occupancy_presence.capacity_bytes());
342        bytes = bytes.saturating_add(self.b_factor.capacity_bytes());
343        bytes = bytes.saturating_add(self.b_factor_presence.capacity_bytes());
344        bytes = bytes.saturating_add(self.formal_charge.capacity_bytes());
345        bytes = bytes.saturating_add(self.formal_charge_presence.capacity_bytes());
346        bytes = bytes.saturating_add(self.atom_site_id.capacity_bytes());
347        bytes = bytes.saturating_add(self.heterogen.capacity_bytes());
348        bytes = bytes.saturating_add(self.diagnostics.capacity_bytes());
349        bytes.saturating_add(
350            self.diagnostics
351                .iter()
352                .map(Diagnostic::retained_bytes)
353                .sum::<usize>(),
354        )
355    }
356}
357
358trait VecCapacityBytes {
359    fn capacity_bytes(&self) -> usize;
360}
361
362impl<T> VecCapacityBytes for Vec<T> {
363    fn capacity_bytes(&self) -> usize {
364        self.capacity().saturating_mul(std::mem::size_of::<T>())
365    }
366}
molframe_core/io/structure_batch.rs

molframe_core/io/
structure_batch.rs
