molframe_pdb/
batch.rs

1//! Pull-based fixed-column reading over bounded source windows.
2
3use crate::{fixed, hybrid36};
4use molframe_core::{
5    Backpressure, Batch, BatchContinuity, BatchDemand, BatchLease, BatchSource, ChunkId, Code,
6    ContinuityLevel, DatasetId, Diagnostic, Element, ExecutionContext, LogicalRow,
7    MissingElementPolicy, Presence, ReadOptions, SourceBytes, StructureAtomRecord, StructureBatch,
8    StructureBatchBuffer, StructureBatchBuilder, StructureBatchError, StructureBatchPool,
9};
10use num_traits::ToPrimitive;
11
12const RETAINED_BYTES_PER_ROW: usize = 128;
13const DICTIONARY_HEADROOM: usize = 64 * 1024;
14
15/// Incremental PDB rows over a seekable or resident byte source.
16#[derive(Debug)]
17pub struct PdbBatchSource<S> {
18    source: S,
19    options: ReadOptions,
20    dataset: DatasetId,
21    next_chunk: ChunkId,
22    next_row: LogicalRow,
23    cursor: u64,
24    window_bytes: usize,
25    model: i32,
26    continuation: ContinuityLevel,
27    finished: bool,
28    pool: StructureBatchPool,
29}
30
31impl<S: SourceBytes> PdbBatchSource<S> {
32    /// Creates a bounded reader. `window_bytes` is an upper bound and never
33    /// grows in response to source length.
34    ///
35    /// # Errors
36    ///
37    /// Returns [`StructureBatchError::DemandTooSmall`] for a zero-sized window.
38    pub fn new(
39        source: S,
40        options: ReadOptions,
41        dataset: DatasetId,
42        first_chunk: ChunkId,
43        first_row: LogicalRow,
44        window_bytes: usize,
45    ) -> Result<Self, StructureBatchError> {
46        if window_bytes == 0 {
47            return Err(StructureBatchError::DemandTooSmall {
48                required: 1,
49                available: 0,
50            });
51        }
52        Ok(Self {
53            source,
54            options,
55            dataset,
56            next_chunk: first_chunk,
57            next_row: first_row,
58            cursor: 0,
59            window_bytes,
60            model: 1,
61            continuation: ContinuityLevel::None,
62            finished: false,
63            pool: StructureBatchPool::new(),
64        })
65    }
66
67    fn row_capacity(demand: BatchDemand) -> Result<u32, StructureBatchError> {
68        let byte_rows =
69            demand.max_bytes.saturating_sub(DICTIONARY_HEADROOM) / RETAINED_BYTES_PER_ROW;
70        let rows = demand.max_rows.min(byte_rows).min(u32::MAX as usize);
71        u32::try_from(rows).map_err(|_| StructureBatchError::DemandTooSmall {
72            required: RETAINED_BYTES_PER_ROW,
73            available: demand.max_bytes,
74        })
75    }
76
77    fn reserve_bytes(capacity: u32, demand: BatchDemand) -> usize {
78        (capacity as usize)
79            .saturating_mul(RETAINED_BYTES_PER_ROW)
80            .saturating_add(DICTIONARY_HEADROOM)
81            .min(demand.max_bytes)
82    }
83
84    fn pull_batches(
85        &mut self,
86        demand: BatchDemand,
87        capacity: u32,
88        reserved: usize,
89        context: &ExecutionContext,
90    ) -> Result<Backpressure<BatchLease<StructureBatch>>, StructureBatchError> {
91        loop {
92            let descriptor = molframe_core::ChunkDescriptor::new(
93                self.dataset,
94                self.next_chunk,
95                self.next_row,
96                capacity,
97            )?;
98            let Some((builder, reservation)) =
99                self.pool
100                    .acquire(descriptor, reserved, demand.max_bytes, context)?
101            else {
102                return Ok(Backpressure::Pending);
103            };
104            let requested = self.window_bytes.min(demand.max_bytes.max(1));
105            let length = self.source.len_hint();
106            let window = self.source.window(self.cursor, requested)?;
107            let bytes = window.bytes();
108            if bytes.is_empty() {
109                self.finished = true;
110                return Ok(Backpressure::Finished);
111            }
112            let source_finished = length
113                .is_some_and(|length| window.end().is_ok_and(|window_end| window_end >= length));
114            let complete = complete_prefix(bytes, source_finished);
115            if complete == 0 {
116                let required = bytes.len().saturating_add(1);
117                return Err(StructureBatchError::RecordExceedsBudget {
118                    required,
119                    available: requested,
120                });
121            }
122            let model_checkpoint = self.model;
123            let parsed = match parse_window(
124                &bytes[..complete],
125                builder,
126                capacity,
127                self.continuation,
128                &self.options,
129                &mut self.model,
130            ) {
131                Ok(parsed) => parsed,
132                Err(error) => {
133                    self.model = model_checkpoint;
134                    return Err(error);
135                }
136            };
137            if parsed.rows == 0 {
138                self.pool.recycle(parsed.batch, reservation);
139                self.cursor = advance(self.cursor, parsed.consumed)?;
140                if source_finished && parsed.consumed == complete {
141                    self.finished = true;
142                    return Ok(Backpressure::Finished);
143                }
144                continue;
145            }
146            let batch = parsed.batch;
147            if !demand.accepts(batch.batch()) {
148                self.model = model_checkpoint;
149                let required = batch.batch().retained_bytes();
150                self.pool.recycle(batch, reservation);
151                return Err(StructureBatchError::DemandTooSmall {
152                    required,
153                    available: demand.max_bytes,
154                });
155            }
156            let next_cursor = advance(self.cursor, parsed.consumed)?;
157            let (next_row, next_chunk) =
158                match next_identity(self.next_row, self.next_chunk, parsed.rows) {
159                    Ok(identity) => identity,
160                    Err(error) => {
161                        self.model = model_checkpoint;
162                        self.pool.recycle(batch, reservation);
163                        return Err(error);
164                    }
165                };
166            let continuation = if parsed.ended_for_demand {
167                ContinuityLevel::Residue
168            } else {
169                ContinuityLevel::None
170            };
171            let lease = match self.pool.lease(batch, reservation) {
172                Ok(lease) => lease,
173                Err(error) => {
174                    self.model = model_checkpoint;
175                    return Err(error);
176                }
177            };
178            self.cursor = next_cursor;
179            self.next_row = next_row;
180            self.next_chunk = next_chunk;
181            self.continuation = continuation;
182            return Ok(Backpressure::Ready(lease));
183        }
184    }
185}
186
187impl<S: SourceBytes> BatchSource for PdbBatchSource<S> {
188    type Batch = StructureBatch;
189    type Error = StructureBatchError;
190
191    fn next_batch(
192        &mut self,
193        demand: BatchDemand,
194        context: &ExecutionContext,
195    ) -> Result<Backpressure<BatchLease<Self::Batch>>, Self::Error> {
196        if self.finished {
197            return Ok(Backpressure::Finished);
198        }
199        if !demand.can_accept_work() {
200            return Ok(Backpressure::Pending);
201        }
202        let capacity = Self::row_capacity(demand)?;
203        if capacity == 0 {
204            return Err(StructureBatchError::DemandTooSmall {
205                required: RETAINED_BYTES_PER_ROW + DICTIONARY_HEADROOM,
206                available: demand.max_bytes,
207            });
208        }
209        self.pull_batches(
210            demand,
211            capacity,
212            Self::reserve_bytes(capacity, demand),
213            context,
214        )
215    }
216}
217
218fn next_identity(
219    row: LogicalRow,
220    chunk: ChunkId,
221    rows: u32,
222) -> Result<(LogicalRow, ChunkId), StructureBatchError> {
223    let row = row
224        .get()
225        .checked_add(u64::from(rows))
226        .map(LogicalRow::new)
227        .ok_or_else(identity_overflow)?;
228    let chunk = chunk.next().ok_or_else(identity_overflow)?;
229    Ok((row, chunk))
230}
231
232struct ParsedWindow {
233    batch: StructureBatchBuffer,
234    consumed: usize,
235    rows: u32,
236    ended_for_demand: bool,
237}
238
239fn parse_window(
240    bytes: &[u8],
241    mut builder: StructureBatchBuilder,
242    capacity: u32,
243    continuity: ContinuityLevel,
244    options: &ReadOptions,
245    model: &mut i32,
246) -> Result<ParsedWindow, StructureBatchError> {
247    let mut consumed = 0usize;
248    let mut rows = 0u32;
249    for raw in bytes.split_inclusive(|byte| *byte == b'\n') {
250        let line_bytes = match raw.strip_suffix(b"\n") {
251            Some(value) => value,
252            None => raw,
253        };
254        let line_bytes = match line_bytes.strip_suffix(b"\r") {
255            Some(value) => value,
256            None => line_bytes,
257        };
258        let line = std::str::from_utf8(line_bytes)
259            .map_err(|_| Diagnostic::new(Code::E1201).with_message("input is not valid text"))?;
260        let record = fixed::record(line);
261        if record == "MODEL" {
262            *model = model_number_or_one(
263                fixed::integer(line, 11, 14).and_then(|value| i32::try_from(value).ok()),
264            );
265        } else if matches!(record, "ATOM" | "HETATM") {
266            if rows == capacity {
267                break;
268            }
269            builder.push(atom_record(line, *model, options))?;
270            rows += 1;
271        }
272        consumed = consumed.saturating_add(raw.len());
273    }
274    let ended_for_demand = rows == capacity && consumed < bytes.len();
275    builder.set_continuity(BatchContinuity {
276        before: continuity,
277        after: if ended_for_demand {
278            ContinuityLevel::Residue
279        } else {
280            ContinuityLevel::None
281        },
282    });
283    Ok(ParsedWindow {
284        batch: builder.finish()?,
285        consumed,
286        rows,
287        ended_for_demand,
288    })
289}
290
291fn atom_record<'a>(line: &'a str, model: i32, options: &ReadOptions) -> StructureAtomRecord<'a> {
292    let raw_atom = fixed::raw(line, 13, 16);
293    let declared = fixed::text(line, 77, 78);
294    let element = match Element::from_symbol(declared) {
295        Some(element) => element,
296        None => match options.missing_element_policy {
297            MissingElementPolicy::PreserveUnknown => Element::UNKNOWN,
298            MissingElementPolicy::InferFromAtomName => Element::infer_from_pdb_atom_name(raw_atom),
299        },
300    };
301    let position = fixed::real(line, 31, 38)
302        .and_then(|value| value.to_f32())
303        .zip(fixed::real(line, 39, 46).and_then(|value| value.to_f32()))
304        .zip(fixed::real(line, 47, 54).and_then(|value| value.to_f32()))
305        .map(|((x, y), z)| [x, y, z]);
306    let sequence =
307        hybrid36::decode(fixed::raw(line, 23, 26), 4).and_then(|value| i32::try_from(value).ok());
308    let atom_site_id =
309        hybrid36::decode(fixed::raw(line, 7, 11), 5).and_then(|value| u32::try_from(value).ok());
310    StructureAtomRecord {
311        model,
312        chain: fixed::text(line, 22, 22),
313        component: fixed::text(line, 18, 20),
314        sequence: value_or_min(sequence),
315        insertion: fixed::text(line, 27, 27),
316        atom: raw_atom.trim(),
317        alternate: fixed::text(line, 17, 17),
318        element,
319        position,
320        occupancy: optional_real(line, 55, 60, 1.0),
321        b_factor: optional_real(line, 61, 66, 0.0),
322        formal_charge: formal_charge(fixed::text(line, 79, 80)),
323        atom_site_id: identifier_or_zero(atom_site_id),
324        heterogen: fixed::record(line) == "HETATM",
325    }
326}
327
328fn value_or_min(value: Option<i32>) -> i32 {
329    let Some(value) = value else {
330        return i32::MIN;
331    };
332    value
333}
334
335fn model_number_or_one(value: Option<i32>) -> i32 {
336    let Some(value) = value else {
337        return 1;
338    };
339    value
340}
341
342fn identifier_or_zero(value: Option<u32>) -> u32 {
343    let Some(value) = value else {
344        return 0;
345    };
346    value
347}
348
349fn optional_real(line: &str, start: usize, end: usize, absent: f32) -> (f32, Presence) {
350    match fixed::real(line, start, end).and_then(|value| value.to_f32()) {
351        Some(value) => (value, Presence::Present),
352        None => (absent, Presence::Unknown),
353    }
354}
355
356fn formal_charge(text: &str) -> (i8, Presence) {
357    let bytes = text.as_bytes();
358    if bytes.len() != 2 || !bytes[0].is_ascii_digit() {
359        return (0, Presence::Inapplicable);
360    }
361    let magnitude = (bytes[0] - b'0').cast_signed();
362    match bytes[1] {
363        b'+' => (magnitude, Presence::Present),
364        b'-' => (-magnitude, Presence::Present),
365        _ => (0, Presence::Unknown),
366    }
367}
368
369fn complete_prefix(bytes: &[u8], source_finished: bool) -> usize {
370    if source_finished {
371        return bytes.len();
372    }
373    bytes
374        .iter()
375        .rposition(|byte| *byte == b'\n')
376        .map_or(0, |position| position + 1)
377}
378
379fn advance(start: u64, count: usize) -> Result<u64, StructureBatchError> {
380    let count = u64::try_from(count).map_err(|_| identity_overflow())?;
381    start.checked_add(count).ok_or_else(identity_overflow)
382}
383
384fn identity_overflow() -> StructureBatchError {
385    StructureBatchError::Diagnostic(
386        Diagnostic::new(Code::E1903).with_message("PDB batch identity exceeds u64"),
387    )
388}
389
390#[cfg(test)]
391#[path = "batch_tests.rs"]
392mod tests;
molframe_pdb/batch.rs

molframe_pdb/
batch.rs
