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