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, ExecutionContext, LogicalRow, Presence, ReadOptions, SourceBytes,
10    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 model = integer_i32(&self.integer, self.present, 3, 1);
401        StructureAtomRecord::admitted(options, &self.text[1], model)
402    }
403
404    fn record(&self, _options: &ReadOptions) -> StructureAtomRecord<'_> {
405        let recorded =
406            |index: usize| (self.present & (1 << (11 + index)) != 0).then(|| self.real[index]);
407        let position = StructureAtomRecord::position_of(recorded(0), recorded(1), recorded(2));
408        StructureAtomRecord {
409            model: integer_i32(&self.integer, self.present, 3, 1),
410            chain: &self.text[5],
411            component: &self.text[4],
412            sequence: integer_i32(&self.integer, self.present, 1, i32::MIN),
413            insertion: &self.text[6],
414            atom: &self.text[2],
415            alternate: &self.text[3],
416            element: StructureAtomRecord::element_of(&self.text[1]),
417            position,
418            occupancy: real_presence(&self.real, self.present, 3, 1.0),
419            b_factor: real_presence(&self.real, self.present, 4, 0.0),
420            formal_charge: integer_presence(&self.integer, self.present, 2),
421            atom_site_id: integer_u32(&self.integer, self.present, 0),
422            heterogen: self.text[0].eq_ignore_ascii_case("HETATM"),
423        }
424    }
425}
426
427fn integer_i32(values: &[i64; 4], present: u32, index: usize, absent: i32) -> i32 {
428    if present & (1 << (7 + index)) == 0 {
429        return absent;
430    }
431    match i32::try_from(values[index]) {
432        Ok(value) => value,
433        Err(_) => absent,
434    }
435}
436
437fn integer_u32(values: &[i64; 4], present: u32, index: usize) -> u32 {
438    if present & (1 << (7 + index)) == 0 {
439        return 0;
440    }
441    let Ok(value) = u32::try_from(values[index]) else {
442        return 0;
443    };
444    value
445}
446
447fn integer_presence(values: &[i64; 4], present: u32, index: usize) -> (i8, Presence) {
448    if present & (1 << (7 + index)) == 0 {
449        return (0, Presence::Unknown);
450    }
451    match i8::try_from(values[index]) {
452        Ok(value) => (value, Presence::Present),
453        Err(_) => (0, Presence::Unknown),
454    }
455}
456
457fn real_presence(values: &[f64; 5], present: u32, index: usize, absent: f32) -> (f32, Presence) {
458    if present & (1 << (11 + index)) == 0 {
459        return (absent, Presence::Unknown);
460    }
461    match values[index].to_f32() {
462        Some(value) => (value, Presence::Present),
463        None => (absent, Presence::Unknown),
464    }
465}
466
467fn is_control(token: &str) -> bool {
468    token.starts_with('_')
469        || token.eq_ignore_ascii_case("loop_")
470        || token.starts_with("data_")
471        || token.starts_with("save_")
472        || token.eq_ignore_ascii_case("stop_")
473}
474
475#[cfg(test)]
476#[path = "batch_tests.rs"]
477mod tests;
molframe_cif/batch.rs

molframe_cif/
batch.rs
