1use crate::parser::{LineStream, parse_declared_sections};
6use crate::{error, types};
7use std::path::Path;
8
9pub struct MemchrLines<'a> {
16 bytes: &'a [u8],
17 pos: usize,
18 peeked: Option<&'a str>,
19 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 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 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 fn clear_peek(&mut self) {
97 if let Some(p) = self.peeked.take() {
98 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 second_line_is_component_label(&self.bytes[self.pos..])
153 }
154}
155
156fn 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
167pub struct ConFrameIterator<'a> {
178 pub(crate) lines: MemchrLines<'a>,
179}
180
181impl<'a> ConFrameIterator<'a> {
182 pub fn new(file_contents: &'a str) -> Self {
188 ConFrameIterator {
189 lines: MemchrLines::new(file_contents),
190 }
191 }
192
193 pub fn from_bytes(bytes: &'a [u8]) -> Self {
199 ConFrameIterator {
200 lines: MemchrLines::from_bytes(bytes),
201 }
202 }
203
204 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 fn read_line_str(&mut self) -> Option<&'a str> {
225 self.lines.clear_peek();
226 self.lines.next_line()
227 }
228
229 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 if let Err(e) = self.advance_lines(6) {
239 return Some(Err(e));
240 }
241 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 return Some(crate::parser::read_zero_atom_tail(|| self.read_line_str()));
253 }
254 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 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 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 if after.iter().all(|b| b.is_ascii_whitespace()) {
299 self.lines.pos = self.lines.bytes.len();
300 break;
301 }
302 if !second_line_is_component_label(after) {
306 break;
307 }
308 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 pub fn forward(&mut self) -> Option<Result<(), error::ParseError>> {
329 self.forward_fast()
331 }
332
333 pub fn byte_pos(&self) -> usize {
335 self.lines.pos
336 }
337
338 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 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 type Item = Result<types::ConFrame, error::ParseError>;
392
393 fn next(&mut self) -> Option<Self::Item> {
399 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 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 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 #[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 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#[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
525pub 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 _ => 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
561pub 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 crate::compression::FileContents::Mapped(map) => Ok(frames_from_bytes(&map, num_threads)?),
576 }
577}
578
579pub 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 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
665pub 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
684pub 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
699pub 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
711pub 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#[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
756pub 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#[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#[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 let boundaries = frame_start_offsets(file_contents);
812 parse_frames_at_offsets(file_contents, &boundaries, num_threads)
813}
814
815#[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 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 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}