Skip to main content

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