molframe_core/io/
structure_batch_builder.rs

1//! Allocation-aware construction and reuse of structure batches.
2
3use super::{BatchContinuity, StructureAtomRecord, StructureBatch, StructureBatchBuffer};
4use crate::Diagnostic;
5use crate::column::Presence;
6use crate::optional::OptionalSymbol;
7use crate::provider::{ChunkDescriptor, ChunkId, DatasetId, LogicalRow, ProviderError};
8use crate::symbol::{AltId, DictionaryFull, Interner};
9
10/// Incremental builder that interns each distinct identifier once per batch.
11#[derive(Debug)]
12pub struct StructureBatchBuilder {
13    buffer: StructureBatchBuffer,
14}
15
16impl StructureBatchBuilder {
17    /// Creates an empty batch with a bounded row reservation.
18    ///
19    /// # Errors
20    ///
21    /// Rejects an unrepresentable descriptor range.
22    pub fn new(
23        dataset: DatasetId,
24        chunk: ChunkId,
25        logical_start: LogicalRow,
26        capacity: u32,
27    ) -> Result<Self, ProviderError> {
28        let descriptor = ChunkDescriptor::new(dataset, chunk, logical_start, 0)?;
29        let capacity = capacity as usize;
30        Ok(Self {
31            buffer: StructureBatchBuffer::new(StructureBatch {
32                descriptor,
33                continuity: BatchContinuity::default(),
34                dictionary: Interner::new(),
35                model: Vec::with_capacity(capacity),
36                chain: Vec::with_capacity(capacity),
37                component: Vec::with_capacity(capacity),
38                sequence: Vec::with_capacity(capacity),
39                insertion: Vec::with_capacity(capacity),
40                atom: Vec::with_capacity(capacity),
41                alternate: Vec::with_capacity(capacity),
42                element: Vec::with_capacity(capacity),
43                position: Vec::with_capacity(capacity),
44                coordinate_presence: Vec::with_capacity(capacity),
45                occupancy: Vec::with_capacity(capacity),
46                occupancy_presence: Vec::with_capacity(capacity),
47                b_factor: Vec::with_capacity(capacity),
48                b_factor_presence: Vec::with_capacity(capacity),
49                formal_charge: Vec::with_capacity(capacity),
50                formal_charge_presence: Vec::with_capacity(capacity),
51                atom_site_id: Vec::with_capacity(capacity),
52                heterogen: Vec::with_capacity(capacity),
53                diagnostics: Vec::new(),
54            }),
55        })
56    }
57
58    /// Resets an exclusively owned batch while preserving column capacity.
59    ///
60    /// The caller must first use [`StructureBatch::can_reuse`] to ensure the
61    /// requested row extent needs no allocation.
62    ///
63    /// # Errors
64    ///
65    /// Rejects an unrepresentable descriptor range.
66    pub(super) fn reuse(
67        mut buffer: StructureBatchBuffer,
68        dataset: DatasetId,
69        chunk: ChunkId,
70        logical_start: LogicalRow,
71    ) -> Result<Self, ProviderError> {
72        let batch = buffer.batch_mut();
73        batch.descriptor = ChunkDescriptor::new(dataset, chunk, logical_start, 0)?;
74        batch.continuity = BatchContinuity::default();
75        batch.dictionary.clear();
76        batch.model.clear();
77        batch.chain.clear();
78        batch.component.clear();
79        batch.sequence.clear();
80        batch.insertion.clear();
81        batch.atom.clear();
82        batch.alternate.clear();
83        batch.element.clear();
84        batch.position.clear();
85        batch.coordinate_presence.clear();
86        batch.occupancy.clear();
87        batch.occupancy_presence.clear();
88        batch.b_factor.clear();
89        batch.b_factor_presence.clear();
90        batch.formal_charge.clear();
91        batch.formal_charge_presence.clear();
92        batch.atom_site_id.clear();
93        batch.heterogen.clear();
94        batch.diagnostics.clear();
95        Ok(Self { buffer })
96    }
97
98    /// Sets hierarchy continuity discovered by the source.
99    pub fn set_continuity(&mut self, continuity: BatchContinuity) {
100        self.buffer.batch_mut().continuity = continuity;
101    }
102
103    /// Appends one row without allocating per atom.
104    ///
105    /// # Errors
106    ///
107    /// Returns [`DictionaryFull`] if the batch-local dictionary is exhausted.
108    pub fn push(&mut self, record: StructureAtomRecord<'_>) -> Result<(), DictionaryFull> {
109        let batch = self.buffer.batch_mut();
110        let chain = batch.dictionary.intern(record.chain)?;
111        let component = batch.dictionary.intern(record.component)?;
112        let atom = batch.dictionary.intern(record.atom)?;
113        let insertion = optional_symbol(&mut batch.dictionary, record.insertion)?;
114        let alternate = optional_alt(&mut batch.dictionary, record.alternate)?;
115        batch.model.push(record.model);
116        batch.chain.push(chain);
117        batch.component.push(component);
118        batch.sequence.push(record.sequence);
119        batch.insertion.push(insertion);
120        batch.atom.push(atom);
121        batch.alternate.push(alternate);
122        batch.element.push(record.element.atomic_number());
123        if let Some(position) = record.position {
124            batch.position.push(position);
125            batch.coordinate_presence.push(Presence::Present);
126        } else {
127            batch.position.push([f32::NAN; 3]);
128            batch.coordinate_presence.push(Presence::Unknown);
129        }
130        batch.occupancy.push(record.occupancy.0);
131        batch.occupancy_presence.push(record.occupancy.1);
132        batch.b_factor.push(record.b_factor.0);
133        batch.b_factor_presence.push(record.b_factor.1);
134        batch.formal_charge.push(record.formal_charge.0);
135        batch.formal_charge_presence.push(record.formal_charge.1);
136        batch.atom_site_id.push(record.atom_site_id);
137        batch.heterogen.push(u8::from(record.heterogen));
138        Ok(())
139    }
140
141    /// Appends a logical-row diagnostic.
142    pub fn push_diagnostic(&mut self, diagnostic: Diagnostic) {
143        self.buffer.batch_mut().diagnostics.push(diagnostic);
144    }
145
146    /// Finalises the descriptor using the actual local extent.
147    ///
148    /// # Errors
149    ///
150    /// Rejects more than `u32::MAX` rows or logical identity overflow.
151    pub fn finish(mut self) -> Result<StructureBatchBuffer, ProviderError> {
152        let batch = self.buffer.batch_mut();
153        let rows =
154            u32::try_from(batch.model.len()).map_err(|_| ProviderError::IdentityOverflow {
155                dataset: batch.descriptor.dataset(),
156                field: "batch rows",
157            })?;
158        batch.descriptor = ChunkDescriptor::new(
159            batch.descriptor.dataset(),
160            batch.descriptor.chunk(),
161            batch.descriptor.logical_start(),
162            rows,
163        )?;
164        Ok(self.buffer)
165    }
166}
167
168fn optional_symbol(
169    dictionary: &mut Interner,
170    text: &str,
171) -> Result<OptionalSymbol, DictionaryFull> {
172    if text.is_empty() {
173        return Ok(OptionalSymbol::NONE);
174    }
175    dictionary.intern(text).map(OptionalSymbol::some)
176}
177
178fn optional_alt(dictionary: &mut Interner, text: &str) -> Result<AltId, DictionaryFull> {
179    if text.is_empty() {
180        return Ok(AltId::BLANK);
181    }
182    let symbol = dictionary.intern(text)?;
183    AltId::labelled(symbol).ok_or(DictionaryFull)
184}
molframe_core/io/structure_batch_builder.rs

molframe_core/io/
structure_batch_builder.rs
