molframe_cif/
batch.rs

1//! Pull-based coordinate projection over bounded CIF token windows.
2
3mod budget;
4mod token;
5
6use budget::{identity_overflow, reserve_bytes, row_capacity};
7use molframe_core::{
8    Backpressure, Batch, BatchContinuity, BatchDemand, BatchLease, BatchSource, ChunkId,
9    ContinuityLevel, DatasetId, Element, ExecutionContext, LogicalRow, Presence, ReadOptions,
10    SourceBytes, StructureAtomRecord, StructureBatch, StructureBatchBuffer, StructureBatchBuilder,
11    StructureBatchError, StructureBatchPool,
12};
13use num_traits::ToPrimitive;
14use token::TokenCursor;
15
16/// Incremental mmCIF coordinate rows.
17#[derive(Debug)]
18pub struct MmcifBatchSource<S: SourceBytes> {
19    tokens: TokenCursor<S>,
20    options: ReadOptions,
21    dataset: DatasetId,
22    chunk: ChunkId,
23    logical_row: LogicalRow,
24    parser: ParserState,
25    continuation: ContinuityLevel,
26    pool: StructureBatchPool,
27}
28
29impl<S: SourceBytes> MmcifBatchSource<S> {
30    /// Creates a reader whose token and source windows never grow with file size.
31    ///
32    /// # Errors
33    ///
34    /// Returns a budget error when the bounded token carry cannot be reserved.
35    pub fn new(
36        source: S,
37        options: ReadOptions,
38        dataset: DatasetId,
39        first_chunk: ChunkId,
40        first_row: LogicalRow,
41        window_bytes: usize,
42        context: &ExecutionContext,
43    ) -> Result<Self, StructureBatchError> {
44        Ok(Self {
45            tokens: TokenCursor::new(source, window_bytes, context)?,
46            options,
47            dataset,
48            chunk: first_chunk,
49            logical_row: first_row,
50            parser: ParserState::new(),
51            continuation: ContinuityLevel::None,
52            pool: StructureBatchPool::new(),
53        })
54    }
55}
56
57impl<S: SourceBytes> BatchSource for MmcifBatchSource<S> {
58    type Batch = StructureBatch;
59    type Error = StructureBatchError;
60
61    fn next_batch(
62        &mut self,
63        demand: BatchDemand,
64        context: &ExecutionContext,
65    ) -> Result<Backpressure<BatchLease<Self::Batch>>, Self::Error> {
66        if self.parser.finished {
67            return Ok(Backpressure::Finished);
68        }
69        if !demand.can_accept_work() {
70            return Ok(Backpressure::Pending);
71        }
72        let capacity = row_capacity(demand)?;
73        let reserve = reserve_bytes(capacity, demand);
74        let descriptor = molframe_core::ChunkDescriptor::new(
75            self.dataset,
76            self.chunk,
77            self.logical_row,
78            capacity,
79        )?;
80        let Some((builder, reservation)) =
81            self.pool
82                .acquire(descriptor, reserve, demand.max_bytes, context)?
83        else {
84            return Ok(Backpressure::Pending);
85        };
86        let token_checkpoint = self.tokens.checkpoint();
87        let parser_checkpoint = self.parser.checkpoint();
88        let parsed = match parse_batch(
89            &mut self.tokens,
90            &mut self.parser,
91            &self.options,
92            builder,
93            capacity,
94            self.continuation,
95        ) {
96            Ok(parsed) => parsed,
97            Err(error) => {
98                self.tokens.restore(token_checkpoint);
99                self.parser.restore(parser_checkpoint);
100                return Err(error);
101            }
102        };
103        let Some(batch) = parsed.batch else {
104            return Ok(Backpressure::Finished);
105        };
106        if !demand.accepts(batch.batch()) {
107            self.tokens.restore(token_checkpoint);
108            self.parser.restore(parser_checkpoint);
109            let required = batch.batch().retained_bytes();
110            self.pool.recycle(batch, reservation);
111            return Err(StructureBatchError::DemandTooSmall {
112                required,
113                available: demand.max_bytes,
114            });
115        }
116        let Some(next_row) = self
117            .logical_row
118            .get()
119            .checked_add(u64::from(parsed.rows))
120            .map(LogicalRow::new)
121        else {
122            self.tokens.restore(token_checkpoint);
123            self.parser.restore(parser_checkpoint);
124            self.pool.recycle(batch, reservation);
125            return Err(identity_overflow());
126        };
127        let Some(next_chunk) = self.chunk.next() else {
128            self.tokens.restore(token_checkpoint);
129            self.parser.restore(parser_checkpoint);
130            self.pool.recycle(batch, reservation);
131            return Err(identity_overflow());
132        };
133        let lease = match self.pool.lease(batch, reservation) {
134            Ok(lease) => lease,
135            Err(error) => {
136                self.tokens.restore(token_checkpoint);
137                self.parser.restore(parser_checkpoint);
138                return Err(error);
139            }
140        };
141        self.logical_row = next_row;
142        self.chunk = next_chunk;
143        self.continuation = parsed.after;
144        Ok(Backpressure::Ready(lease))
145    }
146}
147
148struct ParsedBatch {
149    batch: Option<StructureBatchBuffer>,
150    rows: u32,
151    after: ContinuityLevel,
152}
153
154fn parse_batch<S: SourceBytes>(
155    tokens: &mut TokenCursor<S>,
156    parser: &mut ParserState,
157    options: &ReadOptions,
158    mut builder: StructureBatchBuilder,
159    capacity: u32,
160    before: ContinuityLevel,
161) -> Result<ParsedBatch, StructureBatchError> {
162    let mut rows = 0u32;
163    while rows < capacity && !parser.finished {
164        let Some(token) = tokens.next_token()? else {
165            parser.finished = true;
166            break;
167        };
168        if parser.token(token, options) && parser.row.accepted(options) {
169            builder.push(parser.row.record(options))?;
170            rows += 1;
171        }
172    }
173    if rows == 0 {
174        return Ok(ParsedBatch {
175            batch: None,
176            rows,
177            after: ContinuityLevel::None,
178        });
179    }
180    let after = if parser.finished {
181        ContinuityLevel::None
182    } else {
183        ContinuityLevel::Residue
184    };
185    builder.set_continuity(BatchContinuity { before, after });
186    Ok(ParsedBatch {
187        batch: Some(builder.finish()?),
188        rows,
189        after,
190    })
191}
192
193#[derive(Clone, Copy, Debug, PartialEq, Eq)]
194enum Mode {
195    Search,
196    Tags,
197    Rows,
198}
199
200#[derive(Debug)]
201struct ParserState {
202    mode: Mode,
203    columns: Vec<Option<Field>>,
204    column: usize,
205    row: AtomRow,
206    finished: bool,
207}
208
209#[derive(Clone, Copy, Debug)]
210struct ParserCheckpoint {
211    mode: Mode,
212    column: usize,
213    finished: bool,
214}
215
216impl ParserState {
217    fn new() -> Self {
218        Self {
219            mode: Mode::Search,
220            columns: Vec::new(),
221            column: 0,
222            row: AtomRow::new(),
223            finished: false,
224        }
225    }
226
227    fn checkpoint(&self) -> ParserCheckpoint {
228        ParserCheckpoint {
229            mode: self.mode,
230            column: self.column,
231            finished: self.finished,
232        }
233    }
234
235    fn restore(&mut self, checkpoint: ParserCheckpoint) {
236        self.mode = checkpoint.mode;
237        self.column = checkpoint.column;
238        self.finished = checkpoint.finished;
239    }
240
241    fn token(&mut self, token: &str, options: &ReadOptions) -> bool {
242        match self.mode {
243            Mode::Search => {
244                if token.eq_ignore_ascii_case("loop_") {
245                    self.mode = Mode::Tags;
246                    self.columns.clear();
247                }
248                false
249            }
250            Mode::Tags => self.tag_or_value(token, options),
251            Mode::Rows => self.value(token, options),
252        }
253    }
254
255    fn tag_or_value(&mut self, token: &str, options: &ReadOptions) -> bool {
256        if token.starts_with('_') {
257            self.columns.push(Field::from_tag(token));
258            return false;
259        }
260        if !self.columns.iter().any(Option::is_some) {
261            self.mode = Mode::Search;
262            return false;
263        }
264        self.mode = Mode::Rows;
265        self.value(token, options)
266    }
267
268    fn value(&mut self, token: &str, _options: &ReadOptions) -> bool {
269        if self.column == 0 {
270            if is_control(token) {
271                self.finished = true;
272                return false;
273            }
274            self.row.clear();
275        }
276        let field = self.columns.get(self.column).copied().flatten();
277        if let Some(field) = field {
278            self.row.set(field, token);
279        }
280        self.column += 1;
281        if self.column < self.columns.len() {
282            return false;
283        }
284        self.column = 0;
285        true
286    }
287}
288
289#[derive(Clone, Copy, Debug)]
290enum Field {
291    Group,
292    Id,
293    Element,
294    Atom,
295    Alternate,
296    Component,
297    Chain,
298    Sequence,
299    Insertion,
300    X,
301    Y,
302    Z,
303    Occupancy,
304    BFactor,
305    Charge,
306    Model,
307}
308
309impl Field {
310    fn from_tag(tag: &str) -> Option<Self> {
311        let item = tag.strip_prefix("_atom_site.")?;
312        match item {
313            "group_PDB" => Some(Self::Group),
314            "id" => Some(Self::Id),
315            "type_symbol" => Some(Self::Element),
316            "label_atom_id" => Some(Self::Atom),
317            "label_alt_id" => Some(Self::Alternate),
318            "label_comp_id" => Some(Self::Component),
319            "label_asym_id" => Some(Self::Chain),
320            "label_seq_id" => Some(Self::Sequence),
321            "pdbx_PDB_ins_code" => Some(Self::Insertion),
322            "Cartn_x" => Some(Self::X),
323            "Cartn_y" => Some(Self::Y),
324            "Cartn_z" => Some(Self::Z),
325            "occupancy" => Some(Self::Occupancy),
326            "B_iso_or_equiv" => Some(Self::BFactor),
327            "pdbx_formal_charge" => Some(Self::Charge),
328            "pdbx_PDB_model_num" => Some(Self::Model),
329            _ => None,
330        }
331    }
332}
333
334#[derive(Debug)]
335struct AtomRow {
336    text: [String; 7],
337    integer: [i64; 4],
338    real: [f64; 5],
339    present: u32,
340}
341
342impl AtomRow {
343    fn new() -> Self {
344        Self {
345            text: core::array::from_fn(|_| String::new()),
346            integer: [0; 4],
347            real: [0.0; 5],
348            present: 0,
349        }
350    }
351
352    fn clear(&mut self) {
353        self.present = 0;
354    }
355
356    fn set(&mut self, field: Field, token: &str) {
357        match field {
358            Field::Group => self.set_text(0, token),
359            Field::Element => self.set_text(1, token),
360            Field::Atom => self.set_text(2, token),
361            Field::Alternate => self.set_text(3, token),
362            Field::Component => self.set_text(4, token),
363            Field::Chain => self.set_text(5, token),
364            Field::Insertion => self.set_text(6, token),
365            Field::Id => self.set_integer(0, token),
366            Field::Sequence => self.set_integer(1, token),
367            Field::Charge => self.set_integer(2, token),
368            Field::Model => self.set_integer(3, token),
369            Field::X => self.set_real(0, token),
370            Field::Y => self.set_real(1, token),
371            Field::Z => self.set_real(2, token),
372            Field::Occupancy => self.set_real(3, token),
373            Field::BFactor => self.set_real(4, token),
374        }
375    }
376
377    fn set_text(&mut self, index: usize, token: &str) {
378        self.text[index].clear();
379        if !matches!(token, "." | "?") {
380            self.text[index].push_str(token);
381            self.present |= 1 << index;
382        }
383    }
384
385    fn set_integer(&mut self, index: usize, token: &str) {
386        if let Ok(value) = token.parse() {
387            self.integer[index] = value;
388            self.present |= 1 << (7 + index);
389        }
390    }
391
392    fn set_real(&mut self, index: usize, token: &str) {
393        if let Ok(value) = token.parse() {
394            self.real[index] = value;
395            self.present |= 1 << (11 + index);
396        }
397    }
398
399    fn accepted(&self, options: &ReadOptions) -> bool {
400        let element = Element::from_symbol(&self.text[1]);
401        let hydrogen = element.is_some_and(Element::is_hydrogen);
402        let model = integer_i32(&self.integer, self.present, 3, 1);
403        !(options.discard_hydrogens && hydrogen || options.only_first_model && model != 1)
404    }
405
406    fn record(&self, _options: &ReadOptions) -> StructureAtomRecord<'_> {
407        let position = if self.present & (0b111 << 11) == 0b111 << 11 {
408            self.real[0]
409                .to_f32()
410                .zip(self.real[1].to_f32())
411                .zip(self.real[2].to_f32())
412                .map(|((x, y), z)| [x, y, z])
413        } else {
414            None
415        };
416        StructureAtomRecord {
417            model: integer_i32(&self.integer, self.present, 3, 1),
418            chain: &self.text[5],
419            component: &self.text[4],
420            sequence: integer_i32(&self.integer, self.present, 1, i32::MIN),
421            insertion: &self.text[6],
422            atom: &self.text[2],
423            alternate: &self.text[3],
424            element: match Element::from_symbol(&self.text[1]) {
425                Some(element) => element,
426                None => Element::UNKNOWN,
427            },
428            position,
429            occupancy: real_presence(&self.real, self.present, 3, 1.0),
430            b_factor: real_presence(&self.real, self.present, 4, 0.0),
431            formal_charge: integer_presence(&self.integer, self.present, 2),
432            atom_site_id: integer_u32(&self.integer, self.present, 0),
433            heterogen: self.text[0].eq_ignore_ascii_case("HETATM"),
434        }
435    }
436}
437
438fn integer_i32(values: &[i64; 4], present: u32, index: usize, absent: i32) -> i32 {
439    if present & (1 << (7 + index)) == 0 {
440        return absent;
441    }
442    match i32::try_from(values[index]) {
443        Ok(value) => value,
444        Err(_) => absent,
445    }
446}
447
448fn integer_u32(values: &[i64; 4], present: u32, index: usize) -> u32 {
449    if present & (1 << (7 + index)) == 0 {
450        return 0;
451    }
452    let Ok(value) = u32::try_from(values[index]) else {
453        return 0;
454    };
455    value
456}
457
458fn integer_presence(values: &[i64; 4], present: u32, index: usize) -> (i8, Presence) {
459    if present & (1 << (7 + index)) == 0 {
460        return (0, Presence::Unknown);
461    }
462    match i8::try_from(values[index]) {
463        Ok(value) => (value, Presence::Present),
464        Err(_) => (0, Presence::Unknown),
465    }
466}
467
468fn real_presence(values: &[f64; 5], present: u32, index: usize, absent: f32) -> (f32, Presence) {
469    if present & (1 << (11 + index)) == 0 {
470        return (absent, Presence::Unknown);
471    }
472    match values[index].to_f32() {
473        Some(value) => (value, Presence::Present),
474        None => (absent, Presence::Unknown),
475    }
476}
477
478fn is_control(token: &str) -> bool {
479    token.starts_with('_')
480        || token.eq_ignore_ascii_case("loop_")
481        || token.starts_with("data_")
482        || token.starts_with("save_")
483        || token.eq_ignore_ascii_case("stop_")
484}
485
486#[cfg(test)]
487#[path = "batch_tests.rs"]
488mod tests;
molframe_cif/batch.rs

molframe_cif/
batch.rs
