Skip to main content

readcon_core/
iterators.rs

1//=============================================================================
2// The Public API - A clean iterator for users of our library
3//=============================================================================
4
5use crate::parser::{LineStream, parse_declared_sections};
6use crate::{error, types};
7use std::path::Path;
8
9/// memchr-backed line cursor for the full parse path (not only frame skip).
10///
11/// Profile note: `str::Lines` + `Peekable` showed up under
12/// `ConFrameIterator::next` on multi-atom multi-frame workloads. One cursor
13/// serves `next` / `peek` / `forward_fast` so skip and full parse share the
14/// same O(1) newline scan rather than two desynchronized views of the buffer.
15pub struct MemchrLines<'a> {
16    bytes: &'a [u8],
17    pos: usize,
18    peeked: Option<&'a str>,
19    /// When the buffer was not already a `&str`, header lines are checked
20    /// as they are tokenized. Coordinate rows stay on the byte parser.
21    check_utf8: bool,
22    utf8_bad: bool,
23}
24
25impl<'a> MemchrLines<'a> {
26    pub fn new(text: &'a str) -> Self {
27        Self::from_parts(text.as_bytes(), false)
28    }
29
30    /// Cursor over bytes that are not yet known to be UTF-8.
31    ///
32    /// ASCII coordinate rows are accepted by the numeric scan. Text lines
33    /// are validated when they are materialized, so there is no separate
34    /// whole-buffer UTF-8 pass.
35    pub fn from_bytes(bytes: &'a [u8]) -> Self {
36        Self::from_parts(bytes, true)
37    }
38
39    fn from_parts(bytes: &'a [u8], check_utf8: bool) -> Self {
40        Self {
41            bytes,
42            pos: 0,
43            peeked: None,
44            check_utf8,
45            utf8_bad: false,
46        }
47    }
48
49    #[inline]
50    fn read_one(&mut self) -> Option<&'a str> {
51        if self.utf8_bad || self.pos >= self.bytes.len() {
52            return None;
53        }
54        let rest = &self.bytes[self.pos..];
55        let (line_bytes, advance) = match memchr::memchr(b'\n', rest) {
56            Some(i) => (&rest[..i], i + 1),
57            None => (rest, rest.len()),
58        };
59        let trimmed = if line_bytes.last() == Some(&b'\r') {
60            &line_bytes[..line_bytes.len() - 1]
61        } else {
62            line_bytes
63        };
64        if self.check_utf8 {
65            match std::str::from_utf8(trimmed) {
66                Ok(s) => {
67                    self.pos += advance;
68                    Some(s)
69                }
70                Err(_) => {
71                    // Stay on the rejected line. Moving past it makes a
72                    // whitespace tail look like a clean end of file.
73                    self.utf8_bad = true;
74                    None
75                }
76            }
77        } else {
78            self.pos += advance;
79            // SAFETY: source was `&str`; line is a UTF-8 prefix cut on ASCII `\n`/`\r`.
80            Some(unsafe { std::str::from_utf8_unchecked(trimmed) })
81        }
82    }
83
84    #[inline]
85    pub fn next_line(&mut self) -> Option<&'a str> {
86        if let Some(p) = self.peeked.take() {
87            return Some(p);
88        }
89        self.read_one()
90    }
91
92    #[inline]
93    pub fn peek_line(&mut self) -> Option<&'a str> {
94        if self.peeked.is_none() {
95            self.peeked = self.read_one();
96        }
97        self.peeked
98    }
99
100    /// Drop any peek buffer (required before bulk cursor advances).
101    fn clear_peek(&mut self) {
102        if let Some(p) = self.peeked.take() {
103            // Rewind pos to the start of the peeked line.
104            let start = p.as_ptr() as usize - self.bytes.as_ptr() as usize;
105            self.pos = start;
106        }
107    }
108
109    fn utf8_error(&self) -> Option<error::ParseError> {
110        self.utf8_bad.then(invalid_utf8)
111    }
112
113    /// Keep `err` unless this cursor already rejected a line as UTF-8.
114    fn prefer_utf8(&self, err: error::ParseError) -> error::ParseError {
115        if self.utf8_bad { invalid_utf8() } else { err }
116    }
117}
118
119fn invalid_utf8() -> error::ParseError {
120    error::ParseError::ValidationError("invalid UTF-8".into())
121}
122
123impl<'a> Iterator for MemchrLines<'a> {
124    type Item = &'a str;
125    fn next(&mut self) -> Option<&'a str> {
126        self.next_line()
127    }
128}
129
130impl<'a> LineStream<'a> for MemchrLines<'a> {
131    #[inline]
132    fn next_line(&mut self) -> Option<&'a str> {
133        MemchrLines::next_line(self)
134    }
135    #[inline]
136    fn peek_line(&mut self) -> Option<&'a str> {
137        MemchrLines::peek_line(self)
138    }
139    #[inline]
140    fn next_coord_row(
141        &mut self,
142        spec_version: u32,
143        default_id: u64,
144    ) -> Result<([f64; 3], [bool; 3], u64), crate::error::ParseError> {
145        if let Some(p) = self.peeked.take() {
146            return crate::parser::parse_coord_row_at(p.as_bytes(), 0, spec_version, default_id)
147                .map(|(row, _)| row);
148        }
149        let (row, pos) =
150            crate::parser::parse_coord_row_at(self.bytes, self.pos, spec_version, default_id)?;
151        self.pos = pos;
152        Ok(row)
153    }
154    fn remainder_is_whitespace(&mut self) -> bool {
155        if self.utf8_bad {
156            return false;
157        }
158        if let Some(line) = self.peek_line()
159            && !line.trim().is_empty()
160        {
161            return false;
162        }
163        // `peek_line` may have rejected a line. That is a decode error,
164        // including when every byte after it is ASCII whitespace.
165        if self.utf8_bad {
166            return false;
167        }
168        self.bytes[self.pos..]
169            .iter()
170            .all(|b| b.is_ascii_whitespace())
171    }
172    fn blank_line_opens_section(&mut self) -> bool {
173        // `peek_line` has already moved `pos` to the line after the peek.
174        let Some(peeked) = self.peeked else {
175            return false;
176        };
177        if !peeked.trim().is_empty() {
178            return false;
179        }
180        match probe_section(&self.bytes[self.pos..]) {
181            SectionProbe::Opens => true,
182            SectionProbe::Closed => false,
183            // Leave the rejected bytes on the cursor. `prefer_utf8` then
184            // reports the decode error instead of a missing section.
185            SectionProbe::InvalidUtf8 => {
186                self.utf8_bad = true;
187                false
188            }
189        }
190    }
191
192    fn blank_opens_velocity_block(&mut self) -> bool {
193        match self.peek_line() {
194            Some(line) if line.trim().is_empty() => self.blank_line_opens_section(),
195            _ => true,
196        }
197    }
198}
199
200fn split_raw_line(bytes: &[u8]) -> Option<(&[u8], &[u8])> {
201    if bytes.is_empty() {
202        return None;
203    }
204    match memchr::memchr(b'\n', bytes) {
205        Some(i) => Some((&bytes[..i], &bytes[i + 1..])),
206        None => Some((bytes, &[])),
207    }
208}
209
210fn strip_cr(line: &[u8]) -> &[u8] {
211    if line.last() == Some(&b'\r') {
212        &line[..line.len() - 1]
213    } else {
214        line
215    }
216}
217
218enum SectionProbe {
219    Opens,
220    Closed,
221    InvalidUtf8,
222}
223
224/// A blank line opens a section when the next line is a symbol and the
225/// line after that is a component label (CON-file.5, CON-header.4).
226fn probe_section(after_blank: &[u8]) -> SectionProbe {
227    let Some((raw1, after1)) = split_raw_line(after_blank) else {
228        return SectionProbe::Closed;
229    };
230    let line1 = strip_cr(raw1);
231    if line1.iter().all(|b| matches!(b, b' ' | b'\t')) {
232        return SectionProbe::Closed;
233    }
234    let Ok(text1) = std::str::from_utf8(line1) else {
235        return SectionProbe::InvalidUtf8;
236    };
237    if text1.trim_start().starts_with('{') {
238        return SectionProbe::Closed;
239    }
240    let Some((raw2, _)) = split_raw_line(after1) else {
241        return SectionProbe::Closed;
242    };
243    let Ok(text2) = std::str::from_utf8(strip_cr(raw2)) else {
244        return SectionProbe::InvalidUtf8;
245    };
246    if crate::parser::line_has_component_label(text2) {
247        SectionProbe::Opens
248    } else {
249        SectionProbe::Closed
250    }
251}
252
253fn blank_opens_section(after_blank: &[u8]) -> bool {
254    matches!(probe_section(after_blank), SectionProbe::Opens)
255}
256
257/// An iterator that lazily parses simulation frames from a `.con` or `.convel`
258/// file's contents.
259///
260/// This struct wraps a memchr line cursor over the file buffer and, upon each
261/// iteration, attempts to parse a complete `ConFrame`. A blank line after the
262/// coordinates is a velocity section only when it begins a complete velocity
263/// block. Otherwise that line is line 1 of the next frame, and it may be empty.
264///
265/// The iterator yields items of type `Result<ConFrame, ParseError>`, allowing for
266/// robust error handling for each frame.
267pub struct ConFrameIterator<'a> {
268    pub(crate) lines: MemchrLines<'a>,
269}
270
271impl<'a> ConFrameIterator<'a> {
272    /// Creates a new `ConFrameIterator` from a string slice of the entire file.
273    ///
274    /// # Arguments
275    ///
276    /// * `file_contents` - A string slice containing the text of one or more `.con` frames.
277    pub fn new(file_contents: &'a str) -> Self {
278        ConFrameIterator {
279            lines: MemchrLines::new(file_contents),
280        }
281    }
282
283    /// Like [`Self::new`] for a buffer that has not been UTF-8 validated.
284    ///
285    /// Text lines are checked as they are tokenized. Numeric rows reject a
286    /// non-ASCII byte during the fused parse, which is the UTF-8 check for
287    /// those lines.
288    pub fn from_bytes(bytes: &'a [u8]) -> Self {
289        ConFrameIterator {
290            lines: MemchrLines::from_bytes(bytes),
291        }
292    }
293
294    /// Bulk-skips `n` lines from the shared memchr cursor.
295    ///
296    /// A final line with no trailing newline still counts when it is the
297    /// last line this call needs and its bytes are present. The full
298    /// parser accepts that line. Further lines after it are missing.
299    fn advance_lines(&mut self, n: usize) -> Result<(), error::ParseError> {
300        self.lines.clear_peek();
301        for i in 0..n {
302            let rest = &self.lines.bytes[self.lines.pos..];
303            match memchr::memchr(b'\n', rest) {
304                Some(pos) => self.lines.pos += pos + 1,
305                None => {
306                    if rest.is_empty() {
307                        return Err(error::ParseError::IncompleteFrame);
308                    }
309                    self.lines.pos = self.lines.bytes.len();
310                    if i + 1 == n {
311                        return Ok(());
312                    }
313                    return Err(error::ParseError::IncompleteFrame);
314                }
315            }
316        }
317        Ok(())
318    }
319
320    /// One line from the shared cursor (same as full-parse path).
321    fn read_line_str(&mut self) -> Option<&'a str> {
322        self.lines.clear_peek();
323        self.lines.next_line()
324    }
325
326    /// memchr-backed equivalent of [`Self::forward`]. Skips the next
327    /// frame without fully parsing its atom data. Shares the same line
328    /// cursor as [`Iterator::next`], so skip and full parse interleave safely.
329    pub fn forward_fast(&mut self) -> Option<Result<(), error::ParseError>> {
330        self.lines.clear_peek();
331        if self.lines.pos >= self.lines.bytes.len() {
332            return None;
333        }
334        // Lines 1..=6 of the header are skipped wholesale.
335        if let Err(e) = self.advance_lines(6) {
336            return Some(Err(e));
337        }
338        // Line 7: natm_types.
339        let natm_types: usize = match self.read_line_str() {
340            Some(line) => match crate::parser::parse_count_line(line, 1) {
341                Ok(v) => v[0],
342                Err(e) => return Some(Err(e)),
343            },
344            None => return Some(Err(error::ParseError::IncompleteHeader)),
345        };
346        if natm_types == 0 {
347            // Lines 8 and 9 are whitespace, or absent at EOF. Do not scan
348            // for coordinates or a legacy velocity block (CON-count.1).
349            return Some(crate::parser::read_zero_atom_tail(|| self.read_line_str()));
350        }
351        // Line 8: natms_per_type.
352        let natms_per_type: Vec<usize> = match self.read_line_str() {
353            Some(line) => match crate::parser::parse_count_line(line, natm_types) {
354                Ok(v) => v,
355                Err(e) => return Some(Err(e)),
356            },
357            None => return Some(Err(error::ParseError::IncompleteHeader)),
358        };
359        if natms_per_type.iter().any(|n| *n == 0) {
360            return Some(Err(error::ParseError::InvalidCount(
361                "each per-type count must be >= 1".into(),
362            )));
363        }
364        // Line 9: masses_per_type, consumed.
365        if let Err(e) = self.advance_lines(1) {
366            return Some(Err(e));
367        }
368        let total_atoms: usize = natms_per_type.iter().sum();
369        let coord_block_lines = total_atoms + natm_types * 2;
370        if let Err(e) = self.advance_lines(coord_block_lines) {
371            return Some(Err(e));
372        }
373        // Optional sections: blank line + same-shape block, repeated.
374        self.lines.clear_peek();
375        loop {
376            let rest = &self.lines.bytes[self.lines.pos..];
377            if rest.is_empty() {
378                break;
379            }
380            let next_eol = memchr::memchr(b'\n', rest);
381            let line = match next_eol {
382                Some(pos) => &rest[..pos],
383                None => rest,
384            };
385            let is_blank = line.iter().all(|b| matches!(b, b' ' | b'\t' | b'\r'));
386            if !is_blank {
387                break;
388            }
389            let after = match next_eol {
390                Some(pos) => &rest[pos + 1..],
391                None => &[][..],
392            };
393            // A blank line and then only whitespace is the end of the file,
394            // not a section (CON-frame.5).
395            if after.iter().all(|b| b.is_ascii_whitespace()) {
396                self.lines.pos = self.lines.bytes.len();
397                break;
398            }
399            // A blank that does not begin a section is line 1 of the next
400            // frame, including an empty user line (CON-frame.7, CON-header.4).
401            if !blank_opens_section(after) {
402                break;
403            }
404            // Consume the blank separator and the section block.
405            self.lines.pos += next_eol.map(|p| p + 1).unwrap_or(rest.len());
406            if let Err(e) = self.advance_lines(coord_block_lines) {
407                return Some(Err(e));
408            }
409        }
410        Some(Ok(()))
411    }
412
413    /// Skips the next frame without fully parsing its atomic data.
414    ///
415    /// This is more efficient than `next()` if you only need to advance the
416    /// iterator. It reads the frame's header to determine how many lines to skip,
417    /// including any velocity section if present.
418    ///
419    /// # Returns
420    ///
421    /// * `Some(Ok(()))` on a successful skip.
422    /// * `Some(Err(ParseError::...))` if there's an error parsing the header.
423    /// * `None` if the iterator is already at the end.
424    pub fn forward(&mut self) -> Option<Result<(), error::ParseError>> {
425        // Prefer the shared memchr skip path (same cursor as full parse).
426        self.forward_fast()
427    }
428
429    /// Byte offset of the next frame start in the buffer passed to [`Self::new`].
430    pub fn byte_pos(&self) -> usize {
431        self.lines.pos
432    }
433
434    /// Skip `n` frames without parsing atom data. Returns how many were skipped.
435    /// Named `skip_frames` so it does not collide with [`Iterator::skip`].
436    pub fn skip_frames(&mut self, n: usize) -> Result<usize, error::ParseError> {
437        let mut skipped = 0usize;
438        for _ in 0..n {
439            match self.forward() {
440                None => break,
441                Some(Err(e)) => return Err(e),
442                Some(Ok(())) => skipped += 1,
443            }
444        }
445        Ok(skipped)
446    }
447
448    /// Next frame plus the exact substring of the buffer passed to [`Self::new`].
449    ///
450    /// **Corpus ingest contract:** successive successful spans from the same
451    /// `file_contents` are contiguous (`end` of frame *i* equals `start` of frame
452    /// *i+1*) and, for a buffer that is only multi-frame CON (no prefix garbage),
453    /// concatenating all spans reproduces the trajectory text. Campaign stores
454    /// (`readcon-db`) must persist these spans as authoritative blobs—do not
455    /// re-serialize on the hot ingest path unless the caller supplied in-memory
456    /// [`types::ConFrame`] values without source text.
457    ///
458    /// See also [`crate::index_proj::frame_byte_spans`] and
459    /// [`crate::index_proj::spans_cover_buffer`].
460    pub fn next_with_raw_span(
461        &mut self,
462        file_contents: &'a str,
463    ) -> Option<Result<(types::ConFrame, &'a str), error::ParseError>> {
464        let base = file_contents.as_ptr() as usize;
465        let start = {
466            let line = self.lines.peek_line()?;
467            line.as_ptr() as usize - base
468        };
469        let frame = match self.next()? {
470            Ok(f) => f,
471            Err(e) => return Some(Err(e)),
472        };
473        let end = match self.lines.peek_line() {
474            Some(line) => line.as_ptr() as usize - base,
475            None => file_contents.len(),
476        };
477        debug_assert!(end >= start && end <= file_contents.len());
478        Some(Ok((frame, &file_contents[start..end])))
479    }
480}
481
482impl<'a> Iterator for ConFrameIterator<'a> {
483    /// The type of item yielded by the iterator.
484    ///
485    /// Each item is a `Result` that contains a successfully parsed `ConFrame` or a
486    /// `ParseError` if the frame's data is malformed.
487    type Item = Result<types::ConFrame, error::ParseError>;
488
489    /// Advances the iterator and attempts to parse the next frame.
490    ///
491    /// This method will return `None` only when there are no more lines to consume.
492    /// If there are lines but they do not form a complete frame, it will return
493    /// `Some(Err(ParseError::...))`. Invalid UTF-8 is a validation error, not
494    /// the end of the stream.
495    fn next(&mut self) -> Option<Self::Item> {
496        // A line rejected earlier in this cursor (for example while probing
497        // for a velocity block) must not be reported as a clean EOF.
498        if let Some(err) = self.lines.utf8_error() {
499            return Some(Err(err));
500        }
501        // Trailing whitespace after the last complete frame is not a frame
502        // (CON-frame.1, D2). An empty line that is followed by more text is
503        // line 1 of the next frame (CON-frame.7).
504        if self.lines.remainder_is_whitespace() {
505            if let Some(err) = self.lines.utf8_error() {
506                return Some(Err(err));
507            }
508            while self.lines.next_line().is_some() {}
509            return None;
510        }
511        if let Some(err) = self.lines.utf8_error() {
512            return Some(Err(err));
513        }
514        // Otherwise, attempt to parse the next frame from the available lines.
515        let mut frame = match crate::parser::parse_single_frame_ls(&mut self.lines) {
516            Ok(f) => f,
517            Err(e) => return Some(Err(self.lines.prefer_utf8(e))),
518        };
519        // Optional sections mutate AoS; only re-sync section SoA when needed.
520        // Plain .con assembly already filled positions/ids/masses (no O(N)
521        // post-scan when no velocity/force sections were applied).
522        let sections =
523            match parse_declared_sections(&mut self.lines, &mut frame.header, &mut frame.atom_data)
524            {
525                Ok(n) => n,
526                // A section line that failed UTF-8 makes `next_line` return
527                // `None`. That must not surface as `IncompleteSection`.
528                Err(e) => return Some(Err(self.lines.prefer_utf8(e))),
529            };
530        if sections > 0 {
531            frame.sync_arrays_from_atom_data();
532        }
533        Some(Ok(frame))
534    }
535}
536
537#[cfg(test)]
538mod aos_soa_agreement_tests {
539    use super::*;
540    use std::path::PathBuf;
541
542    #[test]
543    fn iterator_vel_forces_soa_matches_aos() {
544        let p = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
545            .join("resources/test/tiny_cuh2_vel_forces.con");
546        let text = std::fs::read_to_string(&p).expect("fixture");
547        let fr = ConFrameIterator::new(&text)
548            .next()
549            .expect("frame")
550            .expect("parse");
551        let n = fr.atom_data.len();
552        assert!(n > 0);
553        assert_eq!(fr.positions.nrows(), n);
554        let has_vel = fr.atom_data.iter().any(|a| a.velocity().is_some());
555        let has_frc = fr.atom_data.iter().any(|a| a.force().is_some());
556        if has_vel {
557            assert_eq!(
558                fr.velocities.nrows(),
559                n,
560                "SoA velocities must match AoS after section parse"
561            );
562        }
563        if has_frc {
564            assert_eq!(fr.forces.nrows(), n, "SoA forces must match AoS");
565        }
566        for (i, a) in fr.atom_data.iter().enumerate() {
567            let p = fr.positions.as_f64_row(i);
568            assert_eq!([a.x, a.y, a.z], p);
569            if let Some(v) = a.velocity() {
570                assert_eq!(v, fr.velocities.as_f64_row(i));
571            }
572            if let Some(f) = a.force() {
573                assert_eq!(f, fr.forces.as_f64_row(i));
574            }
575        }
576    }
577
578    /// After SoA-primary parse, section sync must not require rewriting positions
579    /// (nrows already equals N); forces SoA still filled from AoS.
580    #[test]
581    fn sync_skips_pos_when_nrows_matches_keeps_force_soa() {
582        let p =
583            PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("resources/test/tiny_cuh2_forces.con");
584        let text = std::fs::read_to_string(&p).expect("fixture");
585        let fr = ConFrameIterator::new(&text)
586            .next()
587            .expect("frame")
588            .expect("parse");
589        let n = fr.atom_data.len();
590        assert_eq!(fr.positions.nrows(), n);
591        assert_eq!(fr.forces.nrows(), n);
592        // Snapshot first position SoA row then re-sync; coords must stay bit-identical
593        // (no needless rewrite would change nothing but we still require agreement).
594        let p0 = fr.positions.as_f64_row(0);
595        let mut fr2 = fr.clone();
596        fr2.sync_arrays_from_atom_data();
597        assert_eq!(fr2.positions.as_f64_row(0), p0);
598        assert_eq!(fr2.forces.nrows(), n);
599        assert_eq!(
600            fr2.forces.as_f64_row(0),
601            fr2.atom_data[0].force().expect("force")
602        );
603    }
604}
605
606/// Reads all frames from a file.
607///
608/// For files smaller than 64 KiB, uses a simple `read_to_string` to avoid
609/// the fixed overhead of mmap (VMA creation, page fault, munmap). For larger
610/// trajectory files, uses memory-mapped I/O to let the OS page cache handle
611/// the data.
612/// Byte-size gate for Rayon multi-frame parse. Avoids an extra O(n) frame-count
613/// scan: phase-1 of [`parse_frames_parallel`] already walks boundaries when we
614/// choose parallel. Below this size, sequential parse wins on small multi-frame
615/// files (pool scheduling overhead).
616#[cfg(feature = "parallel")]
617pub(crate) const PARALLEL_BYTES_THRESHOLD: usize = 48 * 1024;
618
619pub fn read_all_frames(path: &Path) -> Result<Vec<types::ConFrame>, Box<dyn std::error::Error>> {
620    read_all_frames_with_threads(path, None)
621}
622
623/// Collect every frame from `text`.
624///
625/// `num_threads` is `None` for the automatic policy (Rayon when the
626/// `parallel` feature is on and the buffer is at least
627/// [`PARALLEL_BYTES_THRESHOLD`]), `Some(1)` for a sequential iterator,
628/// and `Some(n)` with `n >= 2` for an explicit Rayon pool. Builds
629/// without `parallel` always parse sequentially.
630pub fn frames_from_text(
631    text: &str,
632    num_threads: Option<usize>,
633) -> Result<Vec<types::ConFrame>, error::ParseError> {
634    #[cfg(feature = "parallel")]
635    {
636        let use_parallel = match num_threads {
637            Some(1) => false,
638            Some(n) if n >= 2 => true,
639            // None and Some(0): automatic policy.
640            _ => text.len() >= PARALLEL_BYTES_THRESHOLD,
641        };
642        if use_parallel {
643            let pool = num_threads.filter(|&n| n >= 2);
644            let parts = parse_frames_parallel_with_threads(text, pool);
645            let mut frames = Vec::with_capacity(parts.len());
646            for r in parts {
647                frames.push(r?);
648            }
649            return Ok(frames);
650        }
651    }
652    #[cfg(not(feature = "parallel"))]
653    {
654        let _ = num_threads;
655    }
656    ConFrameIterator::new(text).collect()
657}
658
659/// Like [`read_all_frames`], with an explicit worker count.
660///
661/// `None` is the automatic policy used by [`read_all_frames`]. `Some(1)`
662/// forces sequential parse. `Some(n)` with `n >= 2` pins a Rayon pool
663/// (ignored when the `parallel` feature is off).
664pub fn read_all_frames_with_threads(
665    path: &Path,
666    num_threads: Option<usize>,
667) -> Result<Vec<types::ConFrame>, Box<dyn std::error::Error>> {
668    let contents = crate::compression::read_file_contents(path)?;
669    match contents {
670        crate::compression::FileContents::Owned(text) => Ok(frames_from_text(&text, num_threads)?),
671        // Mapped bytes are not pre-validated. The cursor checks UTF-8 on
672        // text lines and the coordinate scan rejects non-ASCII tokens.
673        crate::compression::FileContents::Mapped(map) => Ok(frames_from_bytes(&map, num_threads)?),
674    }
675}
676
677/// [`frames_from_text`] for bytes that may not be UTF-8 yet.
678pub fn frames_from_bytes(
679    bytes: &[u8],
680    num_threads: Option<usize>,
681) -> Result<Vec<types::ConFrame>, error::ParseError> {
682    #[cfg(feature = "parallel")]
683    {
684        let use_parallel = match num_threads {
685            Some(1) => false,
686            Some(n) if n >= 2 => true,
687            _ => bytes.len() >= PARALLEL_BYTES_THRESHOLD,
688        };
689        if use_parallel {
690            let pool = num_threads.filter(|&n| n >= 2);
691            let parts = parse_frames_parallel_bytes(bytes, pool);
692            let mut frames = Vec::with_capacity(parts.len());
693            for r in parts {
694                frames.push(r?);
695            }
696            return Ok(frames);
697        }
698    }
699    #[cfg(not(feature = "parallel"))]
700    {
701        let _ = num_threads;
702    }
703    let mut frames = Vec::new();
704    for r in ConFrameIterator::from_bytes(bytes) {
705        frames.push(r?);
706    }
707    Ok(frames)
708}
709
710#[cfg(feature = "parallel")]
711fn parse_frames_parallel_bytes(
712    bytes: &[u8],
713    num_threads: Option<usize>,
714) -> Vec<Result<types::ConFrame, error::ParseError>> {
715    use rayon::prelude::*;
716
717    let mut scanner = ConFrameIterator::from_bytes(bytes);
718    let mut boundaries = Vec::new();
719    loop {
720        scanner.lines.clear_peek();
721        let start = scanner.lines.pos;
722        // Whitespace after the final frame does not open another frame.
723        if start >= bytes.len() || bytes[start..].iter().all(u8::is_ascii_whitespace) {
724            break;
725        }
726        boundaries.push(start);
727        match scanner.forward_fast() {
728            Some(Ok(())) => {}
729            Some(Err(_)) | None => break,
730        }
731    }
732    let parse_chunks = || {
733        let num_frames = boundaries.len();
734        (0..num_frames)
735            .into_par_iter()
736            .map(|i| {
737                let start = boundaries[i];
738                let end = if i + 1 < num_frames {
739                    boundaries[i + 1]
740                } else {
741                    bytes.len()
742                };
743                let mut iter = ConFrameIterator::from_bytes(&bytes[start..end]);
744                match iter.next() {
745                    Some(result) => result,
746                    None => Err(error::ParseError::IncompleteFrame),
747                }
748            })
749            .collect()
750    };
751    match num_threads {
752        None => parse_chunks(),
753        Some(n) => {
754            let n = n.max(1);
755            match rayon::ThreadPoolBuilder::new().num_threads(n).build() {
756                Ok(pool) => pool.install(parse_chunks),
757                Err(_) => parse_chunks(),
758            }
759        }
760    }
761}
762
763/// Count frames without building atom payloads (uses [`ConFrameIterator::forward_fast`]
764/// when possible, else [`ConFrameIterator::forward`]).
765///
766/// Prefer this over `read_all_frames(...).len()` when only the frame count is needed.
767pub fn count_frames(path: &Path) -> Result<usize, Box<dyn std::error::Error>> {
768    let contents = crate::compression::read_file_contents(path)?;
769    let text = contents.as_str()?;
770    let mut n = 0usize;
771    let mut iter = ConFrameIterator::new(text);
772    loop {
773        match iter.forward_fast() {
774            Some(Ok(())) => n += 1,
775            Some(Err(e)) => return Err(Box::new(e)),
776            None => break,
777        }
778    }
779    Ok(n)
780}
781
782/// Reads only the first frame from a file.
783///
784/// More efficient than `read_all_frames` for single-frame access because it
785/// stops parsing after the first frame rather than collecting all of them.
786pub fn read_first_frame(path: &Path) -> Result<types::ConFrame, Box<dyn std::error::Error>> {
787    let contents = crate::compression::read_file_contents(path)?;
788    let text = contents.as_str()?;
789    let mut iter = ConFrameIterator::new(text);
790    match iter.next() {
791        Some(Ok(frame)) => Ok(frame),
792        Some(Err(e)) => Err(Box::new(e)),
793        None => Err("No frames found in file".into()),
794    }
795}
796
797/// Skip `index` frames (no atom parse), then parse the next one.
798///
799/// `index == 0` is [`read_first_frame`]. Out of range is an error.
800pub fn read_nth_frame(
801    path: &Path,
802    index: usize,
803) -> Result<types::ConFrame, Box<dyn std::error::Error>> {
804    let contents = crate::compression::read_file_contents(path)?;
805    let text = contents.as_str()?;
806    read_nth_frame_from_text(text, index)
807}
808
809/// Same as [`read_nth_frame`] on an already-loaded buffer.
810pub fn read_nth_frame_from_text(
811    text: &str,
812    index: usize,
813) -> Result<types::ConFrame, Box<dyn std::error::Error>> {
814    let mut iter = ConFrameIterator::new(text);
815    let skipped = iter.skip_frames(index)?;
816    if skipped < index {
817        return Err(format!("frame {index} out of range (only {skipped} frames)").into());
818    }
819    match iter.next() {
820        Some(Ok(frame)) => Ok(frame),
821        Some(Err(e)) => Err(Box::new(e)),
822        None => Err(format!("frame {index} out of range").into()),
823    }
824}
825
826/// Reads all frames using the checked `.rcidx` sidecar next to `path`.
827///
828/// Uncompressed input uses (and refreshes) `path.rcidx`; compressed
829/// input is scanned in memory each time. UTF-8 is validated in
830/// newline-aligned slices.
831#[cfg(feature = "parallel")]
832pub fn read_all_frames_indexed(
833    path: &Path,
834    num_threads: Option<usize>,
835) -> Result<Vec<types::ConFrame>, Box<dyn std::error::Error>> {
836    let contents = crate::compression::read_file_contents(path)?;
837    let text = contents.as_str_sliced()?;
838    let idx = match contents {
839        crate::compression::FileContents::Mapped(_) => {
840            crate::frame_index::load_or_build(path, text, true)?
841        }
842        crate::compression::FileContents::Owned(_) => {
843            let on_disk = std::fs::metadata(path)?.len();
844            crate::frame_index::load_or_build(path, text, on_disk == text.len() as u64)?
845        }
846    };
847    let offs: Vec<usize> = idx.offsets.iter().map(|&o| o as usize).collect();
848    parse_frames_at_offsets(text, &offs, num_threads)
849        .into_iter()
850        .map(|r| r.map_err(|e| Box::new(e) as Box<dyn std::error::Error>))
851        .collect()
852}
853
854/// Frame start byte offsets via [`ConFrameIterator::forward_fast`] (no atom parse).
855/// Stops at the first skip error (that start is still included).
856pub fn frame_start_offsets(file_contents: &str) -> Vec<usize> {
857    let mut boundaries = Vec::new();
858    let mut scanner = ConFrameIterator::new(file_contents);
859    loop {
860        scanner.lines.clear_peek();
861        let start = scanner.lines.pos;
862        if start >= scanner.lines.bytes.len() {
863            break;
864        }
865        boundaries.push(start);
866        match scanner.forward_fast() {
867            Some(Ok(())) => {}
868            Some(Err(_)) | None => break,
869        }
870    }
871    boundaries
872}
873
874/// Parses frames in parallel using rayon, splitting on frame boundaries.
875///
876/// Phase 1: sequential O(N) scan via memchr-backed
877/// [`ConFrameIterator::forward_fast`] to find byte offsets of every
878/// frame's start. The previous implementation built a `Vec<&str>` of
879/// every line and called `lines[..i].iter().map(|l| l.len() +
880/// 1).sum()` on every frame, which is O(N^2) in line count and
881/// dominated runtime on multi-frame trajectories.
882///
883/// Phase 2: parallel parse of each frame slice using rayon on the
884/// **global** Rayon pool (see also [`parse_frames_parallel_with_threads`]
885/// for strong-scaling control of the worker count).
886///
887/// Requires the `parallel` feature.
888#[cfg(feature = "parallel")]
889pub fn parse_frames_parallel(
890    file_contents: &str,
891) -> Vec<Result<types::ConFrame, error::ParseError>> {
892    parse_frames_parallel_with_threads(file_contents, None)
893}
894
895/// Like [`parse_frames_parallel`], but runs phase-2 on an explicit Rayon
896/// pool with `num_threads` workers when `Some(n)` (`n` is clamped to at
897/// least 1). `None` uses the global pool (same as [`parse_frames_parallel`]).
898///
899/// Strong-scaling tests pin worker counts without racing the global pool.
900/// Results are ordered by frame index (stable vs sequential iterator order).
901///
902/// Requires the `parallel` feature.
903#[cfg(feature = "parallel")]
904pub fn parse_frames_parallel_with_threads(
905    file_contents: &str,
906    num_threads: Option<usize>,
907) -> Vec<Result<types::ConFrame, error::ParseError>> {
908    // Phase 1: skip walk (no atom parse) for frame start offsets.
909    let boundaries = frame_start_offsets(file_contents);
910    parse_frames_at_offsets(file_contents, &boundaries, num_threads)
911}
912
913/// Phase 2 of [`parse_frames_parallel_with_threads`]: parse the frames
914/// that start at `boundaries` (strictly increasing byte offsets into
915/// `file_contents`; frame `i` ends where frame `i + 1` starts).
916///
917/// Offsets may come from [`frame_start_offsets`] or from a checked
918/// [`crate::frame_index::FrameIndex`]. Results are ordered by frame index.
919///
920/// Requires the `parallel` feature.
921#[cfg(feature = "parallel")]
922pub fn parse_frames_at_offsets(
923    file_contents: &str,
924    boundaries: &[usize],
925    num_threads: Option<usize>,
926) -> Vec<Result<types::ConFrame, error::ParseError>> {
927    use rayon::prelude::*;
928
929    let parse_chunks = || {
930        let num_frames = boundaries.len();
931        (0..num_frames)
932            .into_par_iter()
933            .map(|i| {
934                let start = boundaries[i];
935                let end = if i + 1 < num_frames {
936                    boundaries[i + 1]
937                } else {
938                    file_contents.len()
939                };
940                let chunk = &file_contents[start..end];
941                let mut iter = ConFrameIterator::new(chunk);
942                match iter.next() {
943                    Some(result) => result,
944                    None => Err(error::ParseError::IncompleteFrame),
945                }
946            })
947            .collect()
948    };
949
950    match num_threads {
951        None => parse_chunks(),
952        Some(n) => {
953            let n = n.max(1);
954            match rayon::ThreadPoolBuilder::new().num_threads(n).build() {
955                Ok(pool) => pool.install(parse_chunks),
956                Err(_) => parse_chunks(),
957            }
958        }
959    }
960}
961
962#[cfg(all(test, feature = "parallel"))]
963mod parallel_strong_scale_tests {
964    use super::*;
965    use std::path::PathBuf;
966
967    fn multi_frame_fixture() -> String {
968        let p = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("resources/test/tiny_cuh2.con");
969        let one = std::fs::read_to_string(p).expect("fixture");
970        // Enough frames for >1 worker to exercise the pool.
971        one.repeat(8)
972    }
973
974    fn sequential_frames(text: &str) -> Vec<types::ConFrame> {
975        ConFrameIterator::new(text)
976            .map(|r| r.expect("seq frame"))
977            .collect()
978    }
979
980    fn frames_payload_key(f: &types::ConFrame) -> (usize, Vec<(String, f64, f64, f64)>) {
981        let atoms: Vec<_> = f
982            .atom_data
983            .iter()
984            .map(|a| (a.symbol.to_string(), a.x, a.y, a.z))
985            .collect();
986        (f.atom_data.len(), atoms)
987    }
988
989    #[test]
990    fn parallel_workers_match_sequential_payloads() {
991        let text = multi_frame_fixture();
992        let seq = sequential_frames(&text);
993        assert!(seq.len() >= 8);
994        let seq_keys: Vec<_> = seq.iter().map(frames_payload_key).collect();
995
996        for workers in [1usize, 2, 4] {
997            let par = parse_frames_parallel_with_threads(&text, Some(workers));
998            assert_eq!(par.len(), seq.len(), "workers={workers}");
999            let par_keys: Vec<_> = par
1000                .into_iter()
1001                .map(|r| frames_payload_key(&r.expect("par frame")))
1002                .collect();
1003            assert_eq!(par_keys, seq_keys, "workers={workers} frame payloads");
1004        }
1005
1006        // Global pool path agrees too.
1007        let par_default = parse_frames_parallel(&text);
1008        let def_keys: Vec<_> = par_default
1009            .into_iter()
1010            .map(|r| frames_payload_key(&r.expect("par")))
1011            .collect();
1012        assert_eq!(def_keys, seq_keys);
1013    }
1014}