1use crate::types::{
2 ConFrame, SECTION_CHARGES, SECTION_DISPLACEMENTS, SECTION_ENERGIES, SECTION_FORCES,
3 SECTION_MAGMOMS, SECTION_SPINS, SECTION_SPREADS, SECTION_VELOCITIES, encode_fixed_bitmask,
4 meta,
5};
6use serde_json::{Value, json};
7use std::collections::BTreeMap;
8use std::fs::{File, OpenOptions};
9use std::io::{self, BufWriter, Read, Seek, SeekFrom, Write};
10use std::path::Path;
11
12#[derive(Debug)]
15pub enum WriteError {
16 Io(io::Error),
17 Format(crate::error::ParseError),
18}
19
20impl std::fmt::Display for WriteError {
21 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
22 match self {
23 WriteError::Io(e) => write!(f, "{e}"),
24 WriteError::Format(e) => write!(f, "{e}"),
25 }
26 }
27}
28
29impl std::error::Error for WriteError {
30 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
31 match self {
32 WriteError::Io(e) => Some(e),
33 WriteError::Format(e) => Some(e),
34 }
35 }
36}
37
38impl From<io::Error> for WriteError {
39 fn from(e: io::Error) -> Self {
40 WriteError::Io(e)
41 }
42}
43
44impl From<WriteError> for io::Error {
45 fn from(e: WriteError) -> Self {
46 match e {
47 WriteError::Io(e) => e,
48 WriteError::Format(err) => io::Error::new(io::ErrorKind::InvalidData, err.to_string()),
49 }
50 }
51}
52
53#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
55pub enum FloatFormat {
56 DecimalPlaces(usize),
58 #[default]
62 RoundTrip,
63}
64
65pub struct ConFrameWriter<W: Write> {
80 writer: BufWriter<W>,
81 float_format: FloatFormat,
82 canonical: bool,
86 metadata_cache: Option<MetadataCacheEntry>,
93 scratch: Vec<u8>,
96 ryu: ryu::Buffer,
98 fmt_slot: [FmtSlot; 3],
101 atom_workers: usize,
103 bound: Option<BoundHeader>,
106 stream_file: Option<File>,
109 force_map: bool,
112 keep_scratch: bool,
116 exact_scratch: bool,
118 keep_bound: bool,
122 bound_vel: Option<bool>,
123 bound_flags: Option<[bool; 8]>,
124 bound_line: Option<String>,
125 bound_no_lattice: Option<bool>,
127}
128
129struct BoundHeader {
130 meta: BTreeMap<String, Value>,
131 spec: u32,
132}
133
134#[derive(Clone, Copy)]
136struct FmtSlot {
137 bits: u64,
138 text: [u8; 24],
139 len: u8,
140 state: u8,
142}
143
144impl FmtSlot {
145 const fn new() -> Self {
146 Self {
147 bits: 0,
148 text: [0; 24],
149 len: 0,
150 state: 0,
151 }
152 }
153}
154
155#[derive(Debug)]
156struct MetadataCacheEntry {
157 spec_version: u32,
161 has_velocities: bool,
162 has_forces: bool,
163 has_energies: bool,
164 has_charges: bool,
165 has_spins: bool,
166 has_magmoms: bool,
167 has_displacements: bool,
168 has_spreads: bool,
169 metadata: std::collections::BTreeMap<String, serde_json::Value>,
170 serialized: String,
172}
173
174impl MetadataCacheEntry {
175 fn matches(
176 &self,
177 spec_version: u32,
178 sections: [bool; 8],
179 metadata: &std::collections::BTreeMap<String, serde_json::Value>,
180 ) -> bool {
181 let [
182 has_velocities,
183 has_forces,
184 has_energies,
185 has_charges,
186 has_spins,
187 has_magmoms,
188 has_displacements,
189 has_spreads,
190 ] = sections;
191 self.spec_version == spec_version
192 && self.has_velocities == has_velocities
193 && self.has_forces == has_forces
194 && self.has_energies == has_energies
195 && self.has_charges == has_charges
196 && self.has_spins == has_spins
197 && self.has_magmoms == has_magmoms
198 && self.has_displacements == has_displacements
199 && self.has_spreads == has_spreads
200 && self.metadata.len() == metadata.len()
201 && self
202 .metadata
203 .iter()
204 .zip(metadata)
205 .all(|((ka, a), (kb, b))| ka == kb && metadata_value_eq(a, b))
206 }
207}
208
209fn metadata_value_eq(a: &serde_json::Value, b: &serde_json::Value) -> bool {
212 use serde_json::Value;
213 match (a, b) {
214 (Value::Number(a), Value::Number(b)) if a.is_f64() && b.is_f64() => {
215 a.as_f64().map(f64::to_bits) == b.as_f64().map(f64::to_bits)
216 }
217 (Value::Array(a), Value::Array(b)) => {
218 a.len() == b.len() && a.iter().zip(b).all(|(a, b)| metadata_value_eq(a, b))
219 }
220 (Value::Object(a), Value::Object(b)) => {
221 a.len() == b.len()
222 && a.iter()
223 .all(|(key, a)| b.get(key).is_some_and(|b| metadata_value_eq(a, b)))
224 }
225 _ => a == b,
226 }
227}
228
229fn require_minor2_section_arity(
232 minor: u64,
233 sections: &[serde_json::Value],
234 metadata: &serde_json::Map<String, serde_json::Value>,
235) -> Result<(), WriteError> {
236 if minor < 2 {
237 return Ok(());
238 }
239 for name in sections.iter().filter_map(|value| value.as_str()) {
240 if crate::types::reserved_section_value_arity(name).is_some() {
241 continue;
242 }
243 let listed = metadata
244 .get(meta::SECTION_ARITY)
245 .and_then(|value| value.as_object())
246 .and_then(|obj| obj.get(name))
247 .and_then(|value| value.as_u64())
248 .is_some_and(|n| n >= 1);
249 if !listed {
250 return Err(WriteError::Format(
251 crate::error::ParseError::ValidationError(format!(
252 "minor >= 2 must list {name} in section_arity"
253 )),
254 ));
255 }
256 }
257 Ok(())
258}
259
260struct StreamOut {
263 file: File,
264 start: u64,
265 armed: bool,
266}
267
268impl Drop for StreamOut {
269 fn drop(&mut self) {
270 if !self.armed {
271 return;
272 }
273 let _ = self.file.set_len(self.start);
274 let _ = self.file.seek(SeekFrom::Start(self.start));
275 }
276}
277
278impl<W: Write> ConFrameWriter<W> {
280 pub fn new(writer: W) -> Self {
286 Self::with_float_format(writer, FloatFormat::default())
287 }
288
289 pub fn with_precision(writer: W, precision: usize) -> Self {
296 Self::with_float_format(writer, FloatFormat::DecimalPlaces(precision))
297 }
298
299 pub fn with_float_format(writer: W, float_format: FloatFormat) -> Self {
304 Self {
305 writer: BufWriter::new(writer),
306 float_format,
307 canonical: false,
308 metadata_cache: None,
309 scratch: Vec::with_capacity(16 * 1024),
310 ryu: ryu::Buffer::new(),
311 fmt_slot: [FmtSlot::new(); 3],
312 atom_workers: 1,
313 bound: None,
314 stream_file: None,
315 force_map: false,
316 keep_scratch: false,
317 exact_scratch: false,
318 keep_bound: false,
319 bound_vel: None,
320 bound_flags: None,
321 bound_line: None,
322 bound_no_lattice: None,
323 }
324 }
325
326 fn retain_scratch(mut self, on: bool) -> Self {
327 self.keep_scratch = on;
328 self
329 }
330
331 pub fn parallel_atoms(mut self, n: usize) -> Self {
337 self.atom_workers = n.max(1);
338 self
339 }
340
341 pub fn map_rows(mut self, on: bool) -> Self {
347 self.force_map = on;
348 self
349 }
350
351 pub fn mapped_file(mut self, file: File) -> Self {
354 self.stream_file = Some(file);
355 self
356 }
357
358 pub fn canonical(mut self, on: bool) -> Self {
363 self.set_canonical(on);
364 self
365 }
366
367 pub fn set_canonical(&mut self, on: bool) {
369 self.canonical = on;
370 if on {
371 self.metadata_cache = None;
372 }
373 }
374
375 pub fn is_canonical(&self) -> bool {
377 self.canonical
378 }
379
380 fn refresh_metadata_cache(
381 &mut self,
382 spec_version: u32,
383 metadata: &BTreeMap<String, Value>,
384 has_vel: bool,
385 has_frc: bool,
386 has_eng: bool,
387 has_chg: bool,
388 has_spn: bool,
389 has_mm: bool,
390 has_dsp: bool,
391 has_spr: bool,
392 ) -> Result<(), WriteError> {
393 let cache_hit = !self.canonical
394 && self.metadata_cache.as_ref().is_some_and(|c| {
395 c.matches(
396 spec_version,
397 [
398 has_vel, has_frc, has_eng, has_chg, has_spn, has_mm, has_dsp, has_spr,
399 ],
400 metadata,
401 )
402 });
403 if cache_hit {
404 return Ok(());
405 }
406
407 let mut meta_obj = serde_json::Map::new();
408 meta_obj.insert(meta::CON_SPEC_VERSION.into(), json!(spec_version));
409 let mut sections = Vec::new();
410 if has_vel {
411 sections.push(json!(SECTION_VELOCITIES));
412 }
413 if has_frc {
414 sections.push(json!(SECTION_FORCES));
415 }
416 if has_eng {
417 sections.push(json!(SECTION_ENERGIES));
418 }
419 if has_chg {
420 sections.push(json!(SECTION_CHARGES));
421 }
422 if has_spn {
423 sections.push(json!(SECTION_SPINS));
424 }
425 if has_mm {
426 sections.push(json!(SECTION_MAGMOMS));
427 }
428 if has_dsp {
429 sections.push(json!(SECTION_DISPLACEMENTS));
430 }
431 if has_spr {
432 sections.push(json!(SECTION_SPREADS));
433 }
434 let validate = metadata
435 .get(meta::VALIDATE)
436 .and_then(|value| value.as_bool())
437 .unwrap_or(false);
438 if !sections.is_empty() || validate {
439 meta_obj.insert(meta::SECTIONS.into(), json!(sections));
440 }
441 for (k, v) in metadata {
442 if k == meta::CON_SPEC_VERSION || k == meta::SECTIONS {
443 continue;
444 }
445 meta_obj.insert(k.clone(), v.clone());
446 }
447 if let Some(u) = meta_obj.get(meta::UNITS).cloned() {
448 let c = crate::units::canonicalize_units_object(&u).map_err(|e| {
450 WriteError::Format(crate::error::ParseError::InvalidUnits(e.to_string()))
451 })?;
452 meta_obj.insert(meta::UNITS.into(), c);
453 }
454 if spec_version >= 3 {
455 match meta_obj.get(meta::UNITS) {
456 None => {
457 return Err(WriteError::Format(crate::error::ParseError::InvalidUnits(
458 "con_spec_version >= 3 requires units.length and units.energy".into(),
459 )));
460 }
461 Some(u) => {
462 crate::units::validate_v3_units_metadata(u).map_err(|e| {
463 WriteError::Format(crate::error::ParseError::InvalidUnits(e.to_string()))
464 })?;
465 }
466 }
467 }
468 let minor = metadata
469 .get(meta::CON_SPEC_MINOR)
470 .and_then(|v| v.as_u64())
471 .unwrap_or(0);
472 if spec_version >= 3 && minor >= 1 && has_vel {
473 let time_ok = meta_obj
474 .get(meta::UNITS)
475 .and_then(|u| u.get("time"))
476 .and_then(|v| v.as_str())
477 .is_some_and(|s| !s.trim().is_empty());
478 if !time_ok {
479 return Err(WriteError::Format(crate::error::ParseError::InvalidUnits(
480 "minor >= 1 requires units.time when a velocities section is present".into(),
481 )));
482 }
483 }
484 require_minor2_section_arity(minor, §ions, &meta_obj)?;
487 let serialized = serde_json::Value::Object(meta_obj).to_string();
488 self.metadata_cache = Some(MetadataCacheEntry {
489 spec_version,
490 has_velocities: has_vel,
491 has_forces: has_frc,
492 has_energies: has_eng,
493 has_charges: has_chg,
494 has_spins: has_spn,
495 has_magmoms: has_mm,
496 has_displacements: has_dsp,
497 has_spreads: has_spr,
498 metadata: metadata.clone(),
499 serialized,
500 });
501 Ok(())
502 }
503
504 pub fn write_frame_binding(&mut self, frame: &mut ConFrame) -> Result<(), WriteError> {
509 self.write_frame_bound(frame)
510 }
511
512 pub fn write_frame_bound(&mut self, frame: &ConFrame) -> Result<(), WriteError> {
518 let spec = frame.header.spec_version.max(crate::CON_SPEC_VERSION);
519 let flags = [
520 frame.has_velocities(),
521 frame.has_forces(),
522 frame.has_energies(),
523 frame.has_charges(),
524 frame.has_spins(),
525 frame.has_magmoms(),
526 frame.has_displacements(),
527 frame.has_spreads(),
528 ];
529 let line = frame.header.prebox_header.metadata_line();
530 let reuse = self.exact_scratch
531 && self.bound.is_some()
532 && self.bound.as_ref().is_some_and(|item| item.spec == spec)
533 && self.bound_flags == Some(flags)
534 && self.bound_line.as_deref() == Some(line);
535 if !reuse {
536 self.bound = Some(BoundHeader {
537 meta: binding_metadata(frame),
538 spec,
539 });
540 if self.exact_scratch {
541 self.bound_flags = Some(flags);
542 self.bound_vel = Some(flags[0]);
543 self.bound_line = Some(line.to_string());
544 self.bound_no_lattice = None;
545 }
546 }
547 self.keep_bound = self.exact_scratch;
548 self.write_frame(frame)
549 }
550
551 pub fn write_frame(&mut self, frame: &ConFrame) -> Result<(), WriteError> {
553 let zero_frame = frame.header.natm_types == 0
554 && frame.atom_data.is_empty()
555 && frame.header.natms_per_type.is_empty()
556 && frame.header.masses_per_type.is_empty();
557 if zero_frame && !frame.header.sections.is_empty() {
558 return Err(WriteError::Format(crate::error::ParseError::InvalidCount(
559 "a zero-atom frame cannot declare sections".into(),
560 )));
561 }
562 if !zero_frame
563 && (frame.header.natm_types == 0
564 || frame.atom_data.is_empty()
565 || frame.header.natms_per_type.contains(&0))
566 {
567 return Err(WriteError::Format(crate::error::ParseError::InvalidCount(
568 "each per-type count must be >= 1".into(),
569 )));
570 }
571 #[cfg(feature = "parallel")]
572 let parallel_rows = self.atom_workers >= 2 && frame.atom_data.len() >= WIDE_ATOM_GRAIN;
573 #[cfg(not(feature = "parallel"))]
574 let parallel_rows = false;
575 if !self.exact_scratch {
578 if parallel_rows {
579 if !ids_are_dense(&frame.atom_data, self.atom_workers) {
582 ensure_unique_ids(&frame.atom_data, self.atom_workers)?;
583 }
584 } else {
585 ensure_unique_ids(&frame.atom_data, 1)?;
586 }
587 }
588 let prec = self.float_format;
589 let stream_file = self.stream_file.take();
590 let keep_bound = self.keep_bound;
591 let bound = self.bound.take();
592 let spec_version = bound
593 .as_ref()
594 .map(|item| item.spec)
595 .unwrap_or(frame.header.spec_version);
596 let skip_refresh = keep_bound && self.metadata_cache.is_some();
599 if !skip_refresh {
600 let owned_meta;
601 let metadata = if let Some(item) = bound.as_ref() {
602 &item.meta
603 } else {
604 owned_meta = frame.header.metadata.to_map();
605 &owned_meta
606 };
607 if spec_version >= 2 {
608 self.refresh_metadata_cache(
609 spec_version,
610 metadata,
611 frame.has_velocities(),
612 frame.has_forces(),
613 frame.has_energies(),
614 frame.has_charges(),
615 frame.has_spins(),
616 frame.has_magmoms(),
617 frame.has_displacements(),
618 frame.has_spreads(),
619 )?;
620 }
621 }
622 if keep_bound {
623 self.bound = bound;
624 }
625 let (boxl, angles) = if self.keep_bound && self.bound_no_lattice == Some(true) {
628 (frame.header.boxl, frame.header.angles)
629 } else {
630 match frame.header.lattice_vectors() {
631 Some(vecs) => {
632 self.bound_no_lattice = Some(false);
633 crate::types::lengths_and_angles(vecs)
634 }
635 None => {
636 self.bound_no_lattice = Some(true);
637 (frame.header.boxl, frame.header.angles)
638 }
639 }
640 };
641 let atoms = frame.atom_data.len();
642 let [
643 has_vel,
644 has_frc,
645 has_eng,
646 has_chg,
647 has_spn,
648 has_mm,
649 has_dsp,
650 has_spr,
651 ] = if self.keep_bound {
652 self.bound_flags.unwrap_or([
653 frame.has_velocities(),
654 frame.has_forces(),
655 frame.has_energies(),
656 frame.has_charges(),
657 frame.has_spins(),
658 frame.has_magmoms(),
659 frame.has_displacements(),
660 frame.has_spreads(),
661 ])
662 } else {
663 [
664 frame.has_velocities(),
665 frame.has_forces(),
666 frame.has_energies(),
667 frame.has_charges(),
668 frame.has_spins(),
669 frame.has_magmoms(),
670 frame.has_displacements(),
671 frame.has_spreads(),
672 ]
673 };
674 let mut blocks = 1usize;
675 if has_vel {
676 blocks += 1;
677 }
678 if has_frc {
679 blocks += 1;
680 }
681 if has_eng {
682 blocks += 1;
683 }
684 if has_chg {
685 blocks += 1;
686 }
687 if has_spn {
688 blocks += 1;
689 }
690 if has_mm {
691 blocks += 1;
692 }
693 if has_dsp {
694 blocks += 1;
695 }
696 if has_spr {
697 blocks += 1;
698 }
699 let mut stream_out = match stream_file {
700 Some(mut file) => {
701 let start = file.stream_position().map_err(WriteError::Io)?;
702 Some(StreamOut {
703 file,
704 start,
705 armed: true,
706 })
707 }
708 None => None,
709 };
710 if !self.exact_scratch {
713 self.scratch
714 .reserve(256 + atoms.saturating_mul(64).saturating_mul(blocks));
715 }
716 if !self.keep_scratch {
717 self.scratch.clear();
718 }
719 push_str_line(&mut self.scratch, &frame.header.prebox_header.user);
720 if spec_version < 2 {
721 push_str_line(
723 &mut self.scratch,
724 frame.header.prebox_header.metadata_line(),
725 );
726 } else {
727 let line = self
728 .metadata_cache
729 .as_ref()
730 .expect("metadata_cache populated above")
731 .serialized
732 .as_bytes();
733 self.scratch.extend_from_slice(line);
734 self.scratch.push(b'\n');
735 }
736 {
737 let buf = &mut self.scratch;
738 let ryu = &mut self.ryu;
739 let slots = &mut self.fmt_slot;
740
741 push_f64(buf, ryu, boxl[0], prec);
743 buf.push(b' ');
744 push_f64(buf, ryu, boxl[1], prec);
745 buf.push(b' ');
746 push_f64(buf, ryu, boxl[2], prec);
747 buf.push(b'\n');
748 push_f64(buf, ryu, angles[0], prec);
749 buf.push(b' ');
750 push_f64(buf, ryu, angles[1], prec);
751 buf.push(b' ');
752 push_f64(buf, ryu, angles[2], prec);
753 buf.push(b'\n');
754 push_str_line(buf, &frame.header.postbox_header[0]);
755 push_str_line(buf, &frame.header.postbox_header[1]);
756 push_u64(buf, frame.header.natm_types as u64);
757 buf.push(b'\n');
758
759 for (i, n) in frame.header.natms_per_type.iter().enumerate() {
760 if i > 0 {
761 buf.push(b' ');
762 }
763 push_u64(buf, *n as u64);
764 }
765 buf.push(b'\n');
766
767 for (i, m) in frame.header.masses_per_type.iter().enumerate() {
768 if i > 0 {
769 buf.push(b' ');
770 }
771 push_f64(buf, ryu, *m, prec);
772 }
773 buf.push(b'\n');
774
775 let save_velocity = has_vel;
780 let direct_velocity = save_velocity
784 && self.atom_workers < 2
785 && !self.force_map
786 && frame
787 .header
788 .natms_per_type
789 .iter()
790 .all(|count| *count < WIDE_ATOM_GRAIN);
791 std::thread_local! {
794 static VEL_ROWS: std::cell::RefCell<Vec<Vec<u8>>> =
795 const { std::cell::RefCell::new(Vec::new()) };
796 }
797 struct RestoreRows(Vec<Vec<u8>>);
798 impl Drop for RestoreRows {
799 fn drop(&mut self) {
800 let rows = std::mem::take(&mut self.0);
801 VEL_ROWS.with(|cell| *cell.borrow_mut() = rows);
802 }
803 }
804 let mut velocity_rows = if direct_velocity {
805 None
806 } else {
807 Some(RestoreRows(VEL_ROWS.with(|cell| cell.take())))
808 };
809 if let Some(velocity_rows) = velocity_rows.as_mut() {
810 let ntypes = frame.header.natms_per_type.len();
811 if velocity_rows.0.len() < ntypes {
812 velocity_rows.0.resize_with(ntypes, Vec::new);
813 }
814 for row in &mut velocity_rows.0 {
815 row.clear();
816 }
817 }
818 let mut atom_idx_offset = 0;
819 for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate() {
820 let symbol = &frame.atom_data[atom_idx_offset].symbol;
821 push_str_line(buf, symbol);
822 push_component_line(buf, "Coordinates of Component", type_idx);
823
824 let mapped = self.force_map
825 || (self.atom_workers >= 2 && num_atoms_in_type >= WIDE_ATOM_GRAIN);
826 append_coord_rows(
827 buf,
828 ryu,
829 slots,
830 &frame.atom_data[atom_idx_offset..atom_idx_offset + num_atoms_in_type],
831 prec,
832 spec_version,
833 self.atom_workers,
834 mapped,
835 self.force_map,
836 stream_out.as_mut().map(|out| &mut out.file),
837 if save_velocity && !direct_velocity {
838 Some(&mut velocity_rows.as_mut().unwrap().0[type_idx])
839 } else {
840 None
841 },
842 )?;
843 atom_idx_offset += num_atoms_in_type;
844 }
845
846 if has_vel {
848 buf.push(b'\n');
849
850 let mut vel_idx_offset = 0;
851 for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
852 {
853 let symbol = &frame.atom_data[vel_idx_offset].symbol;
854 push_str_line(buf, symbol);
855 push_component_line(buf, "Velocities of Component", type_idx);
856 let atoms =
857 &frame.atom_data[vel_idx_offset..vel_idx_offset + num_atoms_in_type];
858 if direct_velocity {
859 append_velocity_lines(buf, ryu, atoms, prec, spec_version);
860 } else {
861 buf.extend_from_slice(&velocity_rows.as_ref().unwrap().0[type_idx]);
862 }
863 vel_idx_offset += num_atoms_in_type;
864 }
865 }
866
867 if has_frc {
869 buf.push(b'\n');
870
871 let mut force_idx_offset = 0;
872 for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
873 {
874 let symbol = &frame.atom_data[force_idx_offset].symbol;
875 push_str_line(buf, symbol);
876 push_component_line(buf, "Forces of Component", type_idx);
877
878 for i in 0..num_atoms_in_type {
879 let atom = &frame.atom_data[force_idx_offset + i];
880 let [fx, fy, fz] = atom.force().unwrap_or([0.0; 3]);
881 push_xyz_line(
882 buf,
883 ryu,
884 slots,
885 [fx, fy, fz],
886 prec,
887 written_fixed_flag(atom.fixed(), spec_version),
888 atom.atom_id,
889 );
890 }
891 force_idx_offset += num_atoms_in_type;
892 }
893 }
894
895 if has_eng {
897 buf.push(b'\n');
898
899 let mut energy_idx_offset = 0;
900 for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
901 {
902 let symbol = &frame.atom_data[energy_idx_offset].symbol;
903 push_str_line(buf, symbol);
904 push_component_line(buf, "Energies of Component", type_idx);
905
906 for i in 0..num_atoms_in_type {
907 let atom = &frame.atom_data[energy_idx_offset + i];
908 let e = atom.energy().unwrap_or(0.0);
909 push_scalar_line(
910 buf,
911 ryu,
912 slots,
913 e,
914 prec,
915 written_fixed_flag(atom.fixed(), spec_version),
916 atom.atom_id,
917 );
918 }
919 energy_idx_offset += num_atoms_in_type;
920 }
921 }
922
923 if has_chg {
924 buf.push(b'\n');
925 let mut off = 0;
926 for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
927 {
928 let symbol = &frame.atom_data[off].symbol;
929 push_str_line(buf, symbol);
930 push_component_line(buf, "Charges of Component", type_idx);
931 for i in 0..num_atoms_in_type {
932 let atom = &frame.atom_data[off + i];
933 let q = atom.charge().unwrap_or(0.0);
934 push_scalar_line(
935 buf,
936 ryu,
937 slots,
938 q,
939 prec,
940 written_fixed_flag(atom.fixed(), spec_version),
941 atom.atom_id,
942 );
943 }
944 off += num_atoms_in_type;
945 }
946 }
947
948 if has_spn {
949 buf.push(b'\n');
950 let mut off = 0;
951 for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
952 {
953 let symbol = &frame.atom_data[off].symbol;
954 push_str_line(buf, symbol);
955 push_component_line(buf, "Spins of Component", type_idx);
956 for i in 0..num_atoms_in_type {
957 let atom = &frame.atom_data[off + i];
958 let s = atom.spin().unwrap_or(0.0);
959 push_scalar_line(
960 buf,
961 ryu,
962 slots,
963 s,
964 prec,
965 written_fixed_flag(atom.fixed(), spec_version),
966 atom.atom_id,
967 );
968 }
969 off += num_atoms_in_type;
970 }
971 }
972
973 if has_mm {
974 buf.push(b'\n');
975 let mut off = 0;
976 for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
977 {
978 let symbol = &frame.atom_data[off].symbol;
979 push_str_line(buf, symbol);
980 push_component_line(buf, "Magmoms of Component", type_idx);
981 for i in 0..num_atoms_in_type {
982 let atom = &frame.atom_data[off + i];
983 let [mx, my, mz] = atom.magmom().unwrap_or([0.0; 3]);
984 push_xyz_line(
985 buf,
986 ryu,
987 slots,
988 [mx, my, mz],
989 prec,
990 written_fixed_flag(atom.fixed(), spec_version),
991 atom.atom_id,
992 );
993 }
994 off += num_atoms_in_type;
995 }
996 }
997
998 if has_dsp {
999 buf.push(b'\n');
1000 let mut off = 0;
1001 for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
1002 {
1003 let symbol = &frame.atom_data[off].symbol;
1004 push_str_line(buf, symbol);
1005 push_component_line(buf, "Displacements of Component", type_idx);
1006 for i in 0..num_atoms_in_type {
1007 let atom = &frame.atom_data[off + i];
1008 let [dx, dy, dz] = atom.displacement().unwrap_or([0.0; 3]);
1009 push_xyz_line(
1010 buf,
1011 ryu,
1012 slots,
1013 [dx, dy, dz],
1014 prec,
1015 written_fixed_flag(atom.fixed(), spec_version),
1016 atom.atom_id,
1017 );
1018 }
1019 off += num_atoms_in_type;
1020 }
1021 }
1022
1023 if has_spr {
1024 buf.push(b'\n');
1025 let mut off = 0;
1026 for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
1027 {
1028 let symbol = &frame.atom_data[off].symbol;
1029 push_str_line(buf, symbol);
1030 push_component_line(buf, "Spreads of Component", type_idx);
1031 for i in 0..num_atoms_in_type {
1032 let atom = &frame.atom_data[off + i];
1033 let [sx, sy, sz] = atom.spread().unwrap_or([0.0; 3]);
1034 if !crate::types::spread_row_ok([sx, sy, sz]) {
1035 let err = if sx.is_finite() && sy.is_finite() && sz.is_finite() {
1036 crate::types::invalid_spread(off + i)
1037 } else {
1038 crate::error::ParseError::NonFiniteValue(format!(
1039 "spreads atom {}",
1040 off + i
1041 ))
1042 };
1043 return Err(WriteError::Format(err));
1044 }
1045 push_xyz_line(
1046 buf,
1047 ryu,
1048 slots,
1049 [sx, sy, sz],
1050 prec,
1051 written_fixed_flag(atom.fixed(), spec_version),
1052 atom.atom_id,
1053 );
1054 }
1055 off += num_atoms_in_type;
1056 }
1057 }
1058 }
1059
1060 if self.keep_scratch {
1061 if let Some(out) = stream_out.as_mut() {
1062 out.armed = false;
1063 }
1064 return Ok(());
1065 }
1066 if let Some(out) = stream_out.as_mut() {
1067 if !self.scratch.is_empty() {
1068 out.file.write_all(&self.scratch)?;
1069 self.scratch.clear();
1070 }
1071 out.file.flush()?;
1072 out.armed = false;
1073 } else if self.scratch.len() >= 64 * 1024 {
1074 self.writer.flush()?;
1076 self.writer.get_mut().write_all(&self.scratch)?;
1077 } else {
1078 self.writer.write_all(&self.scratch)?;
1079 }
1080 Ok(())
1081 }
1082
1083 pub fn flush(&mut self) -> io::Result<()> {
1089 self.writer.flush()
1090 }
1091
1092 pub fn extend<'a>(
1096 &mut self,
1097 frames: impl Iterator<Item = &'a ConFrame>,
1098 ) -> Result<(), WriteError> {
1099 for frame in frames {
1100 self.write_frame(frame)?;
1101 }
1102 Ok(())
1103 }
1104}
1105
1106fn ids_are_dense(atoms: &[crate::types::AtomDatum], workers: usize) -> bool {
1107 let n = atoms.len();
1108 if workers < 2 || n < WIDE_ATOM_GRAIN {
1109 return atoms
1110 .iter()
1111 .enumerate()
1112 .all(|(index, atom)| atom.atom_id == index as u64);
1113 }
1114 #[cfg(feature = "parallel")]
1115 {
1116 let chunk = n.div_ceil(workers);
1117 return crate::parser::on_pool(workers, || {
1118 use rayon::prelude::*;
1119 atoms.par_chunks(chunk).enumerate().all(|(index, group)| {
1120 let base = index * chunk;
1121 group
1122 .iter()
1123 .enumerate()
1124 .all(|(offset, atom)| atom.atom_id == (base + offset) as u64)
1125 })
1126 });
1127 }
1128 #[cfg(not(feature = "parallel"))]
1129 {
1130 let _ = workers;
1131 atoms
1132 .iter()
1133 .enumerate()
1134 .all(|(index, atom)| atom.atom_id == index as u64)
1135 }
1136}
1137
1138fn ensure_unique_ids(atoms: &[crate::types::AtomDatum], workers: usize) -> Result<(), WriteError> {
1139 let n = atoms.len();
1140 if n == 0 {
1141 return Ok(());
1142 }
1143 if ids_are_dense(atoms, workers) {
1145 return Ok(());
1146 }
1147 let mut seen_ids = std::collections::HashSet::with_capacity(n);
1148 for atom in atoms {
1149 if !seen_ids.insert(atom.atom_id) {
1150 return Err(WriteError::Format(
1151 crate::error::ParseError::ValidationError(format!(
1152 "atom id {} is repeated in the frame",
1153 atom.atom_id
1154 )),
1155 ));
1156 }
1157 }
1158 Ok(())
1159}
1160
1161fn written_fixed_flag(fixed: [bool; 3], spec_version: u32) -> u8 {
1162 if spec_version < 2 {
1163 if fixed.iter().any(|axis| *axis) { 1 } else { 0 }
1164 } else {
1165 encode_fixed_bitmask(fixed)
1166 }
1167}
1168
1169fn push_u64(buf: &mut Vec<u8>, n: u64) {
1170 const DIGITS: &[u8; 200] = b"00010203040506070809101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899";
1173 let mut tmp = [0u8; 20];
1174 let mut i = 20;
1175 let mut n = n;
1176 while n >= 100 {
1177 let q = n / 100;
1178 let r = (n % 100) as usize;
1179 i -= 2;
1180 tmp[i] = DIGITS[r * 2];
1181 tmp[i + 1] = DIGITS[r * 2 + 1];
1182 n = q;
1183 }
1184 if n < 10 {
1185 i -= 1;
1186 tmp[i] = b'0' + n as u8;
1187 } else {
1188 let r = n as usize;
1189 i -= 2;
1190 tmp[i] = DIGITS[r * 2];
1191 tmp[i + 1] = DIGITS[r * 2 + 1];
1192 }
1193 buf.extend_from_slice(&tmp[i..]);
1194}
1195
1196fn push_u128(buf: &mut Vec<u8>, mut n: u128) {
1197 if n <= u128::from(u64::MAX) {
1198 push_u64(buf, n as u64);
1199 return;
1200 }
1201 let mut tmp = [0u8; 40];
1202 let mut i = 40;
1203 while n > 0 {
1204 i -= 1;
1205 tmp[i] = b'0' + (n % 10) as u8;
1206 n /= 10;
1207 }
1208 buf.extend_from_slice(&tmp[i..]);
1209}
1210
1211fn push_str_line(buf: &mut Vec<u8>, text: &str) {
1212 buf.extend_from_slice(text.as_bytes());
1213 buf.push(b'\n');
1214}
1215
1216fn push_component_line(buf: &mut Vec<u8>, label: &str, index: usize) {
1217 buf.extend_from_slice(label.as_bytes());
1218 buf.push(b' ');
1219 push_u64(buf, (index + 1) as u64);
1220 buf.push(b'\n');
1221}
1222
1223fn push_f64(buf: &mut Vec<u8>, ryu: &mut ryu::Buffer, value: f64, format: FloatFormat) {
1224 match format {
1225 FloatFormat::DecimalPlaces(precision) => push_f64_prec(buf, value, precision),
1226 FloatFormat::RoundTrip => push_round_trip(buf, ryu, None, value),
1227 }
1228}
1229
1230fn push_round_trip(
1234 buf: &mut Vec<u8>,
1235 ryu: &mut ryu::Buffer,
1236 slot: Option<&mut FmtSlot>,
1237 value: f64,
1238) {
1239 if !value.is_finite() {
1240 let text = ryu.format(value);
1241 buf.extend_from_slice(text.as_bytes());
1242 return;
1243 }
1244 if let Some(slot) = slot {
1245 let bits = value.to_bits();
1246 if slot.bits == bits {
1247 if slot.state == 2 {
1248 buf.extend_from_slice(&slot.text[..slot.len as usize]);
1249 return;
1250 }
1251 if slot.state == 1 {
1252 let n = unsafe { ryu::raw::format64(value, slot.text.as_mut_ptr()) };
1253 debug_assert!(n <= 24);
1254 slot.len = n as u8;
1255 slot.state = 2;
1256 buf.extend_from_slice(&slot.text[..n]);
1257 return;
1258 }
1259 }
1260 write_finite(buf, value);
1261 slot.bits = bits;
1262 slot.state = 1;
1263 return;
1264 }
1265 write_finite(buf, value);
1266}
1267
1268fn write_finite(buf: &mut Vec<u8>, value: f64) {
1269 if let Some((bytes, n)) = encode_fixed_decimal_into(value) {
1270 buf.extend_from_slice(&bytes[..n]);
1271 return;
1272 }
1273 if buf.capacity() - buf.len() < 24 {
1276 let mut tmp = [0u8; 24];
1277 let n = unsafe { ryu::raw::format64(value, tmp.as_mut_ptr()) };
1278 debug_assert!(n <= 24);
1279 buf.extend_from_slice(&tmp[..n]);
1280 return;
1281 }
1282 let n = unsafe { ryu::raw::format64(value, buf.as_mut_ptr().add(buf.len())) };
1283 debug_assert!(n <= 24);
1284 unsafe { buf.set_len(buf.len() + n) };
1285}
1286
1287fn round_trip_width(value: f64) -> usize {
1289 if !value.is_finite() {
1290 return ryu::Buffer::new().format(value).len();
1291 }
1292 let bits = value.to_bits();
1293 let sign = bits >> 63 != 0;
1294 let mag = f64::from_bits(bits & !(1u64 << 63));
1295 if let Some(n) = fixed_decimal_width(mag, sign) {
1296 return n;
1297 }
1298 if let Some((_, n)) = encode_fixed_decimal_into(value) {
1299 return n;
1300 }
1301 let mut tmp = [0u8; 24];
1302 unsafe { ryu::raw::format64(value, tmp.as_mut_ptr()) }
1303}
1304
1305fn fixed_decimal_width(mag: f64, sign: bool) -> Option<usize> {
1308 let sign_len = usize::from(sign);
1309 if mag == 0.0 {
1310 return Some(sign_len + 3);
1311 }
1312 if mag <= 1e16 && mag.fract() == 0.0 {
1313 let n = mag as u64;
1314 if n < 10u64.pow(16) && n as f64 == mag {
1315 return Some(sign_len + count_digits(n) as usize + 2);
1316 }
1317 }
1318 let scaled = mag * 1_000_000.0;
1319 if !scaled.is_finite() || scaled > 1e16 {
1320 return None;
1321 }
1322 let nearest = scaled.round();
1323 let mut n = nearest as u64;
1324 if nearest != n as f64 || (n as f64) / 1_000_000.0 != mag {
1325 return None;
1326 }
1327 let mut places = 6u32;
1328 while places > 0 && n.is_multiple_of(10) {
1329 n /= 10;
1330 places -= 1;
1331 }
1332 if places == 0 {
1333 if n >= 10u64.pow(16) {
1334 return None;
1335 }
1336 return Some(sign_len + count_digits(n) as usize + 2);
1337 }
1338 let digits = count_digits(n);
1339 let kk = digits as i32 - places as i32;
1340 if !(kk > -5 && kk <= 16) {
1341 return None;
1342 }
1343 let body = if kk > 0 {
1344 digits as usize + 1
1345 } else {
1346 2 + (-kk) as usize + digits as usize
1347 };
1348 Some(sign_len + body)
1349}
1350
1351fn encode_fixed_decimal_into(value: f64) -> Option<([u8; 32], usize)> {
1355 let mut out = [0u8; 32];
1356 let n = encode_fixed_decimal(value, &mut out)?;
1357 Some((out, n))
1358}
1359
1360fn encode_fixed_decimal(value: f64, out: &mut [u8; 32]) -> Option<usize> {
1361 let bits = value.to_bits();
1362 let sign = bits >> 63 != 0;
1363 let mag = f64::from_bits(bits & !(1u64 << 63));
1364 if mag == 0.0 {
1365 let text: &[u8] = if sign { b"-0.0" } else { b"0.0" };
1366 out[..text.len()].copy_from_slice(text);
1367 return Some(text.len());
1368 }
1369 if mag <= 1e16 && mag.fract() == 0.0 {
1370 let n = mag as u64;
1371 if n < 10u64.pow(16) && n as f64 == mag {
1372 return Some(write_fixed_digits(out, sign, n, 0));
1373 }
1374 }
1375 let mut scale = 10.0f64;
1376 for d in 1u32..=6 {
1377 let scaled = mag * scale;
1378 if !scaled.is_finite() || scaled > 1e16 {
1379 break;
1380 }
1381 let nearest = scaled.round();
1382 let n = nearest as u64;
1383 if nearest == n as f64 && (n as f64) / scale == mag {
1384 let mut n = n;
1385 let mut places = d;
1386 while places > 0 && n.is_multiple_of(10) {
1387 n /= 10;
1388 places -= 1;
1389 }
1390 if places == 0 {
1391 if n >= 10u64.pow(16) {
1392 return None;
1393 }
1394 return Some(write_fixed_digits(out, sign, n, 0));
1395 }
1396 let digits = count_digits(n);
1397 let kk = digits as i32 - places as i32;
1398 if !(kk > -5 && kk <= 16) {
1399 return None;
1400 }
1401 return Some(write_fixed_digits(out, sign, n, places));
1402 }
1403 scale *= 10.0;
1404 }
1405 None
1406}
1407
1408fn count_digits(mut n: u64) -> u32 {
1409 let mut digits = 1u32;
1410 while n >= 10 {
1411 n /= 10;
1412 digits += 1;
1413 }
1414 digits
1415}
1416
1417fn write_fixed_digits(out: &mut [u8; 32], sign: bool, mut n: u64, places: u32) -> usize {
1418 let mut tmp = [0u8; 20];
1419 let mut i = tmp.len();
1420 if n == 0 {
1421 tmp[i - 1] = b'0';
1422 i -= 1;
1423 } else {
1424 while n > 0 {
1425 i -= 1;
1426 tmp[i] = b'0' + (n % 10) as u8;
1427 n /= 10;
1428 }
1429 }
1430 let digs = &tmp[i..];
1431 let mut at = 0usize;
1432 if sign {
1433 out[at] = b'-';
1434 at += 1;
1435 }
1436 if places == 0 {
1437 out[at..at + digs.len()].copy_from_slice(digs);
1438 at += digs.len();
1439 out[at] = b'.';
1440 out[at + 1] = b'0';
1441 return at + 2;
1442 }
1443 let kk = digs.len() as i32 - places as i32;
1444 if kk > 0 {
1445 let cut = kk as usize;
1446 out[at..at + cut].copy_from_slice(&digs[..cut]);
1447 at += cut;
1448 out[at] = b'.';
1449 at += 1;
1450 out[at..at + digs.len() - cut].copy_from_slice(&digs[cut..]);
1451 at += digs.len() - cut;
1452 return at;
1453 }
1454 out[at] = b'0';
1455 out[at + 1] = b'.';
1456 at += 2;
1457 for _ in 0..(-kk) {
1458 out[at] = b'0';
1459 at += 1;
1460 }
1461 out[at..at + digs.len()].copy_from_slice(digs);
1462 at + digs.len()
1463}
1464
1465fn push_f64_prec(buf: &mut Vec<u8>, v: f64, prec: usize) {
1468 if !v.is_finite() || prec != 6 {
1471 let _ = write!(buf, "{v:.prec$}");
1472 return;
1473 }
1474 if v.is_sign_negative() {
1475 buf.push(b'-');
1476 }
1477 let ax = v.abs();
1478 if prec == 0 {
1479 push_u128(buf, ax.round() as u128);
1480 return;
1481 }
1482 let scale = 10u128.pow(prec as u32);
1483 let n = (ax * scale as f64).round() as u128;
1484 let int_part = n / scale;
1485 let frac = n % scale;
1486 push_u128(buf, int_part);
1487 buf.push(b'.');
1488 let mut tmp = [b'0'; 20];
1489 let mut x = frac;
1490 let mut i = prec;
1491 while i > 0 {
1492 i -= 1;
1493 tmp[i] = b'0' + (x % 10) as u8;
1494 x /= 10;
1495 }
1496 buf.extend_from_slice(&tmp[..prec]);
1497}
1498
1499const WIDE_ATOM_GRAIN: usize = 4096;
1500
1501fn take_format_buf(cap: usize) -> Vec<u8> {
1505 let mut pool = FORMAT_BUFS.lock().unwrap_or_else(|err| err.into_inner());
1506 if let Some(mut buf) = pool.pop() {
1507 if buf.capacity() < cap {
1508 buf.reserve(cap);
1509 }
1510 buf.clear();
1511 return buf;
1512 }
1513 drop(pool);
1514 Vec::with_capacity(cap)
1515}
1516
1517fn give_format_buf(buf: Vec<u8>) {
1518 let mut pool = FORMAT_BUFS.lock().unwrap_or_else(|err| err.into_inner());
1519 if pool.len() < 128 {
1520 pool.push(buf);
1521 }
1522}
1523
1524static FORMAT_BUFS: std::sync::Mutex<Vec<Vec<u8>>> = std::sync::Mutex::new(Vec::new());
1525
1526fn coord_slab(atoms: &[crate::types::AtomDatum]) -> Option<*const f64> {
1531 let n = atoms.len();
1532 if n < 2 {
1533 return None;
1534 }
1535 let base = atoms[0].coord_ptr();
1536 if base.is_null() {
1537 return None;
1538 }
1539 let aligned = |index: usize| unsafe { atoms[index].coord_ptr() == base.add(index * 3) };
1542 if aligned(1) && aligned(n / 2) && aligned(n - 1) {
1543 Some(base)
1544 } else {
1545 None
1546 }
1547}
1548
1549fn load_xyz(atom: &crate::types::AtomDatum, slab: Option<*const f64>, index: usize) -> [f64; 3] {
1550 if let Some(base) = slab {
1551 unsafe {
1552 let ptr = base.add(index * 3);
1553 [*ptr, *ptr.add(1), *ptr.add(2)]
1554 }
1555 } else {
1556 [atom.x(), atom.y(), atom.z()]
1557 }
1558}
1559
1560fn push_atom_motion(
1561 coord: &mut Vec<u8>,
1562 vel: Option<&mut Vec<u8>>,
1563 ryu: &mut ryu::Buffer,
1564 slots: &mut [FmtSlot; 3],
1565 vel_slots: &mut [FmtSlot; 3],
1566 xyz: [f64; 3],
1567 atom: &crate::types::AtomDatum,
1568 prec: FloatFormat,
1569 spec_version: u32,
1570) {
1571 let fixed = written_fixed_flag(atom.fixed(), spec_version);
1572 push_xyz_line(coord, ryu, slots, xyz, prec, fixed, atom.atom_id);
1573 if let Some(vel) = vel {
1574 let [vx, vy, vz] = atom.velocity().unwrap_or([0.0; 3]);
1575 push_xyz_line(vel, ryu, vel_slots, [vx, vy, vz], prec, fixed, atom.atom_id);
1578 }
1579}
1580
1581pub(crate) fn binding_header_bytes(frame: &ConFrame) -> Result<Vec<u8>, WriteError> {
1585 let spec = frame.header.spec_version.max(crate::CON_SPEC_VERSION);
1586 let meta = binding_metadata(frame);
1587 let mut writer = ConFrameWriter::new(Vec::<u8>::new());
1588 writer.refresh_metadata_cache(
1589 spec,
1590 &meta,
1591 frame.has_velocities(),
1592 frame.has_forces(),
1593 frame.has_energies(),
1594 frame.has_charges(),
1595 frame.has_spins(),
1596 frame.has_magmoms(),
1597 frame.has_displacements(),
1598 frame.has_spreads(),
1599 )?;
1600 let (boxl, angles) = match frame.header.lattice_vectors() {
1601 Some(vecs) => crate::types::lengths_and_angles(vecs),
1602 None => (frame.header.boxl, frame.header.angles),
1603 };
1604 let prec = FloatFormat::RoundTrip;
1605 let mut buf = Vec::with_capacity(256);
1606 push_str_line(&mut buf, &frame.header.prebox_header.user);
1607 if spec < 2 {
1608 push_str_line(&mut buf, frame.header.prebox_header.metadata_line());
1609 } else {
1610 let line = writer
1611 .metadata_cache
1612 .as_ref()
1613 .expect("metadata")
1614 .serialized
1615 .as_bytes();
1616 buf.extend_from_slice(line);
1617 buf.push(b'\n');
1618 }
1619 let ryu = &mut writer.ryu;
1620 push_f64(&mut buf, ryu, boxl[0], prec);
1621 buf.push(b' ');
1622 push_f64(&mut buf, ryu, boxl[1], prec);
1623 buf.push(b' ');
1624 push_f64(&mut buf, ryu, boxl[2], prec);
1625 buf.push(b'\n');
1626 push_f64(&mut buf, ryu, angles[0], prec);
1627 buf.push(b' ');
1628 push_f64(&mut buf, ryu, angles[1], prec);
1629 buf.push(b' ');
1630 push_f64(&mut buf, ryu, angles[2], prec);
1631 buf.push(b'\n');
1632 push_str_line(&mut buf, &frame.header.postbox_header[0]);
1633 push_str_line(&mut buf, &frame.header.postbox_header[1]);
1634 push_u64(&mut buf, frame.header.natm_types as u64);
1635 buf.push(b'\n');
1636 for (i, n) in frame.header.natms_per_type.iter().enumerate() {
1637 if i > 0 {
1638 buf.push(b' ');
1639 }
1640 push_u64(&mut buf, *n as u64);
1641 }
1642 buf.push(b'\n');
1643 for (i, m) in frame.header.masses_per_type.iter().enumerate() {
1644 if i > 0 {
1645 buf.push(b' ');
1646 }
1647 push_f64(&mut buf, ryu, *m, prec);
1648 }
1649 buf.push(b'\n');
1650 Ok(buf)
1651}
1652
1653pub(crate) fn format_coord_range(frame: &ConFrame, begin: usize, end: usize) -> Vec<u8> {
1656 let spec = frame.header.spec_version.max(crate::CON_SPEC_VERSION);
1657 let prec = FloatFormat::RoundTrip;
1658 let mut buf = Vec::with_capacity((end - begin).saturating_mul(32) + 64);
1659 let mut ryu = ryu::Buffer::new();
1660 let mut slots = [FmtSlot::new(); 3];
1661 let mut vel_slots = [FmtSlot::new(); 3];
1662 let mut offset = 0usize;
1663 for (type_idx, &count) in frame.header.natms_per_type.iter().enumerate() {
1664 let type_end = offset + count;
1665 let from = begin.max(offset);
1666 let to = end.min(type_end);
1667 if from < to {
1668 if from == offset {
1669 let symbol = frame.atom_data[offset].symbol.as_ref();
1670 push_str_line(&mut buf, symbol);
1671 push_component_line(&mut buf, "Coordinates of Component", type_idx);
1672 }
1673 let atoms = &frame.atom_data[from..to];
1674 let slab = coord_slab(atoms);
1675 for (index, atom) in atoms.iter().enumerate() {
1676 push_atom_motion(
1677 &mut buf,
1678 None,
1679 &mut ryu,
1680 &mut slots,
1681 &mut vel_slots,
1682 load_xyz(atom, slab, index),
1683 atom,
1684 prec,
1685 spec,
1686 );
1687 }
1688 }
1689 offset = type_end;
1690 }
1691 buf
1692}
1693
1694pub fn format_coordinate_lines(atoms: &[crate::types::AtomDatum]) -> Result<Vec<u8>, WriteError> {
1699 let mut buf = take_format_buf(atoms.len().saturating_mul(48));
1700 let mut ryu = ryu::Buffer::new();
1701 let mut slots = [FmtSlot::new(); 3];
1702 append_coord_rows(
1703 &mut buf,
1704 &mut ryu,
1705 &mut slots,
1706 atoms,
1707 FloatFormat::RoundTrip,
1708 2,
1709 1,
1710 false,
1711 false,
1712 None,
1713 None,
1714 )?;
1715 Ok(buf)
1716}
1717
1718pub fn recycle_format_buffer(buf: Vec<u8>) {
1721 give_format_buf(buf);
1722}
1723
1724fn append_velocity_lines(
1725 buf: &mut Vec<u8>,
1726 ryu: &mut ryu::Buffer,
1727 atoms: &[crate::types::AtomDatum],
1728 prec: FloatFormat,
1729 spec_version: u32,
1730) {
1731 let mut vel_slots = [FmtSlot::new(); 3];
1732 for atom in atoms {
1733 let [vx, vy, vz] = atom.velocity().unwrap_or([0.0; 3]);
1734 push_xyz_line(
1735 buf,
1736 ryu,
1737 &mut vel_slots,
1738 [vx, vy, vz],
1739 prec,
1740 written_fixed_flag(atom.fixed(), spec_version),
1741 atom.atom_id,
1742 );
1743 }
1744}
1745
1746fn append_coord_rows(
1747 buf: &mut Vec<u8>,
1748 ryu: &mut ryu::Buffer,
1749 slots: &mut [FmtSlot; 3],
1750 atoms: &[crate::types::AtomDatum],
1751 prec: FloatFormat,
1752 spec_version: u32,
1753 workers: usize,
1754 mapped: bool,
1755 force_map: bool,
1756 stream: Option<&mut File>,
1757 mut vel: Option<&mut Vec<u8>>,
1758) -> Result<(), WriteError> {
1759 let want_vel = vel.is_some();
1760 if let Some(file) = stream {
1761 if mapped && !atoms.is_empty() {
1762 file.write_all(buf).map_err(WriteError::Io)?;
1763 buf.clear();
1764 let vel_bytes = stream_coord_rows(
1765 file,
1766 atoms,
1767 workers,
1768 prec,
1769 spec_version,
1770 want_vel,
1771 force_map,
1772 )?;
1773 if let Some(dest) = vel.as_mut() {
1774 dest.extend_from_slice(&vel_bytes);
1775 }
1776 return Ok(());
1777 }
1778 }
1779 if workers < 2 || atoms.len() < WIDE_ATOM_GRAIN {
1780 let mut vel_slots = [FmtSlot::new(); 3];
1781 let slab = coord_slab(atoms);
1782 for (index, atom) in atoms.iter().enumerate() {
1783 push_atom_motion(
1784 buf,
1785 vel.as_deref_mut(),
1786 ryu,
1787 slots,
1788 &mut vel_slots,
1789 load_xyz(atom, slab, index),
1790 atom,
1791 prec,
1792 spec_version,
1793 );
1794 }
1795 return Ok(());
1796 }
1797 #[cfg(feature = "parallel")]
1798 {
1799 let groups = workers.saturating_mul(4).max(workers);
1804 let chunk = atoms.len().div_ceil(groups).max(1);
1805 let nparts = atoms.len().div_ceil(chunk);
1806 let bad_ids = std::sync::atomic::AtomicBool::new(false);
1807 let next = std::sync::atomic::AtomicUsize::new(0);
1808 let parts: Vec<std::sync::Mutex<Option<(Vec<u8>, Vec<u8>)>>> =
1809 (0..nparts).map(|_| std::sync::Mutex::new(None)).collect();
1810 crate::parser::on_pool(workers, || {
1811 use rayon::prelude::*;
1812 (0..workers).into_par_iter().for_each(|_| {
1813 let mut ryu = ryu::Buffer::new();
1814 let mut slots = [FmtSlot::new(); 3];
1815 let mut vel_slots = [FmtSlot::new(); 3];
1816 loop {
1817 let index = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1818 if index >= nparts {
1819 break;
1820 }
1821 let start = index * chunk;
1822 let end = (start + chunk).min(atoms.len());
1823 let group = &atoms[start..end];
1824 let mut local = Vec::with_capacity(group.len().saturating_mul(48));
1825 let mut local_vel = if want_vel {
1826 Vec::with_capacity(group.len().saturating_mul(48))
1827 } else {
1828 Vec::new()
1829 };
1830 let slab = coord_slab(group);
1831 for (offset, atom) in group.iter().enumerate() {
1832 if atom.atom_id != (start + offset) as u64 {
1833 bad_ids.store(true, std::sync::atomic::Ordering::Relaxed);
1834 }
1835 push_atom_motion(
1836 &mut local,
1837 if want_vel { Some(&mut local_vel) } else { None },
1838 &mut ryu,
1839 &mut slots,
1840 &mut vel_slots,
1841 load_xyz(atom, slab, offset),
1842 atom,
1843 prec,
1844 spec_version,
1845 );
1846 }
1847 *parts[index].lock().expect("format part") = Some((local, local_vel));
1848 }
1849 });
1850 });
1851 if bad_ids.load(std::sync::atomic::Ordering::Relaxed) {
1852 ensure_unique_ids(atoms, 1)?;
1853 }
1854 let mut coord_len = 0usize;
1855 let mut vel_len = 0usize;
1856 for part in &parts {
1857 let guard = part.lock().expect("format part");
1858 if let Some((coord, rows)) = guard.as_ref() {
1859 coord_len += coord.len();
1860 vel_len += rows.len();
1861 }
1862 }
1863 let coord_at = buf.len();
1864 buf.reserve(coord_len);
1865 if let Some(dest) = vel.as_deref_mut() {
1866 dest.reserve(vel_len);
1867 }
1868 if coord_len > 0 {
1869 let coord_ptr = buf.as_mut_ptr() as usize;
1870 let mut jobs = Vec::with_capacity(nparts);
1871 let mut at = coord_at;
1872 for part in &parts {
1873 let guard = part.lock().expect("format part");
1874 let Some((coord, _)) = guard.as_ref() else {
1875 continue;
1876 };
1877 let n = coord.len();
1878 if n == 0 {
1879 continue;
1880 }
1881 jobs.push((coord.as_ptr() as usize, coord_ptr + at, n));
1882 at += n;
1883 }
1884 let next_copy = std::sync::atomic::AtomicUsize::new(0);
1885 crate::parser::on_pool(workers, || {
1886 use rayon::prelude::*;
1887 (0..workers).into_par_iter().for_each(|_| {
1888 loop {
1889 let index = next_copy.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1890 if index >= jobs.len() {
1891 break;
1892 }
1893 let (src, dest, n) = jobs[index];
1894 unsafe {
1895 std::ptr::copy_nonoverlapping(src as *const u8, dest as *mut u8, n);
1896 }
1897 }
1898 });
1899 });
1900 unsafe { buf.set_len(coord_at + coord_len) };
1904 }
1905 if let Some(dest) = vel.as_deref_mut() {
1906 for part in &parts {
1907 let guard = part.lock().expect("format part");
1908 if let Some((_, rows)) = guard.as_ref() {
1909 dest.extend_from_slice(rows);
1910 }
1911 }
1912 }
1913 return Ok(());
1914 }
1915 #[cfg(not(feature = "parallel"))]
1916 {
1917 let _ = workers;
1918 let mut vel_slots = [FmtSlot::new(); 3];
1919 let slab = coord_slab(atoms);
1920 for (index, atom) in atoms.iter().enumerate() {
1921 push_atom_motion(
1922 buf,
1923 vel.as_deref_mut(),
1924 ryu,
1925 slots,
1926 &mut vel_slots,
1927 load_xyz(atom, slab, index),
1928 atom,
1929 prec,
1930 spec_version,
1931 );
1932 }
1933 return Ok(());
1934 }
1935}
1936
1937fn physical_cores() -> usize {
1946 let logical = std::thread::available_parallelism()
1947 .map(|n| n.get())
1948 .unwrap_or(1);
1949 let siblings =
1950 std::fs::read_to_string("/sys/devices/system/cpu/cpu0/topology/thread_siblings_list")
1951 .ok()
1952 .map(|text| {
1953 text.split([',', '-'])
1954 .filter(|piece| !piece.trim().is_empty())
1955 .count()
1956 })
1957 .filter(|n| *n > 0)
1958 .unwrap_or(1);
1959 (logical / siblings).max(1)
1960}
1961
1962fn side_copy_workers(workers: usize) -> bool {
1963 #[cfg(miri)]
1965 {
1966 let _ = workers;
1967 return false;
1968 }
1969 #[cfg(not(miri))]
1970 {
1971 workers >= 2 && workers.saturating_mul(2) <= physical_cores()
1972 }
1973}
1974
1975pub(crate) fn advise_shared_store(file: &File, ptr: *mut u8, len: usize) {
1984 #[cfg(all(target_os = "linux", not(miri)))]
1985 if len >= 2 * 1024 * 1024 && mount_requests_huge(file) {
1986 unsafe {
1987 let _ = libc::madvise(ptr.cast::<libc::c_void>(), len, libc::MADV_HUGEPAGE);
1988 }
1989 }
1990 #[cfg(any(not(target_os = "linux"), miri))]
1991 {
1992 let _ = (file, ptr, len);
1993 }
1994}
1995
1996#[cfg(all(target_os = "linux", not(miri)))]
1997fn mount_requests_huge(file: &File) -> bool {
1998 use std::os::fd::AsRawFd;
1999 let Ok(linked) = std::fs::read_link(format!("/proc/self/fd/{}", file.as_raw_fd())) else {
2000 return false;
2001 };
2002 let Ok(table) = std::fs::read_to_string("/proc/mounts") else {
2003 return false;
2004 };
2005 let path = linked.to_string_lossy();
2006 let mut best_len = 0usize;
2007 let mut huge = false;
2008 for line in table.lines() {
2009 let mut fields = line.split_whitespace();
2010 let _device = fields.next();
2011 let Some(raw_mount) = fields.next() else {
2012 continue;
2013 };
2014 let mount = raw_mount.replace("\\040", " ");
2015 if mount.len() >= best_len && path.starts_with(&mount) {
2016 let _fstype = fields.next();
2017 let options = fields.next().unwrap_or("");
2018 best_len = mount.len();
2019 huge = options.split(',').any(|opt| {
2020 matches!(
2021 opt,
2022 "huge=within_size" | "huge=always" | "huge=advise" | "huge=force"
2023 )
2024 });
2025 }
2026 }
2027 huge
2028}
2029
2030pub(crate) fn file_is_tmpfs(file: &File) -> bool {
2031 #[cfg(all(not(miri), target_os = "linux"))]
2034 {
2035 use std::os::fd::AsRawFd;
2036 let mut stat = unsafe { std::mem::zeroed::<libc::statfs>() };
2037 let rc = unsafe { libc::fstatfs(file.as_raw_fd(), &mut stat) };
2038 return rc == 0 && stat.f_type == 0x01021994;
2039 }
2040 #[cfg(any(miri, not(target_os = "linux")))]
2041 {
2042 let _ = file;
2043 false
2044 }
2045}
2046
2047fn stream_coord_pipeline(
2054 file: &mut File,
2055 atoms: &[crate::types::AtomDatum],
2056 workers: usize,
2057 prec: FloatFormat,
2058 spec_version: u32,
2059 want_vel: bool,
2060) -> Result<Vec<u8>, WriteError> {
2061 let workers = workers.max(2).min(atoms.len().max(1));
2062 let formatters = workers - 1;
2063 let groups = if file_is_tmpfs(file) {
2068 workers.saturating_mul(32)
2069 } else {
2070 workers.saturating_mul(8)
2071 }
2072 .max(workers);
2073 let chunk = atoms.len().div_ceil(groups).max(1);
2074 let nparts = atoms.len().div_ceil(chunk);
2075 let base = file.stream_position().map_err(WriteError::Io)?;
2076 let out = file.try_clone().map_err(WriteError::Io)?;
2077 let next = std::sync::atomic::AtomicUsize::new(0);
2078 let bad_ids = std::sync::atomic::AtomicBool::new(false);
2079 let slots: Vec<std::sync::Mutex<Option<(Vec<u8>, Vec<u8>)>>> =
2080 (0..nparts).map(|_| std::sync::Mutex::new(None)).collect();
2081 let flag = &bad_ids;
2082 let (end, mut vel_bytes) = std::thread::scope(|scope| -> Result<_, WriteError> {
2083 let mut handles = Vec::with_capacity(formatters);
2084 for worker in 0..formatters {
2085 let next = &next;
2086 let slots = &slots;
2087 handles.push(scope.spawn(move || {
2088 crate::parser::pin_worker(worker + 1);
2090 let mut ryu = ryu::Buffer::new();
2091 let mut slots_fmt = [FmtSlot::new(); 3];
2092 let mut vel_slots = [FmtSlot::new(); 3];
2093 loop {
2094 let index = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2095 if index >= nparts {
2096 break;
2097 }
2098 let start = index * chunk;
2099 let end = (start + chunk).min(atoms.len());
2100 let group = &atoms[start..end];
2101 let mut coord = Vec::with_capacity(group.len().saturating_mul(48));
2102 let mut vel = if want_vel {
2103 Vec::with_capacity(group.len().saturating_mul(48))
2104 } else {
2105 Vec::new()
2106 };
2107 let slab = coord_slab(group);
2108 for (offset, atom) in group.iter().enumerate() {
2109 if atom.atom_id != (start + offset) as u64 {
2110 flag.store(true, std::sync::atomic::Ordering::Relaxed);
2111 }
2112 push_atom_motion(
2113 &mut coord,
2114 if want_vel { Some(&mut vel) } else { None },
2115 &mut ryu,
2116 &mut slots_fmt,
2117 &mut vel_slots,
2118 load_xyz(atom, slab, offset),
2119 atom,
2120 prec,
2121 spec_version,
2122 );
2123 }
2124 *slots[index].lock().expect("format slot") = Some((coord, vel));
2125 }
2126 }));
2127 }
2128 let mut offset = base;
2129 let mut vel_bytes = Vec::new();
2130 #[cfg(all(not(miri), target_os = "linux"))]
2134 let mut old_affinity = unsafe { std::mem::zeroed::<libc::cpu_set_t>() };
2135 #[cfg(all(not(miri), target_os = "linux"))]
2136 unsafe {
2137 libc::sched_getaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &mut old_affinity);
2138 }
2139 crate::parser::pin_worker(0);
2140 let written = (|| -> Result<_, WriteError> {
2141 for index in 0..nparts {
2142 let (coord, vel) = loop {
2143 if let Some(part) = slots[index].lock().expect("format slot").take() {
2144 break part;
2145 }
2146 std::hint::spin_loop();
2147 };
2148 pwrite_full(&out, &coord, offset).map_err(WriteError::Io)?;
2149 offset = offset.checked_add(coord.len() as u64).ok_or_else(|| {
2150 WriteError::Io(io::Error::new(
2151 io::ErrorKind::InvalidData,
2152 "coordinate block is too large",
2153 ))
2154 })?;
2155 if want_vel {
2156 vel_bytes.extend_from_slice(&vel);
2157 }
2158 }
2159 Ok(())
2160 })();
2161 #[cfg(all(not(miri), target_os = "linux"))]
2162 unsafe {
2163 libc::sched_setaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &old_affinity);
2164 }
2165 written?;
2166 for handle in handles {
2167 if handle.join().is_err() {
2168 return Err(WriteError::Io(io::Error::other("formatter panicked")));
2169 }
2170 }
2171 Ok((offset, vel_bytes))
2172 })?;
2173 file.set_len(end).map_err(WriteError::Io)?;
2176 file.seek(SeekFrom::Start(end)).map_err(WriteError::Io)?;
2177 if bad_ids.load(std::sync::atomic::Ordering::Relaxed) {
2178 ensure_unique_ids(atoms, 1)?;
2179 }
2180 if !want_vel {
2181 vel_bytes.clear();
2182 }
2183 Ok(vel_bytes)
2184}
2185
2186fn mapped_line_budget(prec: FloatFormat) -> usize {
2192 match prec {
2193 FloatFormat::RoundTrip => 128,
2194 FloatFormat::DecimalPlaces(places) => {
2195 let float = 1usize
2196 .saturating_add(309)
2197 .saturating_add(1)
2198 .saturating_add(places);
2199 float.saturating_mul(3).saturating_add(4 + 1 + 1 + 20 + 1)
2200 }
2201 }
2202}
2203
2204fn stream_coord_rows(
2205 file: &mut File,
2206 atoms: &[crate::types::AtomDatum],
2207 workers: usize,
2208 prec: FloatFormat,
2209 spec_version: u32,
2210 want_vel: bool,
2211 force_map: bool,
2212) -> Result<Vec<u8>, WriteError> {
2213 let workers = workers.max(1).min(atoms.len().max(1));
2217 let side_copies = side_copy_workers(workers);
2218 if !side_copies && !force_map {
2219 return stream_coord_pipeline(file, atoms, workers, prec, spec_version, want_vel);
2220 }
2221 let groups = if side_copies {
2222 workers.saturating_mul(8).max(workers)
2223 } else {
2224 workers.saturating_mul(4).max(workers)
2225 };
2226 let chunk = atoms.len().div_ceil(groups).max(1);
2227 let nparts = atoms.len().div_ceil(chunk);
2228 let base = file.stream_position().map_err(WriteError::Io)?;
2229 let line_budget = mapped_line_budget(prec);
2230 let upper = atoms.len().checked_mul(line_budget).ok_or_else(|| {
2231 WriteError::Io(io::Error::new(
2232 io::ErrorKind::InvalidData,
2233 "coordinate block is too large",
2234 ))
2235 })?;
2236 let reserve_end = base.checked_add(upper as u64).ok_or_else(|| {
2237 WriteError::Io(io::Error::new(
2238 io::ErrorKind::InvalidData,
2239 "coordinate block is too large",
2240 ))
2241 })?;
2242 let mut map = {
2243 file.set_len(reserve_end).map_err(WriteError::Io)?;
2244 Some(unsafe { memmap2::MmapMut::map_mut(&*file).map_err(WriteError::Io)? })
2245 };
2246 let dest_base = map
2247 .as_mut()
2248 .map(|mapping| mapping.as_mut_ptr() as usize + base as usize)
2249 .unwrap_or(0);
2250 let lengths: Vec<std::sync::atomic::AtomicU64> = (0..nparts)
2251 .map(|_| std::sync::atomic::AtomicU64::new(0))
2252 .collect();
2253 let bad_ids = std::sync::atomic::AtomicBool::new(false);
2254 let failed = std::sync::atomic::AtomicBool::new(false);
2255 let next_job = std::sync::atomic::AtomicUsize::new(0);
2256 let produced = std::sync::atomic::AtomicUsize::new(0);
2257 let copies: std::sync::Mutex<Vec<(Vec<u8>, usize)>> = std::sync::Mutex::new(Vec::new());
2258 let ready = std::sync::Condvar::new();
2259 let vel_parts: Vec<std::sync::Mutex<Option<Vec<u8>>>> =
2260 (0..nparts).map(|_| std::sync::Mutex::new(None)).collect();
2261 let flag = &bad_ids;
2262 std::thread::scope(|scope| -> Result<(), WriteError> {
2263 let mut copiers = Vec::with_capacity(workers);
2264 if side_copies {
2265 for copier in 0..workers {
2266 let copies = &copies;
2267 let ready = &ready;
2268 let produced = &produced;
2269 copiers.push(scope.spawn(move || {
2270 crate::parser::pin_worker(workers + copier);
2272 loop {
2273 let job = {
2274 let mut queue = copies.lock().expect("copies");
2275 loop {
2276 if let Some(job) = queue.pop() {
2277 break Some(job);
2278 }
2279 if produced.load(std::sync::atomic::Ordering::Acquire) == nparts {
2280 break None;
2281 }
2282 queue = ready.wait(queue).expect("copies");
2283 }
2284 };
2285 let Some((coord, offset)) = job else {
2286 break;
2287 };
2288 let n = coord.len();
2289 if n > 0 {
2290 unsafe {
2291 std::ptr::copy_nonoverlapping(
2292 coord.as_ptr(),
2293 (dest_base + offset) as *mut u8,
2294 n,
2295 );
2296 }
2297 }
2298 give_format_buf(coord);
2299 }
2300 }));
2301 }
2302 }
2303 let mut formatters = Vec::with_capacity(workers);
2304 for worker in 0..workers {
2305 let lengths = &lengths;
2306 let next_job = &next_job;
2307 let vel_parts = &vel_parts;
2308 let copies = &copies;
2309 let ready = &ready;
2310 let produced = &produced;
2311 let failed = &failed;
2312 formatters.push(scope.spawn(move || -> Result<(), WriteError> {
2313 crate::parser::pin_worker(worker);
2314 let mut ryu = ryu::Buffer::new();
2315 let mut slots = [FmtSlot::new(); 3];
2316 let mut vel_slots = [FmtSlot::new(); 3];
2317 loop {
2318 let index = next_job.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2319 if index >= nparts {
2320 break;
2321 }
2322 let start = index * chunk;
2323 let end = (start + chunk).min(atoms.len());
2324 let group = &atoms[start..end];
2325 let mut coord = if side_copies {
2326 take_format_buf(group.len().saturating_mul(48))
2327 } else {
2328 Vec::with_capacity(group.len().saturating_mul(48))
2329 };
2330 let mut vel = if want_vel {
2331 Vec::with_capacity(group.len().saturating_mul(48))
2332 } else {
2333 Vec::new()
2334 };
2335 let slab = coord_slab(group);
2336 for (offset, atom) in group.iter().enumerate() {
2337 if atom.atom_id != (start + offset) as u64 {
2338 flag.store(true, std::sync::atomic::Ordering::Relaxed);
2339 }
2340 push_atom_motion(
2341 &mut coord,
2342 if want_vel { Some(&mut vel) } else { None },
2343 &mut ryu,
2344 &mut slots,
2345 &mut vel_slots,
2346 load_xyz(atom, slab, offset),
2347 atom,
2348 prec,
2349 spec_version,
2350 );
2351 }
2352 let n = coord.len();
2353 if n > group.len().saturating_mul(line_budget) {
2354 failed.store(true, std::sync::atomic::Ordering::Release);
2358 lengths[index].store(1, std::sync::atomic::Ordering::Release);
2359 return Err(WriteError::Io(io::Error::other(
2360 "coordinate line exceeded the mapped reserve",
2361 )));
2362 }
2363 lengths[index].store(n as u64 + 1, std::sync::atomic::Ordering::Release);
2364 let mut offset = 0usize;
2365 for prev in 0..index {
2366 loop {
2367 if failed.load(std::sync::atomic::Ordering::Acquire) {
2368 return Err(WriteError::Io(io::Error::other(
2369 "coordinate line exceeded the mapped reserve",
2370 )));
2371 }
2372 let known = lengths[prev].load(std::sync::atomic::Ordering::Acquire);
2373 if known != 0 {
2374 if failed.load(std::sync::atomic::Ordering::Acquire) {
2375 return Err(WriteError::Io(io::Error::other(
2376 "coordinate line exceeded the mapped reserve",
2377 )));
2378 }
2379 offset += (known - 1) as usize;
2380 break;
2381 }
2382 std::hint::spin_loop();
2383 }
2384 }
2385 if side_copies {
2386 let mut queue = copies.lock().expect("copies");
2387 queue.push((coord, offset));
2388 } else {
2389 let n = coord.len();
2390 if n > 0 {
2391 unsafe {
2392 std::ptr::copy_nonoverlapping(
2393 coord.as_ptr(),
2394 (dest_base + offset) as *mut u8,
2395 n,
2396 );
2397 }
2398 }
2399 }
2400 if side_copies {
2401 produced.fetch_add(1, std::sync::atomic::Ordering::Release);
2402 ready.notify_one();
2403 }
2404 if want_vel {
2405 *vel_parts[index].lock().expect("vel part") = Some(vel);
2406 }
2407 }
2408 Ok(())
2409 }));
2410 }
2411 let formatted = (|| -> Result<(), WriteError> {
2412 for handle in formatters {
2413 handle.join().unwrap_or_else(|_| {
2414 Err(WriteError::Io(io::Error::other("formatter panicked")))
2415 })?;
2416 }
2417 Ok(())
2418 })();
2419 produced.store(nparts, std::sync::atomic::Ordering::Release);
2422 ready.notify_all();
2423 for handle in copiers {
2424 handle
2425 .join()
2426 .map_err(|_| WriteError::Io(io::Error::other("copy panicked")))?;
2427 }
2428 formatted?;
2429 Ok(())
2430 })?;
2431 let mut cursor = 0u64;
2432 for length in &lengths {
2433 cursor += length
2434 .load(std::sync::atomic::Ordering::Relaxed)
2435 .saturating_sub(1);
2436 }
2437 drop(map);
2438 let end = base + cursor;
2439 file.set_len(end).map_err(WriteError::Io)?;
2441 file.seek(SeekFrom::Start(end)).map_err(WriteError::Io)?;
2442 if bad_ids.load(std::sync::atomic::Ordering::Relaxed) {
2443 ensure_unique_ids(atoms, 1)?;
2444 }
2445 let mut vel_bytes = Vec::new();
2446 if want_vel {
2447 for part in &vel_parts {
2448 if let Some(vel) = part.lock().expect("vel part").as_ref() {
2449 vel_bytes.extend_from_slice(vel);
2450 }
2451 }
2452 }
2453 drop(vel_parts);
2454 Ok(vel_bytes)
2455}
2456
2457fn push_xyz_line(
2458 buf: &mut Vec<u8>,
2459 ryu: &mut ryu::Buffer,
2460 slots: &mut [FmtSlot; 3],
2461 xyz: [f64; 3],
2462 prec: FloatFormat,
2463 fixed: u8,
2464 atom_id: u64,
2465) {
2466 let [x, y, z] = xyz;
2467 match prec {
2468 FloatFormat::RoundTrip => {
2469 push_round_trip(buf, ryu, Some(&mut slots[0]), x);
2470 buf.push(b' ');
2471 push_round_trip(buf, ryu, Some(&mut slots[1]), y);
2472 buf.push(b' ');
2473 push_round_trip(buf, ryu, Some(&mut slots[2]), z);
2474 }
2475 FloatFormat::DecimalPlaces(_) => {
2476 push_f64(buf, ryu, x, prec);
2477 buf.push(b' ');
2478 push_f64(buf, ryu, y, prec);
2479 buf.push(b' ');
2480 push_f64(buf, ryu, z, prec);
2481 }
2482 }
2483 buf.push(b' ');
2484 buf.push(b'0' + fixed);
2486 buf.push(b' ');
2487 push_u64(buf, atom_id);
2488 buf.push(b'\n');
2489}
2490
2491fn push_scalar_line(
2492 buf: &mut Vec<u8>,
2493 ryu: &mut ryu::Buffer,
2494 slots: &mut [FmtSlot; 3],
2495 v: f64,
2496 prec: FloatFormat,
2497 fixed: u8,
2498 atom_id: u64,
2499) {
2500 match prec {
2501 FloatFormat::RoundTrip => push_round_trip(buf, ryu, Some(&mut slots[0]), v),
2502 FloatFormat::DecimalPlaces(_) => push_f64(buf, ryu, v, prec),
2503 }
2504 buf.push(b' ');
2505 push_u64(buf, u64::from(fixed));
2506 buf.push(b' ');
2507 push_u64(buf, atom_id);
2508 buf.push(b'\n');
2509}
2510
2511pub(crate) fn stamp_legacy_velocity_time(
2519 meta: &mut BTreeMap<String, Value>,
2520 spec_version: u32,
2521 has_velocity: bool,
2522) {
2523 if !has_velocity {
2524 return;
2525 }
2526 let minor = meta
2527 .get(meta::CON_SPEC_MINOR)
2528 .and_then(Value::as_u64)
2529 .unwrap_or(0);
2530 if spec_version >= 3 && minor >= 1 {
2531 return;
2532 }
2533 let mut units = match meta.get(meta::UNITS).cloned() {
2538 Some(Value::Object(map)) => Value::Object(map),
2539 Some(_) => return,
2540 None => crate::units::default_v3_units_json(),
2541 };
2542 if let Some(object) = units.as_object_mut() {
2543 object.insert(
2544 "time".to_string(),
2545 Value::String(crate::units::NATURAL_TIME_UNIT.to_string()),
2546 );
2547 meta.insert(meta::UNITS.to_string(), units);
2548 if meta.contains_key(meta::CON_SPEC_MINOR) {
2550 meta.insert(meta::CON_SPEC_MINOR.to_string(), Value::from(1u64));
2551 }
2552 }
2553}
2554
2555pub fn binding_metadata(frame: &ConFrame) -> BTreeMap<String, Value> {
2561 let mut meta = frame.header.metadata.to_map();
2562 meta.remove(meta::CON_SPEC_VERSION);
2563 meta.remove(meta::SECTIONS);
2564 let has_velocity = frame.has_velocities();
2565 stamp_legacy_velocity_time(&mut meta, frame.header.spec_version, has_velocity);
2566 if !meta.contains_key(meta::UNITS) {
2569 meta.insert(
2570 meta::UNITS.to_string(),
2571 crate::units::default_v3_units_json(),
2572 );
2573 }
2574 if !meta.contains_key(meta::CON_SPEC_MINOR) {
2575 meta.insert(meta::CON_SPEC_MINOR.to_string(), Value::from(1u64));
2576 }
2577 meta
2578}
2579
2580pub fn binding_rewrite_frame(frame: &ConFrame) -> Result<ConFrame, crate::error::ParseError> {
2582 let mut out = frame.clone();
2583 out.header.metadata = crate::types::FrameMeta::from_map(binding_metadata(frame));
2584 if out.header.spec_version < crate::CON_SPEC_VERSION {
2587 out.header.spec_version = crate::CON_SPEC_VERSION;
2588 }
2589 Ok(out)
2590}
2591
2592impl ConFrameWriter<File> {
2594 pub fn from_path<P: AsRef<Path>>(path: P) -> io::Result<Self> {
2598 let file = File::create(path)?;
2599 Ok(Self::new(file))
2600 }
2601
2602 pub fn from_path_with_precision<P: AsRef<Path>>(path: P, precision: usize) -> io::Result<Self> {
2604 let file = File::create(path)?;
2605 Ok(Self::with_precision(file, precision))
2606 }
2607
2608 pub fn from_path_append<P: AsRef<Path>>(path: P) -> io::Result<Self> {
2645 Ok(Self::new(open_uncompressed_append(path.as_ref())?))
2646 }
2647
2648 pub fn from_path_append_with_precision<P: AsRef<Path>>(
2650 path: P,
2651 precision: usize,
2652 ) -> io::Result<Self> {
2653 Ok(Self::with_precision(
2654 open_uncompressed_append(path.as_ref())?,
2655 precision,
2656 ))
2657 }
2658}
2659
2660pub(crate) fn open_uncompressed_append(path: &Path) -> io::Result<File> {
2662 if let Some(ext) = path.extension().and_then(|ext| ext.to_str())
2663 && (ext == "gz" || ext == "zst")
2664 {
2665 return Err(io::Error::new(
2666 io::ErrorKind::InvalidInput,
2667 "append is only supported for an uncompressed file",
2668 ));
2669 }
2670 let mut file = OpenOptions::new()
2671 .create(true)
2672 .read(true)
2673 .append(true)
2674 .open(path)?;
2675 let len = file.seek(SeekFrom::End(0))?;
2676 if len > 0 {
2677 file.seek(SeekFrom::End(-1))?;
2678 let mut last = [0u8; 1];
2679 file.read_exact(&mut last)?;
2680 if last[0] != b'\n' {
2681 file.write_all(b"\n")?;
2682 }
2683 }
2684 Ok(file)
2685}
2686
2687pub fn write_frames_to_vec(frames: &[ConFrame]) -> Result<Vec<u8>, WriteError> {
2691 write_frames_to_vec_with_threads(frames, 1, None, false)
2692}
2693
2694pub fn write_frames_to_vec_with_threads(
2701 frames: &[ConFrame],
2702 workers: usize,
2703 precision: Option<usize>,
2704 canonical: bool,
2705) -> Result<Vec<u8>, WriteError> {
2706 let refs: Vec<&ConFrame> = frames.iter().collect();
2707 write_frame_refs_with_threads(&refs, workers, precision, canonical)
2708}
2709
2710fn temporary_sibling(path: &Path) -> io::Result<(std::path::PathBuf, File)> {
2716 static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2717 let dir = match path.parent() {
2718 Some(dir) if !dir.as_os_str().is_empty() => dir.to_path_buf(),
2719 _ => std::path::PathBuf::from("."),
2720 };
2721 let tmp = dir.join(format!(
2722 ".readcon-write-{}-{}.tmp",
2723 std::process::id(),
2724 NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
2725 ));
2726 let file = OpenOptions::new()
2727 .read(true)
2728 .write(true)
2729 .create(true)
2730 .truncate(true)
2731 .open(&tmp)?;
2732 Ok((tmp, file))
2733}
2734
2735fn commit_temporary(
2736 tmp: &Path,
2737 dest: &Path,
2738 result: Result<(), WriteError>,
2739) -> Result<(), WriteError> {
2740 match result {
2741 Ok(()) => match std::fs::rename(tmp, dest) {
2742 Ok(()) => Ok(()),
2743 Err(err) => {
2744 let _ = std::fs::remove_file(tmp);
2745 Err(WriteError::Io(err))
2746 }
2747 },
2748 Err(err) => {
2749 let _ = std::fs::remove_file(tmp);
2750 Err(err)
2751 }
2752 }
2753}
2754
2755pub fn write_frame_refs_to_path(
2763 path: &std::path::Path,
2764 frames: &[&ConFrame],
2765 workers: usize,
2766 precision: Option<usize>,
2767 canonical: bool,
2768 bind: bool,
2769) -> Result<(), WriteError> {
2770 if workers < 2 {
2773 return write_serial_frames(path, frames, precision, canonical, bind);
2774 }
2775 if workers >= 2 && frames.len() >= 2 {
2776 return pipeline_frame_groups(path, frames, workers, precision, canonical, bind);
2777 }
2778 if workers >= 2 && frames.len() == 1 && frames[0].atom_data.len() >= WIDE_ATOM_GRAIN {
2779 let mut writer = match precision {
2780 Some(precision) => ConFrameWriter::with_precision(Vec::<u8>::new(), precision),
2781 None => ConFrameWriter::new(Vec::<u8>::new()),
2782 }
2783 .canonical(canonical)
2784 .parallel_atoms(workers);
2785 let (tmp, file) = temporary_sibling(path).map_err(WriteError::Io)?;
2788 writer.stream_file = Some(file);
2789 let result = if bind {
2790 writer.write_frame_bound(frames[0])
2791 } else {
2792 writer.write_frame(frames[0])
2793 };
2794 drop(writer);
2795 return commit_temporary(&tmp, path, result);
2796 }
2797 let parts = formatted_parts(frames, workers, precision, canonical, bind)?;
2798 if parts.len() == 1 {
2799 std::fs::write(path, &parts[0]).map_err(WriteError::Io)
2803 } else {
2804 pwrite_parts(path, &parts).map_err(WriteError::Io)
2805 }
2806}
2807
2808pub fn write_frame_refs_with_threads(
2812 frames: &[&ConFrame],
2813 workers: usize,
2814 precision: Option<usize>,
2815 canonical: bool,
2816) -> Result<Vec<u8>, WriteError> {
2817 let mut parts = formatted_parts(frames, workers, precision, canonical, false)?;
2818 if parts.len() == 1 {
2821 return Ok(parts.pop().unwrap_or_default());
2822 }
2823 let mut out = Vec::new();
2824 for part in parts {
2825 out.extend_from_slice(&part);
2826 }
2827 Ok(out)
2828}
2829
2830pub(crate) fn frame_byte_weight(frame: &ConFrame) -> usize {
2836 let atoms = frame.atom_data.len().max(1);
2837 if !frame.header.sections.is_empty() {
2841 return atoms.saturating_mul(1 + frame.header.sections.len());
2842 }
2843 let mut blocks = 1usize;
2844 if frame.velocities.is_present() {
2845 blocks += 1;
2846 }
2847 if frame.forces.is_present() {
2848 blocks += 1;
2849 }
2850 if frame.atom_energies.is_present() {
2851 blocks += 1;
2852 }
2853 if frame.charges.is_present() {
2854 blocks += 1;
2855 }
2856 if frame.spins.is_present() {
2857 blocks += 1;
2858 }
2859 if frame.magmoms.is_present() {
2860 blocks += 1;
2861 }
2862 if frame.displacements.is_present() {
2863 blocks += 1;
2864 }
2865 if frame.spreads.is_present() {
2866 blocks += 1;
2867 }
2868 if blocks == 1
2869 && let Some(atom) = frame.atom_data.first()
2870 {
2871 if atom.has_velocity() {
2872 blocks += 1;
2873 }
2874 if atom.has_forces() {
2875 blocks += 1;
2876 }
2877 if atom.has_energy() {
2878 blocks += 1;
2879 }
2880 if atom.has_charge() {
2881 blocks += 1;
2882 }
2883 if atom.has_spin() {
2884 blocks += 1;
2885 }
2886 if atom.has_magmom() {
2887 blocks += 1;
2888 }
2889 if atom.has_displacement() {
2890 blocks += 1;
2891 }
2892 if atom.has_spread() {
2893 blocks += 1;
2894 }
2895 }
2896 atoms.saturating_mul(blocks)
2897}
2898
2899fn equal_frame_ranges(n: usize, groups: usize) -> Vec<(usize, usize)> {
2900 let groups = groups.max(1).min(n.max(1));
2901 let mut out = Vec::with_capacity(groups);
2902 if n == 0 {
2903 return out;
2904 }
2905 for group in 0..groups {
2906 let start = group * n / groups;
2907 let end = (group + 1) * n / groups;
2908 if start < end {
2909 out.push((start, end));
2910 }
2911 }
2912 out
2913}
2914
2915pub(crate) fn group_frame_ranges(frames: &[&ConFrame], groups: usize) -> Vec<(usize, usize)> {
2922 let n = frames.len();
2923 if n == 0 {
2924 return Vec::new();
2925 }
2926 let groups = groups.max(1).min(n);
2927 if n == 1 || !sampled_weights_vary(frames) {
2928 return equal_frame_ranges(n, groups);
2929 }
2930 let weights: Vec<usize> = frames
2931 .iter()
2932 .map(|frame| frame_byte_weight(frame))
2933 .collect();
2934 pack_frames_by_weight(&weights, groups)
2935}
2936
2937fn sampled_weights_vary(frames: &[&ConFrame]) -> bool {
2938 let anchor = frame_byte_weight(frames[0]);
2939 let last = frames.len() - 1;
2940 let samples = 16.min(frames.len());
2941 for i in 0..samples {
2942 let index = if samples == 1 {
2943 0
2944 } else {
2945 i * last / (samples - 1)
2946 };
2947 if frame_byte_weight(frames[index]) != anchor {
2948 return true;
2949 }
2950 }
2951 false
2952}
2953
2954pub(crate) fn pack_frames_by_weight(weights: &[usize], workers: usize) -> Vec<(usize, usize)> {
2960 let n = weights.len();
2961 if n == 0 {
2962 return Vec::new();
2963 }
2964 let workers = workers.max(1).min(n);
2965 let mut prefix = vec![0u64; n + 1];
2966 for (i, weight) in weights.iter().enumerate() {
2967 prefix[i + 1] = prefix[i] + (*weight as u64).max(1);
2968 }
2969 let total = prefix[n];
2970 let mut out = Vec::with_capacity(workers);
2971 let mut start = 0usize;
2972 for worker in 0..workers {
2973 let end = if worker + 1 == workers {
2974 n
2975 } else {
2976 let target = ((worker as u64 + 1) * total) / workers as u64;
2977 let latest = n - (workers - worker - 1);
2978 let mut end = start + 1;
2979 while end < latest && prefix[end] < target {
2980 end += 1;
2981 }
2982 end
2983 };
2984 out.push((start, end));
2985 start = end;
2986 }
2987 out
2988}
2989
2990pub(crate) fn pipeline_format_threads(workers: usize) -> usize {
2995 workers.max(1)
2996}
2997
2998pub(crate) fn direct_slice_write(workers: usize, total_weight: usize) -> bool {
3002 workers >= 8 && total_weight >= 3_000_000
3003}
3004
3005fn written_box(frame: &ConFrame) -> ([f64; 3], [f64; 3]) {
3007 match frame.header.lattice_vectors() {
3008 Some(vecs) => crate::types::lengths_and_angles(vecs),
3009 None => (frame.header.boxl, frame.header.angles),
3010 }
3011}
3012
3013fn add_widths(sum: &mut usize, values: &[f64]) {
3014 *sum += sum_spelling_widths(values);
3015}
3016
3017fn milli_width_table() -> &'static [u8] {
3020 const N: usize = 1_000_001;
3021 static TABLE: std::sync::OnceLock<Box<[u8]>> = std::sync::OnceLock::new();
3022 TABLE.get_or_init(|| {
3023 let mut table = vec![0u8; N].into_boxed_slice();
3024 for n in 0..N {
3025 table[n] = spell_width((n as f64) / 10_000.0) as u8;
3026 }
3027 table
3028 })
3029}
3030
3031fn spell_width(value: f64) -> usize {
3032 if !value.is_finite() {
3033 return ryu::Buffer::new().format(value).len();
3034 }
3035 if let Some((_, n)) = encode_fixed_decimal_into(value) {
3036 return n;
3037 }
3038 let mut tmp = [0u8; 24];
3039 unsafe { ryu::raw::format64(value, tmp.as_mut_ptr()) }
3040}
3041
3042fn sum_spelling_widths(values: &[f64]) -> usize {
3044 let table = milli_width_table();
3045 let mut sum = 0usize;
3046 for &value in values {
3047 let bits = value.to_bits();
3048 let neg = bits >> 63 != 0;
3049 let mag = f64::from_bits(bits & !(1u64 << 63));
3050 let scaled = mag * 10_000.0;
3051 let nearest = scaled.round();
3052 if nearest <= 1_000_000.0 {
3053 let n = nearest as usize;
3054 if ((n as f64) / 10_000.0).to_bits() == mag.to_bits() {
3055 sum += table[n] as usize + usize::from(neg);
3056 continue;
3057 }
3058 }
3059 sum += spell_width(value);
3060 }
3061 sum
3062}
3063
3064fn packed_float_widths(frame: &ConFrame) -> usize {
3066 let mut sum = 0usize;
3067 let (boxl, angles) = written_box(frame);
3068 add_widths(&mut sum, &boxl);
3069 add_widths(&mut sum, &angles);
3070 add_widths(&mut sum, &frame.header.masses_per_type);
3071 if let Some(base) = coord_slab(&frame.atom_data) {
3072 let n = frame.atom_data.len() * 3;
3073 let coords = unsafe { std::slice::from_raw_parts(base, n) };
3074 add_widths(&mut sum, coords);
3075 } else {
3076 for atom in &frame.atom_data {
3077 sum += round_trip_width(atom.x());
3078 sum += round_trip_width(atom.y());
3079 sum += round_trip_width(atom.z());
3080 }
3081 }
3082 sum += optional_rows_width(frame.velocities.is_present().then(|| &*frame.velocities));
3083 sum += optional_rows_width(frame.forces.is_present().then(|| &*frame.forces));
3084 sum += optional_rows_width(frame.magmoms.is_present().then(|| &*frame.magmoms));
3085 sum += optional_rows_width(
3086 frame
3087 .displacements
3088 .is_present()
3089 .then(|| &*frame.displacements),
3090 );
3091 sum += optional_rows_width(frame.spreads.is_present().then(|| &*frame.spreads));
3092 sum += optional_column_width(
3093 frame
3094 .atom_energies
3095 .is_present()
3096 .then(|| &*frame.atom_energies),
3097 );
3098 sum += optional_column_width(frame.charges.is_present().then(|| &*frame.charges));
3099 sum += optional_column_width(frame.spins.is_present().then(|| &*frame.spins));
3100 if !frame.velocities.is_present() && frame.has_velocities() {
3102 for atom in &frame.atom_data {
3103 let [x, y, z] = atom.velocity().unwrap_or([0.0; 3]);
3104 sum += round_trip_width(x) + round_trip_width(y) + round_trip_width(z);
3105 }
3106 }
3107 if !frame.forces.is_present() && frame.has_forces() {
3108 for atom in &frame.atom_data {
3109 let [x, y, z] = atom.force().unwrap_or([0.0; 3]);
3110 sum += round_trip_width(x) + round_trip_width(y) + round_trip_width(z);
3111 }
3112 }
3113 sum
3114}
3115
3116fn optional_rows_width(storage: Option<&crate::storage_dtype::Array2Storage>) -> usize {
3117 let Some(storage) = storage else {
3118 return 0;
3119 };
3120 let Some(values) = storage.f64_slice() else {
3121 return 0;
3122 };
3123 let mut sum = 0usize;
3124 add_widths(&mut sum, values);
3125 sum
3126}
3127
3128fn optional_column_width(storage: Option<&crate::storage_dtype::Array1Storage>) -> usize {
3129 let Some(storage) = storage else {
3130 return 0;
3131 };
3132 let Some(values) = storage.f64_slice() else {
3133 return 0;
3134 };
3135 let mut sum = 0usize;
3136 add_widths(&mut sum, values);
3137 sum
3138}
3139
3140fn same_packed_shape(reference: &ConFrame, frame: &ConFrame) -> bool {
3142 if frame.atom_data.len() != reference.atom_data.len()
3143 || frame.atom_data.is_empty()
3144 || frame.header.spec_version != reference.header.spec_version
3145 || frame.header.natm_types != reference.header.natm_types
3146 || frame.header.natms_per_type != reference.header.natms_per_type
3147 || frame.header.masses_per_type.len() != reference.header.masses_per_type.len()
3148 || frame.header.prebox_header.user != reference.header.prebox_header.user
3149 || frame.header.prebox_header.metadata_line()
3150 != reference.header.prebox_header.metadata_line()
3151 || frame.header.postbox_header != reference.header.postbox_header
3152 || frame.has_velocities() != reference.has_velocities()
3153 || frame.has_forces() != reference.has_forces()
3154 || frame.has_energies() != reference.has_energies()
3155 || frame.has_charges() != reference.has_charges()
3156 || frame.has_spins() != reference.has_spins()
3157 || frame.has_magmoms() != reference.has_magmoms()
3158 || frame.has_displacements() != reference.has_displacements()
3159 || frame.has_spreads() != reference.has_spreads()
3160 || !ids_are_dense(&frame.atom_data, 1)
3161 {
3162 return false;
3163 }
3164 let mut offset = 0usize;
3165 for &count in &frame.header.natms_per_type {
3166 if frame.atom_data[offset].symbol != reference.atom_data[offset].symbol {
3167 return false;
3168 }
3169 offset += count;
3170 }
3171 true
3172}
3173
3174fn packed_frame_lengths(
3181 frames: &[&ConFrame],
3182 workers: usize,
3183 precision: Option<usize>,
3184 canonical: bool,
3185 bind: bool,
3186) -> Option<Vec<usize>> {
3187 if precision.is_some() || frames.len() < 2 {
3189 return None;
3190 }
3191 let reference = frames[0];
3192 if reference.atom_data.is_empty() || !ids_are_dense(&reference.atom_data, 1) {
3193 return None;
3194 }
3195 let workers = workers.max(1).min(frames.len());
3196 let shape_ok = std::thread::scope(|scope| {
3197 let mut guards = Vec::with_capacity(workers);
3198 for thread in 0..workers {
3199 let start = thread * frames.len() / workers;
3200 let end = (thread + 1) * frames.len() / workers;
3201 guards.push(scope.spawn(move || {
3202 frames[start..end]
3203 .iter()
3204 .all(|frame| same_packed_shape(reference, frame))
3205 }));
3206 }
3207 guards
3208 .into_iter()
3209 .all(|guard| guard.join().unwrap_or(false))
3210 });
3211 if !shape_ok {
3212 return None;
3213 }
3214 let widths = std::thread::scope(|scope| {
3215 let mut guards = Vec::with_capacity(workers);
3216 for thread in 0..workers {
3217 let start = thread * frames.len() / workers;
3218 let end = (thread + 1) * frames.len() / workers;
3219 guards.push(scope.spawn(move || {
3220 #[cfg(feature = "parallel")]
3221 crate::parser::pin_worker(thread);
3222 frames[start..end]
3223 .iter()
3224 .map(|frame| packed_float_widths(frame))
3225 .collect::<Vec<_>>()
3226 }));
3227 }
3228 let mut widths = Vec::with_capacity(frames.len());
3229 for guard in guards {
3230 widths.extend(guard.join().unwrap_or_default());
3231 }
3232 widths
3233 });
3234 if widths.len() != frames.len() {
3235 return None;
3236 }
3237 let sample = format_group_bytes(
3238 std::slice::from_ref(&reference),
3239 1,
3240 precision,
3241 canonical,
3242 bind,
3243 )
3244 .ok()?;
3245 let overhead = sample.len().checked_sub(widths[0])?;
3246 let lengths: Vec<usize> = widths.iter().map(|width| overhead + width).collect();
3247 for &index in [0, frames.len() / 2, frames.len() - 1].iter() {
3248 let got = format_group_bytes(
3249 std::slice::from_ref(&frames[index]),
3250 1,
3251 precision,
3252 canonical,
3253 bind,
3254 )
3255 .ok()?
3256 .len();
3257 if got != lengths[index] {
3258 return None;
3259 }
3260 }
3261 Some(lengths)
3262}
3263
3264fn format_group_into(
3267 group: &[&ConFrame],
3268 dest: &mut [u8],
3269 precision: Option<usize>,
3270 canonical: bool,
3271 bind: bool,
3272) -> Result<(), WriteError> {
3273 if dest.is_empty() {
3274 return Ok(());
3275 }
3276 struct Held {
3277 writer: ConFrameWriter<Vec<u8>>,
3278 }
3279 impl Drop for Held {
3280 fn drop(&mut self) {
3281 std::mem::forget(std::mem::take(&mut self.writer.scratch));
3283 }
3284 }
3285 let cap = dest.len();
3286 let ptr = dest.as_mut_ptr();
3287 let mut held = Held {
3288 writer: match precision {
3289 Some(precision) => ConFrameWriter::with_precision(Vec::new(), precision),
3290 None => ConFrameWriter::new(Vec::new()),
3291 }
3292 .canonical(canonical)
3293 .retain_scratch(true),
3294 };
3295 held.writer.exact_scratch = true;
3296 let previous = std::mem::replace(&mut held.writer.scratch, unsafe {
3297 Vec::from_raw_parts(ptr, 0, cap)
3298 });
3299 drop(previous);
3300 let result = write_group(&mut held.writer, group, bind);
3301 let written = held.writer.scratch.len();
3302 let stable = held.writer.scratch.as_ptr() == ptr;
3303 result?;
3304 if written != cap || !stable {
3305 return Err(WriteError::Io(io::Error::other(
3306 "packed frame length drifted",
3307 )));
3308 }
3309 Ok(())
3310}
3311
3312fn place_packed(
3314 file: &File,
3315 frames: &[&ConFrame],
3316 lengths: &[usize],
3317 workers: usize,
3318 precision: Option<usize>,
3319 canonical: bool,
3320 bind: bool,
3321) -> Result<(), WriteError> {
3322 let total: u64 = lengths.iter().map(|len| *len as u64).sum();
3323 file.set_len(total).map_err(WriteError::Io)?;
3324 if total == 0 {
3325 return Ok(());
3326 }
3327 let ranges = pack_frames_by_weight(lengths, workers);
3328 let mut prefix = vec![0usize; lengths.len() + 1];
3329 for (index, len) in lengths.iter().enumerate() {
3330 prefix[index + 1] = prefix[index] + len;
3331 }
3332 if cfg!(miri) {
3335 let mut parts = Vec::with_capacity(ranges.len());
3336 for &(start, end) in &ranges {
3337 parts.push(format_group_bytes(
3338 &frames[start..end],
3339 1,
3340 precision,
3341 canonical,
3342 bind,
3343 )?);
3344 }
3345 return store_disjoint_parts(file, &parts).map_err(WriteError::Io);
3346 }
3347 let result = std::thread::scope(|scope| {
3350 let mut guards = Vec::with_capacity(ranges.len());
3351 for (thread, &(start, end)) in ranges.iter().enumerate() {
3352 let offset = prefix[start];
3353 let span = prefix[end] - prefix[start];
3354 let piece = &frames[start..end];
3355 let cloned = file.try_clone().map_err(WriteError::Io)?;
3356 guards.push(scope.spawn(move || -> Result<(), WriteError> {
3357 #[cfg(feature = "parallel")]
3358 crate::parser::pin_worker(thread);
3359 if span == 0 {
3360 return Ok(());
3361 }
3362 let page = 4096usize;
3363 let map_off = offset & !(page - 1);
3364 let delta = offset - map_off;
3365 let map_len = delta + span;
3366 let mut map = unsafe {
3367 memmap2::MmapOptions::new()
3368 .offset(map_off as u64)
3369 .len(map_len)
3370 .map_mut(&cloned)
3371 .map_err(WriteError::Io)?
3372 };
3373 advise_shared_store(&cloned, map.as_mut_ptr(), map.len());
3374 let dest = &mut map[delta..delta + span];
3375 format_group_into(piece, dest, precision, canonical, bind)?;
3376 drop(map);
3377 Ok(())
3378 }));
3379 }
3380 for guard in guards {
3381 match guard.join() {
3382 Ok(result) => result?,
3383 Err(_) => {
3384 return Err(WriteError::Io(io::Error::other("frame formatter panicked")));
3385 }
3386 }
3387 }
3388 Ok(())
3389 });
3390 result
3391}
3392
3393fn write_serial_frames(
3397 path: &std::path::Path,
3398 frames: &[&ConFrame],
3399 precision: Option<usize>,
3400 canonical: bool,
3401 bind: bool,
3402) -> Result<(), WriteError> {
3403 let (tmp, file) = temporary_sibling(path).map_err(WriteError::Io)?;
3404 let mut writer = match precision {
3405 Some(precision) => ConFrameWriter::with_precision(file, precision),
3406 None => ConFrameWriter::new(file),
3407 }
3408 .canonical(canonical);
3409 let result =
3410 write_group(&mut writer, frames, bind).and_then(|_| writer.flush().map_err(WriteError::Io));
3411 drop(writer);
3412 commit_temporary(&tmp, path, result)
3413}
3414
3415fn pipeline_frame_groups(
3416 path: &std::path::Path,
3417 frames: &[&ConFrame],
3418 workers: usize,
3419 precision: Option<usize>,
3420 canonical: bool,
3421 bind: bool,
3422) -> Result<(), WriteError> {
3423 if (2..=workers).contains(&frames.len()) {
3427 let mut parts: Vec<Result<Vec<u8>, WriteError>> = Vec::with_capacity(frames.len());
3428 std::thread::scope(|scope| {
3429 let mut guards = Vec::with_capacity(frames.len());
3430 for (index, frame) in frames.iter().enumerate() {
3431 guards.push(scope.spawn(move || {
3432 #[cfg(feature = "parallel")]
3433 crate::parser::pin_worker(index);
3434 format_group_bytes(std::slice::from_ref(frame), 1, precision, canonical, bind)
3435 }));
3436 }
3437 for guard in guards {
3438 parts.push(match guard.join() {
3439 Ok(part) => part,
3440 Err(_) => Err(WriteError::Io(io::Error::other("frame formatter panicked"))),
3441 });
3442 }
3443 });
3444 let mut formatted = Vec::with_capacity(parts.len());
3445 for part in parts {
3446 formatted.push(part?);
3447 }
3448 let (tmp, file) = temporary_sibling(path).map_err(WriteError::Io)?;
3449 let result = store_disjoint_parts(&file, &formatted).map_err(WriteError::Io);
3450 return commit_temporary(&tmp, path, result);
3451 }
3452 let heavy = direct_slice_write(
3456 workers,
3457 frame_byte_weight(frames[0]).saturating_mul(frames.len()),
3458 );
3459 if heavy
3460 && let Some(lengths) = cached_packed_lengths(frames, workers, precision, canonical, bind)
3461 {
3462 let (tmp, file) = temporary_sibling(path).map_err(WriteError::Io)?;
3463 let result = place_packed(&file, frames, &lengths, workers, precision, canonical, bind);
3464 return commit_temporary(&tmp, path, result);
3465 }
3466 let groups = workers.saturating_mul(4).max(workers);
3470 let ranges = group_frame_ranges(frames, groups);
3471 let pieces: Vec<&[&ConFrame]> = ranges
3472 .iter()
3473 .map(|&(start, end)| &frames[start..end])
3474 .collect();
3475 let format_threads = pipeline_format_threads(workers).min(pieces.len().max(1));
3476 let (tx, rx) = std::sync::mpsc::sync_channel::<(usize, Result<Vec<u8>, WriteError>)>(2);
3479 let next_job = std::sync::atomic::AtomicUsize::new(0);
3480 let (tmp, file) = temporary_sibling(path).map_err(WriteError::Io)?;
3481 let result = std::thread::scope(|scope| {
3482 for _ in 0..format_threads {
3483 let tx = tx.clone();
3484 let pieces = &pieces;
3485 let next_job = &next_job;
3486 scope.spawn(move || {
3487 loop {
3488 let index = next_job.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3489 if index >= pieces.len() {
3490 break;
3491 }
3492 let message = format_group_bytes(pieces[index], 1, precision, canonical, bind);
3493 if tx.send((index, message)).is_err() {
3494 break;
3495 }
3496 }
3497 });
3498 }
3499 drop(tx);
3500 let mut file = file;
3501 let mut next = 0usize;
3502 let mut pending = BTreeMap::<usize, Vec<u8>>::new();
3503 while next < pieces.len() {
3504 let (index, part) = rx
3505 .recv()
3506 .map_err(|err| WriteError::Io(io::Error::other(err.to_string())))?;
3507 pending.insert(index, part?);
3508 while let Some(part) = pending.remove(&next) {
3509 file.write_all(&part).map_err(WriteError::Io)?;
3510 next += 1;
3511 }
3512 }
3513 file.flush().map_err(WriteError::Io)?;
3514 Ok(())
3515 });
3516 commit_temporary(&tmp, path, result)
3517}
3518
3519struct PackedCache {
3520 first: usize,
3521 last: usize,
3522 count: usize,
3523 precision: Option<usize>,
3524 canonical: bool,
3525 bind: bool,
3526 generation: u64,
3527 lengths: Vec<usize>,
3528}
3529
3530static PACKED_GEN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
3531static PACKED_FIRST: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
3532static PACKED_LAST: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
3533
3534pub(crate) fn invalidate_packed_lengths() {
3536 PACKED_GEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3537}
3538
3539pub(crate) fn invalidate_packed_lengths_if(frame: usize) {
3542 let first = PACKED_FIRST.load(std::sync::atomic::Ordering::Relaxed);
3543 if first != 0
3544 && (frame == first || frame == PACKED_LAST.load(std::sync::atomic::Ordering::Relaxed))
3545 {
3546 invalidate_packed_lengths();
3547 PACKED_FIRST.store(0, std::sync::atomic::Ordering::Relaxed);
3548 }
3549}
3550
3551fn cached_packed_lengths(
3552 frames: &[&ConFrame],
3553 workers: usize,
3554 precision: Option<usize>,
3555 canonical: bool,
3556 bind: bool,
3557) -> Option<Vec<usize>> {
3558 let first = frames[0] as *const ConFrame as usize;
3559 let last = *frames.last()? as *const ConFrame as usize;
3560 let count = frames.len();
3561 static CACHE: std::sync::Mutex<Option<PackedCache>> = std::sync::Mutex::new(None);
3562 let generation = PACKED_GEN.load(std::sync::atomic::Ordering::Relaxed);
3563 if let Some(hit) = CACHE.lock().unwrap_or_else(|err| err.into_inner()).as_ref()
3564 && hit.generation == generation
3565 && hit.first == first
3566 && hit.last == last
3567 && hit.count == count
3568 && hit.precision == precision
3569 && hit.canonical == canonical
3570 && hit.bind == bind
3571 {
3572 return Some(hit.lengths.clone());
3573 }
3574 let generation_before = PACKED_GEN.load(std::sync::atomic::Ordering::Relaxed);
3575 let lengths = packed_frame_lengths(frames, workers, precision, canonical, bind)?;
3576 let generation_after = PACKED_GEN.load(std::sync::atomic::Ordering::Relaxed);
3577 if generation_before == generation_after {
3578 PACKED_FIRST.store(first, std::sync::atomic::Ordering::Relaxed);
3579 PACKED_LAST.store(last, std::sync::atomic::Ordering::Relaxed);
3580 *CACHE.lock().unwrap_or_else(|err| err.into_inner()) = Some(PackedCache {
3581 first,
3582 last,
3583 count,
3584 precision,
3585 canonical,
3586 bind,
3587 generation: generation_after,
3588 lengths: lengths.clone(),
3589 });
3590 }
3591 Some(lengths)
3592}
3593
3594fn store_disjoint_parts(file: &File, parts: &[Vec<u8>]) -> io::Result<()> {
3598 let mut offsets = Vec::with_capacity(parts.len());
3599 let mut total = 0u64;
3600 for part in parts {
3601 offsets.push(total);
3602 total += part.len() as u64;
3603 }
3604 file.set_len(total)?;
3605 if total == 0 {
3606 return Ok(());
3607 }
3608 if file_is_tmpfs(file) {
3609 store_mapped_parts(file, parts, &offsets)
3610 } else {
3611 std::thread::scope(|scope| {
3612 let mut guards = Vec::new();
3613 for (part, at) in parts.iter().zip(offsets) {
3614 let cloned = file.try_clone()?;
3615 guards.push(scope.spawn(move || pwrite_full(&cloned, part, at)));
3616 }
3617 for guard in guards {
3618 guard
3619 .join()
3620 .unwrap_or_else(|_| Err(io::Error::other("pwrite worker")))?;
3621 }
3622 Ok(())
3623 })
3624 }
3625}
3626
3627fn store_mapped_parts(file: &File, parts: &[Vec<u8>], offsets: &[u64]) -> io::Result<()> {
3628 let total = offsets
3629 .last()
3630 .copied()
3631 .unwrap_or(0)
3632 .saturating_add(parts.last().map(|part| part.len() as u64).unwrap_or(0));
3633 let mut map = unsafe {
3634 memmap2::MmapOptions::new()
3635 .len(total as usize)
3636 .map_mut(file)?
3637 };
3638 let base = map.as_mut_ptr() as usize;
3639 std::thread::scope(|scope| {
3640 for (part, at) in parts.iter().zip(offsets) {
3641 let src = part.as_ptr() as usize;
3642 let len = part.len();
3643 let dest = base + *at as usize;
3644 scope.spawn(move || {
3645 if len == 0 {
3646 return;
3647 }
3648 unsafe {
3651 std::ptr::copy_nonoverlapping(src as *const u8, dest as *mut u8, len);
3652 }
3653 });
3654 }
3655 });
3656 Ok(())
3657}
3658
3659fn write_group<W: Write>(
3660 writer: &mut ConFrameWriter<W>,
3661 group: &[&ConFrame],
3662 bind: bool,
3663) -> Result<(), WriteError> {
3664 if bind {
3665 for frame in group {
3666 writer.write_frame_bound(frame)?;
3667 }
3668 return Ok(());
3669 }
3670 writer.extend(group.iter().copied())
3671}
3672
3673fn format_group_bytes(
3674 group: &[&ConFrame],
3675 workers: usize,
3676 precision: Option<usize>,
3677 canonical: bool,
3678 bind: bool,
3679) -> Result<Vec<u8>, WriteError> {
3680 let mut buffer = Vec::new();
3681 let mut writer = match precision {
3682 Some(precision) => ConFrameWriter::with_precision(&mut buffer, precision),
3683 None => ConFrameWriter::new(&mut buffer),
3684 }
3685 .canonical(canonical)
3686 .retain_scratch(true);
3687 if workers >= 2 {
3688 writer = writer.parallel_atoms(workers);
3689 }
3690 write_group(&mut writer, group, bind)?;
3691 Ok(std::mem::take(&mut writer.scratch))
3692}
3693
3694fn formatted_parts(
3695 frames: &[&ConFrame],
3696 workers: usize,
3697 precision: Option<usize>,
3698 canonical: bool,
3699 bind: bool,
3700) -> Result<Vec<Vec<u8>>, WriteError> {
3701 let wide = frames.len() == 1 && frames[0].atom_data.len() >= WIDE_ATOM_GRAIN;
3702 if workers < 2 || (frames.len() < 2 && !wide) {
3703 return Ok(vec![format_group_bytes(
3704 frames, 1, precision, canonical, bind,
3705 )?]);
3706 }
3707 if wide {
3708 return Ok(vec![format_group_bytes(
3709 frames, workers, precision, canonical, bind,
3710 )?]);
3711 }
3712 #[cfg(feature = "parallel")]
3713 {
3714 let chunk = frames.len().div_ceil(workers).max(1);
3715 return crate::parser::on_pool(workers, || {
3716 use rayon::prelude::*;
3717 frames
3718 .par_chunks(chunk)
3719 .map(|group| format_group_bytes(group, 1, precision, canonical, bind))
3720 .collect()
3721 });
3722 }
3723 #[cfg(not(feature = "parallel"))]
3724 {
3725 let _ = workers;
3726 Ok(vec![format_group_bytes(
3727 frames, 1, precision, canonical, bind,
3728 )?])
3729 }
3730}
3731
3732fn pwrite_parts(path: &std::path::Path, parts: &[Vec<u8>]) -> io::Result<()> {
3733 let slices: Vec<&[u8]> = parts.iter().map(|part| part.as_slice()).collect();
3734 pwrite_borrowed(path, &slices)
3735}
3736
3737fn pwrite_borrowed(path: &std::path::Path, parts: &[&[u8]]) -> io::Result<()> {
3738 let mut offset = 0u64;
3739 let mut offsets = Vec::with_capacity(parts.len());
3740 for part in parts {
3741 offsets.push(offset);
3742 offset += part.len() as u64;
3743 }
3744 let file = OpenOptions::new()
3745 .write(true)
3746 .create(true)
3747 .truncate(true)
3748 .open(path)?;
3749 file.set_len(offset)?;
3750 std::thread::scope(|scope| {
3751 let mut handles = Vec::new();
3752 for (part, at) in parts.iter().zip(offsets) {
3753 let file = file.try_clone()?;
3754 handles.push(scope.spawn(move || pwrite_full(&file, part, at)));
3755 }
3756 for handle in handles {
3757 handle
3758 .join()
3759 .unwrap_or_else(|_| Err(io::Error::other("pwrite worker")))?;
3760 }
3761 Ok(())
3762 })
3763}
3764
3765fn pwrite_full(file: &File, buf: &[u8], mut offset: u64) -> io::Result<()> {
3766 if buf.is_empty() {
3767 return Ok(());
3768 }
3769 #[cfg(miri)]
3772 {
3773 use std::io::{Seek, SeekFrom, Write};
3774 let mut out = file.try_clone()?;
3775 out.seek(SeekFrom::Start(offset))?;
3776 out.write_all(buf)?;
3777 return Ok(());
3778 }
3779 #[cfg(all(not(miri), unix))]
3780 {
3781 use std::os::unix::io::AsRawFd;
3782 let fd = file.as_raw_fd();
3783 let mut written = 0usize;
3784 while written < buf.len() {
3785 let n = unsafe {
3786 libc::pwrite(
3787 fd,
3788 buf[written..].as_ptr() as *const libc::c_void,
3789 buf.len() - written,
3790 offset as libc::off_t,
3791 )
3792 };
3793 if n < 0 {
3794 let err = io::Error::last_os_error();
3795 if err.kind() == io::ErrorKind::Interrupted {
3796 continue;
3797 }
3798 return Err(err);
3799 }
3800 if n == 0 {
3801 return Err(io::Error::new(io::ErrorKind::WriteZero, "pwrite"));
3802 }
3803 written += n as usize;
3804 offset += n as u64;
3805 }
3806 Ok(())
3807 }
3808 #[cfg(all(not(miri), windows))]
3809 {
3810 use std::os::windows::fs::FileExt;
3811 file.seek_write(buf, offset)?;
3812 Ok(())
3813 }
3814}
3815
3816fn write_at_offset(file: &File, buf: &[u8], offset: u64) -> io::Result<usize> {
3817 #[cfg(unix)]
3818 {
3819 use std::os::unix::fs::FileExt;
3820 return file.write_at(buf, offset);
3821 }
3822 #[cfg(windows)]
3823 {
3824 use std::os::windows::fs::FileExt;
3825 file.seek_write(buf, offset)
3826 }
3827}
3828
3829pub fn write_frames_to_string(frames: &[ConFrame]) -> Result<String, WriteError> {
3833 let bytes = write_frames_to_vec(frames)?;
3834 String::from_utf8(bytes).map_err(|err| {
3835 WriteError::Io(io::Error::new(
3836 io::ErrorKind::InvalidData,
3837 err.utf8_error().to_string(),
3838 ))
3839 })
3840}
3841
3842impl ConFrameWriter<flate2::write::GzEncoder<File>> {
3844 pub fn from_path_gzip<P: AsRef<Path>>(path: P) -> io::Result<Self> {
3846 let encoder = crate::compression::gzip_writer(path.as_ref())?;
3847 Ok(Self::new(encoder))
3848 }
3849
3850 pub fn from_path_gzip_with_precision<P: AsRef<Path>>(
3852 path: P,
3853 precision: usize,
3854 ) -> io::Result<Self> {
3855 let encoder = crate::compression::gzip_writer(path.as_ref())?;
3856 Ok(Self::with_precision(encoder, precision))
3857 }
3858}
3859
3860#[cfg(feature = "zstd")]
3863impl ConFrameWriter<zstd::stream::write::AutoFinishEncoder<'static, File>> {
3864 pub fn from_path_zstd<P: AsRef<Path>>(path: P) -> io::Result<Self> {
3866 let encoder = crate::compression::zstd_writer(path.as_ref())?;
3867 Ok(Self::new(encoder))
3868 }
3869
3870 pub fn from_path_zstd_with_precision<P: AsRef<Path>>(
3872 path: P,
3873 precision: usize,
3874 ) -> io::Result<Self> {
3875 let encoder = crate::compression::zstd_writer(path.as_ref())?;
3876 Ok(Self::with_precision(encoder, precision))
3877 }
3878}
3879
3880#[cfg(test)]
3881mod float_format_tests {
3882 use super::{encode_fixed_decimal_into, push_f64_prec, push_u64, write_finite};
3883
3884 fn formatted(v: f64, prec: usize) -> String {
3885 let mut buf = Vec::new();
3886 push_f64_prec(&mut buf, v, prec);
3887 String::from_utf8(buf).expect("utf8")
3888 }
3889
3890 #[test]
3891 fn matches_std_fixed_precision() {
3892 let vals = [
3893 0.0,
3894 -0.0,
3895 1.0,
3896 -1.0,
3897 0.9045,
3898 6.975_299_999_999_995,
3899 63.546,
3900 1.008,
3901 15.3456,
3902 90.0,
3903 218.0,
3904 1e-6,
3905 -1.23456789,
3906 10.0,
3907 ];
3908 for prec in [0usize, 6] {
3909 for v in vals {
3910 let got = formatted(v, prec);
3911 let exp = format!("{v:.prec$}");
3912 assert_eq!(got, exp, "v={v:?} prec={prec}");
3913 }
3914 }
3915 }
3916
3917 #[test]
3918 fn fixed_decimal_matches_ryu_including_edges() {
3919 let mut ryu_buf = ryu::Buffer::new();
3920 let mut values = vec![
3921 0.0,
3922 -0.0,
3923 1.0,
3924 -1.0,
3925 10.0,
3926 97.0,
3927 -97.0,
3928 63.546,
3929 0.5,
3930 0.125,
3931 1.5,
3932 -0.1,
3933 0.1 + 0.2,
3934 1.0 / 3.0,
3935 1e15,
3936 -1e15,
3937 1e16,
3938 1e20,
3939 -1e20,
3940 1e-6,
3941 1e-10,
3942 f64::MIN_POSITIVE,
3943 f64::MAX,
3944 f64::MIN,
3945 f64::INFINITY,
3946 f64::NEG_INFINITY,
3947 f64::NAN,
3948 ];
3949 let mut state = 0x1234_5678_9abc_def0u64;
3950 for _ in 0..256 {
3951 state = state.wrapping_mul(0x5851_f42d_4c95_7f2d).wrapping_add(1);
3952 values.push(f64::from_bits(state));
3953 }
3954 for value in values {
3955 let mut got = Vec::new();
3956 let mut scratch = ryu::Buffer::new();
3957 super::push_round_trip(&mut got, &mut scratch, None, value);
3958 let expect = ryu_buf.format(value);
3959 if value.is_nan() {
3960 assert_eq!(std::str::from_utf8(&got).unwrap(), expect);
3961 continue;
3962 }
3963 assert_eq!(std::str::from_utf8(&got).unwrap(), expect, "{value:?}");
3964 }
3965 }
3966
3967 #[test]
3968 fn round_trip_spelling_matches_ryu_on_repeated_and_unique_rows() {
3969 use crate::iterators::ConFrameIterator;
3970 use crate::writer::ConFrameWriter;
3971 let mut text = String::from(
3972 "bench\n{\"con_spec_version\":2}\n10.0 10.0 10.0\n90.0 90.0 90.0\n0 0\n0 0 1\n1\n6\n12.0\nCu\nCoordinates of Component 1\n",
3973 );
3974 for i in 0..6 {
3975 let x = (i % 2) as f64 * 1.5 + 0.25;
3976 let y = 3.25;
3977 let z = if i == 3 {
3978 -0.0
3979 } else {
3980 (i as f64) * 0.0010000001
3981 };
3982 text.push_str(&format!(" {x:.17} {y:.17} {z:.17} 0 {i}\n"));
3983 }
3984 text.push_str("\nCu\nVelocities of Component 1\n");
3986 for i in 0..6 {
3987 let vx = 0.125;
3988 let vy = if i % 2 == 0 { 0.5 } else { -0.5 };
3989 let vz = (i as f64) * 1.0e-4;
3990 text.push_str(&format!(" {vx:.17} {vy:.17} {vz:.17} 0 {i}\n"));
3991 }
3992 let frames: Vec<_> = ConFrameIterator::new(&text)
3993 .map(|frame| frame.expect("parse"))
3994 .collect();
3995 let mut values = Vec::new();
3996 for atom in &frames[0].atom_data {
3997 values.extend([atom.x(), atom.y(), atom.z()]);
3998 }
3999 for atom in &frames[0].atom_data {
4000 if let Some(v) = atom.velocity() {
4001 values.extend(v);
4002 }
4003 }
4004 let mut buf = Vec::new();
4005 {
4006 let mut writer = ConFrameWriter::new(&mut buf);
4007 writer.extend(frames.iter()).unwrap();
4008 }
4009 let written = String::from_utf8(buf).unwrap();
4010 let mut ryu = ryu::Buffer::new();
4011 let mut spelled = Vec::new();
4012 for line in written.lines() {
4013 let toks: Vec<_> = line.split_whitespace().collect();
4014 if toks.len() >= 5
4015 && toks[0].contains('.')
4016 && toks[3].bytes().all(|b| b.is_ascii_digit())
4017 {
4018 spelled.extend(toks[..3].iter().map(|tok| (*tok).to_string()));
4019 }
4020 }
4021 assert_eq!(spelled.len(), values.len());
4022 for (got, value) in spelled.iter().zip(values) {
4023 let expect = if value.is_finite() {
4024 ryu.format_finite(value).to_string()
4025 } else {
4026 ryu.format(value).to_string()
4027 };
4028 assert_eq!(got, &expect, "value {value:?}");
4029 }
4030 }
4031
4032 #[test]
4033 fn fixed_decimal_matches_ryu_or_falls_back_to_it() {
4034 let mut values = vec![
4035 0.0,
4036 -0.0,
4037 1.0,
4038 -2.0,
4039 10.0,
4040 99.0,
4041 100.0,
4042 0.1,
4043 0.01,
4044 0.001,
4045 0.0001,
4046 0.0002,
4047 -0.0001,
4048 1.5,
4049 0.25,
4050 0.125,
4051 63.546,
4052 12.0,
4053 90.0,
4054 3.25,
4055 0.5,
4056 -0.5,
4057 1.0e-4,
4058 1.25,
4059 218.0,
4060 6.975_299_999_999_995,
4061 1.008,
4062 15.3456,
4063 -1.23456789,
4064 1e-6,
4065 1e16,
4066 1e20,
4067 ];
4068 for i in 0..200 {
4069 values.push(i as f64);
4070 values.push((i as f64) * 0.001);
4071 values.push((i as f64) * 0.0001);
4072 values.push(-(i as f64) * 0.001);
4073 }
4074 let mut ryu = ryu::Buffer::new();
4075 for value in values {
4076 let mut buf = Vec::new();
4077 write_finite(&mut buf, value);
4078 let expect = ryu.format_finite(value);
4079 assert_eq!(std::str::from_utf8(&buf).unwrap(), expect, "{value:?}");
4080 if encode_fixed_decimal_into(value).is_none() {
4081 continue;
4082 }
4083 }
4084 }
4085
4086 #[test]
4087 fn u64_matches_decimal_display() {
4088 let mut values = vec![0, 1, 9, 10, 11, 99, 100, 101, 999, 1000, 50000, u64::MAX];
4089 values.extend(0..200);
4090 for n in values {
4091 let mut buf = Vec::new();
4092 push_u64(&mut buf, n);
4093 assert_eq!(std::str::from_utf8(&buf).unwrap(), n.to_string());
4094 }
4095 }
4096}
4097
4098#[cfg(test)]
4099mod section_arity_tests {
4100 use super::require_minor2_section_arity;
4101 use serde_json::json;
4102
4103 #[test]
4104 fn minor2_lists_a_new_section_in_section_arity() {
4105 let widgets = vec![json!("widgets")];
4106 let err = require_minor2_section_arity(2, &widgets, &serde_json::Map::new()).unwrap_err();
4107 let msg = err.to_string();
4108 assert!(msg.contains("section_arity"), "{msg}");
4109
4110 let mut listed = serde_json::Map::new();
4111 listed.insert("section_arity".into(), json!({"widgets": 3}));
4112 require_minor2_section_arity(2, &widgets, &listed).unwrap();
4113
4114 let reserved = vec![json!("velocities"), json!("energies")];
4115 require_minor2_section_arity(2, &reserved, &serde_json::Map::new()).unwrap();
4116 require_minor2_section_arity(1, &widgets, &serde_json::Map::new()).unwrap();
4117 }
4118}
4119
4120#[cfg(test)]
4121mod weight_partition {
4122 use super::pack_frames_by_weight;
4123 use proptest::prelude::*;
4124
4125 fn is_ordered_partition(n: usize, workers: usize, ranges: &[(usize, usize)]) -> bool {
4128 if n == 0 {
4129 return ranges.is_empty();
4130 }
4131 let workers = workers.max(1).min(n);
4132 if ranges.len() != workers {
4133 return false;
4134 }
4135 let mut cursor = 0usize;
4136 for &(start, end) in ranges {
4137 if start != cursor || end <= start || end > n {
4138 return false;
4139 }
4140 cursor = end;
4141 }
4142 cursor == n
4143 }
4144
4145 proptest! {
4146 #![proptest_config(ProptestConfig {
4147 cases: 64,
4148 failure_persistence: None,
4149 ..ProptestConfig::default()
4150 })]
4151 #[test]
4152 fn byte_weights_do_not_drop_or_reorder_frames(
4153 weights in prop::collection::vec(0usize..50_000, 0..48),
4154 workers in 1usize..17,
4155 ) {
4156 let ranges = pack_frames_by_weight(&weights, workers);
4157 prop_assert!(
4158 is_ordered_partition(weights.len(), workers, &ranges),
4159 "weights={weights:?} workers={workers} ranges={ranges:?}"
4160 );
4161 }
4162 }
4163}
4164
4165#[cfg(all(test, feature = "parallel"))]
4166mod parallel_write_tests {
4167 use super::{
4168 direct_slice_write, frame_byte_weight, pack_frames_by_weight, pipeline_format_threads,
4169 };
4170
4171 #[test]
4174 fn two_workers_both_format_frame_groups() {
4175 assert_eq!(pipeline_format_threads(1), 1);
4176 assert_eq!(pipeline_format_threads(2), 2);
4177 assert_eq!(pipeline_format_threads(8), 8);
4178 }
4179
4180 #[test]
4183 fn serial_and_narrow_writes_skip_the_direct_slice() {
4184 let long = 3_000_000;
4185 assert!(!direct_slice_write(1, long));
4186 assert!(!direct_slice_write(2, long));
4187 assert!(!direct_slice_write(4, long));
4188 assert!(direct_slice_write(8, long));
4189 assert!(!direct_slice_write(8, long - 1));
4190 }
4191
4192 #[test]
4195 fn a_heavy_frame_is_not_glued_to_a_light_one() {
4196 assert_eq!(
4197 pack_frames_by_weight(&[1, 1, 1, 100], 2),
4198 vec![(0, 3), (3, 4)]
4199 );
4200 assert_eq!(
4201 pack_frames_by_weight(&[100, 1, 1, 1], 2),
4202 vec![(0, 1), (1, 4)]
4203 );
4204 assert_eq!(
4205 pack_frames_by_weight(&[4, 4, 4, 4, 4, 4, 4, 4], 4),
4206 vec![(0, 2), (2, 4), (4, 6), (6, 8)]
4207 );
4208 }
4209
4210 #[test]
4211 fn velocity_lines_weigh_more_than_the_same_coordinates() {
4212 let con = std::fs::read_to_string(concat!(
4213 env!("CARGO_MANIFEST_DIR"),
4214 "/resources/test/tiny_cuh2.con"
4215 ))
4216 .unwrap();
4217 let vel = std::fs::read_to_string(concat!(
4218 env!("CARGO_MANIFEST_DIR"),
4219 "/resources/test/tiny_cuh2.convel"
4220 ))
4221 .unwrap();
4222 let plain = crate::iterators::frames_from_text(&con, Some(1)).unwrap();
4223 let moving = crate::iterators::frames_from_text(&vel, Some(1)).unwrap();
4224 assert!(frame_byte_weight(&moving[0]) > frame_byte_weight(&plain[0]));
4225 assert_eq!(
4226 frame_byte_weight(&moving[0]),
4227 plain[0].atom_data.len().max(1) * 2
4228 );
4229 }
4230
4231 #[test]
4234 fn packed_groups_match_the_serial_bytes() {
4235 let con = std::fs::read_to_string(concat!(
4236 env!("CARGO_MANIFEST_DIR"),
4237 "/resources/test/tiny_cuh2.con"
4238 ))
4239 .unwrap();
4240 let vel = std::fs::read_to_string(concat!(
4241 env!("CARGO_MANIFEST_DIR"),
4242 "/resources/test/tiny_cuh2.convel"
4243 ))
4244 .unwrap();
4245 let text = format!("{con}{con}{vel}");
4246 let frames = crate::iterators::frames_from_text(&text, Some(1)).unwrap();
4247 assert!(frames.len() > 2);
4248 let refs: Vec<&crate::types::ConFrame> = frames.iter().collect();
4249 let mut serial = Vec::new();
4250 {
4251 let mut writer = super::ConFrameWriter::new(&mut serial);
4252 for frame in &frames {
4253 writer.write_frame_bound(frame).unwrap();
4254 }
4255 writer.flush().unwrap();
4256 }
4257 let path = std::env::temp_dir().join(format!("rc-pack-{}.con", std::process::id()));
4258 super::write_frame_refs_to_path(&path, &refs, 2, None, false, true).unwrap();
4259 let got = std::fs::read(&path).unwrap();
4260 assert_eq!(got, serial);
4261 let _ = std::fs::remove_file(&path);
4262 }
4263
4264 #[test]
4269 fn mixed_con_and_convel_groups_stay_within_ten_percent() {
4270 let con = std::fs::read_to_string(concat!(
4271 env!("CARGO_MANIFEST_DIR"),
4272 "/resources/test/tiny_cuh2.con"
4273 ))
4274 .unwrap();
4275 let vel = std::fs::read_to_string(concat!(
4276 env!("CARGO_MANIFEST_DIR"),
4277 "/resources/test/tiny_cuh2.convel"
4278 ))
4279 .unwrap();
4280 let text = format!("{}{}", con.repeat(24), vel.repeat(24));
4281 let frames = crate::iterators::frames_from_text(&text, Some(1)).unwrap();
4282 assert_eq!(frames.len(), 48);
4283 let refs: Vec<&crate::types::ConFrame> = frames.iter().collect();
4284 let weights: Vec<usize> = refs.iter().map(|frame| frame_byte_weight(frame)).collect();
4285 let groups = 4usize;
4286 let ranges = super::group_frame_ranges(&refs, groups);
4287 let group_weight = |ranges: &[(usize, usize)]| -> Vec<u64> {
4288 ranges
4289 .iter()
4290 .map(|&(start, end)| weights[start..end].iter().map(|w| *w as u64).sum())
4291 .collect()
4292 };
4293 let weighted = group_weight(&ranges);
4294 let total: u64 = weighted.iter().sum();
4295 let mean = total as f64 / weighted.len() as f64;
4296 let max = *weighted.iter().max().unwrap() as f64;
4297 assert!(
4298 max <= mean * 1.10,
4299 "heaviest group {max} mean {mean} weights {weighted:?}"
4300 );
4301 let equal = group_weight(&super::equal_frame_ranges(frames.len(), groups));
4302 let equal_max = *equal.iter().max().unwrap() as f64;
4303 let equal_mean = total as f64 / equal.len() as f64;
4304 assert!(
4305 equal_max > equal_mean * 1.10,
4306 "equal frame counts should miss the 10% band: {equal:?}"
4307 );
4308
4309 let mut serial = Vec::new();
4310 {
4311 let mut writer = super::ConFrameWriter::new(&mut serial);
4312 for frame in &frames {
4313 writer.write_frame_bound(frame).unwrap();
4314 }
4315 writer.flush().unwrap();
4316 }
4317 let path = std::env::temp_dir().join(format!("rc-mix-{}.con", std::process::id()));
4318 super::write_frame_refs_to_path(&path, &refs, 4, None, false, true).unwrap();
4319 assert_eq!(std::fs::read(&path).unwrap(), serial);
4320 let _ = std::fs::remove_file(&path);
4321 }
4322
4323 #[test]
4327 fn disjoint_parts_rebuild_the_concatenation() {
4328 let path = std::env::temp_dir().join(format!("rc-parts-{}.con", std::process::id()));
4329 let file = std::fs::OpenOptions::new()
4330 .read(true)
4331 .write(true)
4332 .create(true)
4333 .truncate(true)
4334 .open(&path)
4335 .unwrap();
4336 let parts = vec![
4337 b"aa".to_vec(),
4338 b"bbb".to_vec(),
4339 Vec::new(),
4340 b"cdef".to_vec(),
4341 ];
4342 super::store_disjoint_parts(&file, &parts).unwrap();
4343 drop(file);
4344 assert_eq!(std::fs::read(&path).unwrap(), b"aabbbcdef");
4345 let _ = std::fs::remove_file(&path);
4346 }
4347
4348 #[test]
4349 fn round_trip_width_matches_the_spelling() {
4350 let mut samples = vec![
4351 0.0,
4352 -0.0,
4353 1.0,
4354 -2.0,
4355 20.0,
4356 90.0,
4357 63.546,
4358 0.001,
4359 0.000001,
4360 1e-7,
4361 1e-20,
4362 1e22,
4363 -1.25,
4364 9.983,
4365 f64::MIN_POSITIVE,
4366 f64::MAX,
4367 f64::NAN,
4368 f64::INFINITY,
4369 f64::NEG_INFINITY,
4370 ];
4371 for i in 0..2000 {
4372 samples.push((i as f64) * 0.001);
4373 samples.push(-((i as f64) + 0.5) * 0.017);
4374 }
4375 let mut buf = Vec::new();
4376 let mut ryu = ryu::Buffer::new();
4377 for value in samples {
4378 buf.clear();
4379 super::push_round_trip(&mut buf, &mut ryu, None, value);
4380 assert_eq!(
4381 super::round_trip_width(value),
4382 buf.len(),
4383 "spelling {buf:?} of {value}"
4384 );
4385 }
4386 }
4387
4388 #[test]
4391 fn packed_lengths_match_formatted_frames() {
4392 use crate::types::ConFrameBuilder;
4393 let frame = |shift: f64, weird: f64| {
4394 let mut builder = ConFrameBuilder::new([20.0, 20.0, 20.0], [90.0, 90.0, 90.0]);
4395 builder
4396 .set_units(serde_json::json!({
4397 "length": "angstrom",
4398 "energy": "eV",
4399 "mass": "amu",
4400 "time": crate::units::NATURAL_TIME_UNIT,
4401 }))
4402 .unwrap();
4403 builder
4404 .add_atom("Cu", shift, shift + 0.001, weird, [false; 3], 0, 63.546)
4405 .with_velocity([shift * 0.1, -0.25, 0.0]);
4406 builder
4407 .add_atom("Cu", 1.5, -2.0, 0.0, [false; 3], 1, 63.546)
4408 .with_velocity([0.0, 1e-20, -0.0]);
4409 builder.build().unwrap()
4410 };
4411 let frames: Vec<_> = (0..6)
4412 .map(|i| {
4413 frame(
4414 i as f64 * 0.017,
4415 if i == 3 { 1.0e22 } else { 0.125 + i as f64 },
4416 )
4417 })
4418 .collect();
4419 let refs: Vec<&crate::types::ConFrame> = frames.iter().collect();
4420 let lengths = super::packed_frame_lengths(&refs, 4, None, false, true)
4421 .expect("uniform frames have exact lengths");
4422 for (frame, len) in refs.iter().zip(&lengths) {
4423 let formatted =
4424 super::format_group_bytes(std::slice::from_ref(frame), 1, None, false, true)
4425 .unwrap();
4426 assert_eq!(formatted.len(), *len);
4427 }
4428 let mut dest = vec![0u8; lengths.iter().sum()];
4429 let ptr = dest.as_mut_ptr();
4430 super::format_group_into(&refs, &mut dest, None, false, true).unwrap();
4431 assert_eq!(dest.as_ptr(), ptr);
4432 let mut serial = Vec::new();
4433 {
4434 let mut writer = super::ConFrameWriter::new(&mut serial);
4435 for frame in &frames {
4436 writer.write_frame_bound(frame).unwrap();
4437 }
4438 writer.flush().unwrap();
4439 }
4440 assert_eq!(dest, serial);
4441 let path = std::env::temp_dir().join(format!("rc-packed-{}.con", std::process::id()));
4442 super::write_frame_refs_to_path(&path, &refs, 4, None, false, true).unwrap();
4443 assert_eq!(std::fs::read(&path).unwrap(), serial);
4444 let _ = std::fs::remove_file(&path);
4445
4446 let mut other = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
4447 other
4448 .set_units(serde_json::json!({
4449 "length": "angstrom",
4450 "energy": "eV",
4451 "mass": "amu",
4452 "time": crate::units::NATURAL_TIME_UNIT,
4453 }))
4454 .unwrap();
4455 other
4456 .add_atom("H", 0.0, 0.0, 0.0, [false; 3], 0, 1.0)
4457 .with_velocity([0.0, 0.0, 0.0]);
4458 let odd = other.build().unwrap();
4459 let mixed = [refs[0], &odd];
4460 assert!(super::packed_frame_lengths(&mixed, 2, None, false, true).is_none());
4461 super::write_frame_refs_to_path(&path, &mixed, 2, None, false, true).unwrap();
4462 let mut serial = Vec::new();
4463 {
4464 let mut writer = super::ConFrameWriter::new(&mut serial);
4465 writer.write_frame_bound(mixed[0]).unwrap();
4466 writer.write_frame_bound(mixed[1]).unwrap();
4467 writer.flush().unwrap();
4468 }
4469 assert_eq!(std::fs::read(&path).unwrap(), serial);
4470 let _ = std::fs::remove_file(&path);
4471 }
4472
4473 #[test]
4476 fn a_coordinate_edit_reprices_the_packed_lengths() {
4477 use crate::types::ConFrameBuilder;
4478 let mut frames = Vec::new();
4479 for i in 0..4 {
4480 let mut builder = ConFrameBuilder::new([20.0, 20.0, 20.0], [90.0, 90.0, 90.0]);
4481 builder
4482 .set_units(serde_json::json!({
4483 "length": "angstrom",
4484 "energy": "eV",
4485 "mass": "amu",
4486 "time": crate::units::NATURAL_TIME_UNIT,
4487 }))
4488 .unwrap();
4489 builder
4490 .add_atom("Cu", i as f64, 0.25, 0.5, [false; 3], 0, 63.546)
4491 .with_velocity([0.1, 0.0, 0.0]);
4492 builder
4493 .add_atom("Cu", 1.5, -2.0, 0.0, [false; 3], 1, 63.546)
4494 .with_velocity([0.0, 0.2, 0.0]);
4495 frames.push(builder.build().unwrap());
4496 }
4497 let refs: Vec<&crate::types::ConFrame> = frames.iter().collect();
4498 let before = super::cached_packed_lengths(&refs, 2, None, false, true).unwrap();
4499 let old = before[2];
4500 frames[2].atom_data[0].set_x(1.0e22);
4501 let refs: Vec<&crate::types::ConFrame> = frames.iter().collect();
4502 let after = super::cached_packed_lengths(&refs, 2, None, false, true).unwrap();
4503 let formatted =
4504 super::format_group_bytes(std::slice::from_ref(&refs[2]), 1, None, false, true)
4505 .unwrap();
4506 assert_ne!(old, formatted.len());
4507 assert_eq!(after[2], formatted.len());
4508 }
4509}
4510
4511#[cfg(all(test, feature = "parallel"))]
4512mod parallel_write_bytes {
4513 use super::{write_frames_to_vec, write_frames_to_vec_with_threads};
4514 use crate::iterators::frames_from_text;
4515
4516 fn wide_text(n: usize) -> String {
4517 let mut text = String::new();
4518 text.push_str(
4519 "grid\n{\"con_spec_version\":2}\n10.0 10.0 10.0\n90.0 90.0 90.0\n0 0\n0 0\n1\n",
4520 );
4521 text.push_str(&format!("{n}\n63.546\nCu\nCoordinates of Component 1\n"));
4522 for i in 0..n {
4523 text.push_str(&format!("{i}.125 0.25 0.5 0 {i}\n"));
4524 }
4525 text
4526 }
4527
4528 #[test]
4529 fn parallel_bytes_match_the_serial_writer() {
4530 let wide = frames_from_text(&wide_text(5000), Some(1)).unwrap();
4531 let serial = write_frames_to_vec(&wide).unwrap();
4532 let parallel = write_frames_to_vec_with_threads(&wide, 4, None, false).unwrap();
4533 assert_eq!(parallel, serial);
4534 let mut with_vel = wide_text(5000);
4535 with_vel.push_str("\nCu\nVelocities of Component 1\n");
4536 for i in 0..5000 {
4537 with_vel.push_str(&format!("{}.0001 0.0002 -0.0001 0 {i}\n", i % 7));
4538 }
4539 let moving = frames_from_text(&with_vel, Some(1)).unwrap();
4540 let serial = write_frames_to_vec(&moving).unwrap();
4541 let parallel = write_frames_to_vec_with_threads(&moving, 4, None, false).unwrap();
4542 assert_eq!(parallel, serial);
4543 let one = frames_from_text(
4544 &std::fs::read_to_string(concat!(
4545 env!("CARGO_MANIFEST_DIR"),
4546 "/resources/test/cuh2.con"
4547 ))
4548 .unwrap()
4549 .repeat(8),
4550 Some(1),
4551 )
4552 .unwrap();
4553 let serial = write_frames_to_vec(&one).unwrap();
4554 let parallel = write_frames_to_vec_with_threads(&one, 4, None, false).unwrap();
4555 assert_eq!(parallel, serial);
4556 }
4557
4558 #[test]
4559 fn pwrite_matches_the_concatenated_writer() {
4560 let one = frames_from_text(
4561 &std::fs::read_to_string(concat!(
4562 env!("CARGO_MANIFEST_DIR"),
4563 "/resources/test/cuh2.con"
4564 ))
4565 .unwrap()
4566 .repeat(8),
4567 Some(1),
4568 )
4569 .unwrap();
4570 let serial = write_frames_to_vec(&one).unwrap();
4571 let refs: Vec<&crate::types::ConFrame> = one.iter().collect();
4572 let path = std::env::temp_dir().join(format!(
4573 "readcon-pwrite-{}-{}.con",
4574 std::process::id(),
4575 one.len()
4576 ));
4577 super::write_frame_refs_to_path(&path, &refs, 4, None, false, false).unwrap();
4578 let got = std::fs::read(&path).unwrap();
4579 let _ = std::fs::remove_file(&path);
4580 assert_eq!(got, serial);
4581 }
4582
4583 #[test]
4584 fn stamped_pwrite_matches_the_binding_writer() {
4585 let one = frames_from_text(
4586 &std::fs::read_to_string(concat!(
4587 env!("CARGO_MANIFEST_DIR"),
4588 "/resources/test/cuh2.con"
4589 ))
4590 .unwrap()
4591 .repeat(8),
4592 Some(1),
4593 )
4594 .unwrap();
4595 let mut serial = Vec::new();
4596 {
4597 let mut writer = super::ConFrameWriter::new(&mut serial);
4598 for frame in &one {
4599 let mut stamped = frame.clone();
4600 writer.write_frame_binding(&mut stamped).unwrap();
4601 }
4602 }
4603 let refs: Vec<&crate::types::ConFrame> = one.iter().collect();
4604 let path = std::env::temp_dir().join(format!(
4605 "readcon-pwrite-bind-{}-{}.con",
4606 std::process::id(),
4607 one.len()
4608 ));
4609 super::write_frame_refs_to_path(&path, &refs, 4, None, false, true).unwrap();
4610 let got = std::fs::read(&path).unwrap();
4611 let _ = std::fs::remove_file(&path);
4612 assert_eq!(got, serial);
4613 }
4614
4615 #[test]
4616 fn a_wide_file_write_matches_the_serial_bytes_when_ids_are_not_dense() {
4617 let n = 4096usize;
4618 let mut text = String::from(
4619 "grid\n{\"con_spec_version\":2}\n10.0 10.0 10.0\n90.0 90.0 90.0\n0 0\n0 0\n1\n",
4620 );
4621 text.push_str(&format!("{n}\n63.546\nCu\nCoordinates of Component 1\n"));
4622 for i in 0..n {
4623 text.push_str(&format!("{i}.125 0.25 0.5 0 {}\n", i + 10));
4624 }
4625 text.push_str("\nCu\nVelocities of Component 1\n");
4626 for i in 0..n {
4627 text.push_str(&format!("{}.0001 0.0002 -0.0001 0 {}\n", i % 7, i + 10));
4628 }
4629 let frames = frames_from_text(&text, Some(1)).unwrap();
4630 let serial = write_frames_to_vec(&frames).unwrap();
4631 let parallel = write_frames_to_vec_with_threads(&frames, 8, None, false).unwrap();
4632 assert_eq!(parallel, serial);
4633 let refs: Vec<&crate::types::ConFrame> = frames.iter().collect();
4634 let path = std::env::temp_dir().join(format!(
4635 "readcon-wide-stream-{}-{n}.con",
4636 std::process::id()
4637 ));
4638 super::write_frame_refs_to_path(&path, &refs, 8, None, false, false).unwrap();
4639 let got = std::fs::read(&path).unwrap();
4640 let _ = std::fs::remove_file(&path);
4641 assert_eq!(got, serial);
4642 }
4643
4644 #[test]
4645 fn a_mixed_wide_frame_matches_the_serial_parse() {
4646 let n = 7000usize;
4647 let mut text = String::from(
4648 "grid\n{\"con_spec_version\":2}\n10.0 10.0 10.0\n90.0 90.0 90.0\n0 0\n0 0\n2\n",
4649 );
4650 text.push_str(&format!("{n} {n}\n63.546 1.008\n"));
4651 text.push_str("Cu\nCoordinates of Component 1\n");
4652 for i in 0..n {
4653 text.push_str(&format!("{i}.125 0.25 0.5 0 {i}\n"));
4654 }
4655 text.push_str("H\nCoordinates of Component 2\n");
4656 for i in 0..n {
4657 let id = n + i;
4658 text.push_str(&format!("1.5 2.25 {i}.5 0 {id}\n"));
4659 }
4660 assert!(text.len() >= (256 << 10), "frame is {} bytes", text.len());
4661 let serial = frames_from_text(&text, Some(1)).unwrap();
4662 let parallel = frames_from_text(&text, Some(8)).unwrap();
4663 assert_eq!(parallel.len(), 1);
4664 assert_eq!(parallel[0].atom_data.len(), n * 2);
4665 assert_eq!(parallel[0].atom_data[0].symbol.as_ref(), "Cu");
4666 assert_eq!(parallel[0].atom_data[n].symbol.as_ref(), "H");
4667 assert_eq!(
4668 parallel[0].atom_data[n].atom_id,
4669 serial[0].atom_data[n].atom_id
4670 );
4671 assert_eq!(parallel, serial);
4672 }
4673
4674 #[test]
4675 fn coordinate_lines_match_the_serial_frame_body() {
4676 let frames = frames_from_text(&wide_text(32), Some(1)).unwrap();
4677 let lines = super::format_coordinate_lines(&frames[0].atom_data).unwrap();
4678 let serial = write_frames_to_vec(&frames).unwrap();
4679 let text = std::str::from_utf8(&serial).unwrap();
4680 let marker = "Coordinates of Component 1\n";
4681 let start = text.find(marker).unwrap() + marker.len();
4682 assert_eq!(lines, text.as_bytes()[start..]);
4683 }
4684
4685 #[test]
4686 fn pwrite_offsets_follow_an_exclusive_scan() {
4687 let path =
4688 std::env::temp_dir().join(format!("readcon-pwrite-scan-{}.bin", std::process::id()));
4689 let parts: [&[u8]; 4] = [b"abc", b"", b"def\n", b"g"];
4690 super::pwrite_borrowed(&path, &parts).unwrap();
4691 let got = std::fs::read(&path).unwrap();
4692 let _ = std::fs::remove_file(&path);
4693 assert_eq!(got, b"abcdef\ng");
4694 }
4695
4696 #[test]
4697 fn mapped_write_matches_the_serial_write() {
4698 use crate::types::ConFrameBuilder;
4699 for n in [1usize, 3, 7, 4097, 8191] {
4700 let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
4701 for i in 0..n {
4702 let x = (i % 97) as f64 * 0.125;
4703 let y = (i % 13) as f64;
4704 let z = if i % 2 == 0 { i as f64 * 0.001 } else { -0.5 };
4705 builder.add_atom("Cu", x, y, z, [false, i % 5 == 0, false], i as u64, 63.546);
4706 }
4707 let frame = builder.build().expect("frame");
4708 let serial = super::write_frames_to_vec(std::slice::from_ref(&frame)).expect("serial");
4709 for workers in [1usize, 2, 3, 8] {
4710 let path = std::env::temp_dir().join(format!(
4711 "readcon-map-{n}-{workers}-{}.con",
4712 std::process::id()
4713 ));
4714 let mut writer = super::ConFrameWriter::new(Vec::<u8>::new())
4715 .parallel_atoms(workers)
4716 .map_rows(true);
4717 writer.stream_file = Some(
4718 std::fs::OpenOptions::new()
4719 .read(true)
4720 .write(true)
4721 .create(true)
4722 .truncate(true)
4723 .open(&path)
4724 .expect("create"),
4725 );
4726 writer.write_frame(&frame).expect("mapped");
4727 drop(writer);
4728 let got = std::fs::read(&path).expect("read");
4729 let _ = std::fs::remove_file(&path);
4730 assert_eq!(got, serial, "n={n} workers={workers}");
4731 }
4732 }
4733 }
4734
4735 fn copper_frame(
4736 n: usize,
4737 symbol: &str,
4738 mass: f64,
4739 id_at: impl Fn(usize) -> u64,
4740 ) -> crate::types::ConFrame {
4741 use crate::types::ConFrameBuilder;
4742 let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
4743 for i in 0..n {
4744 builder.add_atom(
4745 symbol,
4746 i as f64 * 0.125,
4747 0.25,
4748 0.5,
4749 [false, false, false],
4750 id_at(i),
4751 mass,
4752 );
4753 }
4754 builder.build().expect("frame")
4755 }
4756
4757 #[test]
4758 fn two_wide_frames_on_two_threads_keep_up_with_one_frame() {
4759 let n = 200_000usize;
4760 let one = copper_frame(n, "Cu", 63.546, |i| i as u64);
4761 let frames = [one.clone(), one.clone()];
4762 let refs: Vec<&crate::types::ConFrame> = frames.iter().collect();
4763 let serial = super::write_frames_to_vec(&frames).expect("serial");
4764 let path =
4765 std::env::temp_dir().join(format!("readcon-weak-write-{}.con", std::process::id()));
4766 super::write_frame_refs_to_path(&path, &refs, 2, None, false, false).expect("parallel");
4767 let got = std::fs::read(&path).expect("read");
4768 let _ = std::fs::remove_file(&path);
4769 assert_eq!(got, serial);
4770 let time = |frames: &[crate::types::ConFrame], workers: usize| {
4774 let refs: Vec<&crate::types::ConFrame> = frames.iter().collect();
4775 let path = std::env::temp_dir().join(format!(
4776 "readcon-weak-time-{}-{workers}.con",
4777 std::process::id()
4778 ));
4779 let _ = super::write_frame_refs_to_path(&path, &refs, workers, None, false, false);
4780 let mut samples = Vec::with_capacity(5);
4781 for _ in 0..5 {
4782 let _ = std::fs::remove_file(&path);
4783 let started = std::time::Instant::now();
4784 super::write_frame_refs_to_path(&path, &refs, workers, None, false, false)
4785 .expect("write");
4786 samples.push(started.elapsed().as_secs_f64());
4787 }
4788 let _ = std::fs::remove_file(&path);
4789 samples.sort_by(|left, right| left.partial_cmp(right).unwrap());
4790 samples[samples.len() / 2]
4791 };
4792 let one_s = time(std::slice::from_ref(&one), 1);
4793 let two_s = time(&frames, 2);
4794 #[cfg(not(coverage))]
4795 assert!(
4796 two_s < one_s * 1.7,
4797 "two frames on two threads took {two_s:.4}s against one frame {one_s:.4}s"
4798 );
4799 #[cfg(coverage)]
4800 let _ = (one_s, two_s);
4801 }
4802
4803 #[test]
4804 fn parallel_write_rejects_an_id_repeated_in_another_type() {
4805 use crate::types::ConFrameBuilder;
4806 let n = 4096usize;
4807 let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
4808 for i in 0..n {
4809 builder.add_atom(
4810 "Cu",
4811 i as f64 * 0.125,
4812 0.25,
4813 0.5,
4814 [false, false, false],
4815 i as u64,
4816 63.546,
4817 );
4818 }
4819 for i in 0..n {
4820 builder.add_atom(
4821 "H",
4822 i as f64 * 0.125,
4823 1.5,
4824 0.5,
4825 [false, false, false],
4826 i as u64,
4827 1.008,
4828 );
4829 }
4830 let frame = builder.build().expect("frame");
4831 let serial = write_frames_to_vec(std::slice::from_ref(&frame));
4832 let parallel =
4833 write_frames_to_vec_with_threads(std::slice::from_ref(&frame), 4, None, false);
4834 assert!(serial.is_err(), "serial writer accepted a repeated id");
4835 assert!(parallel.is_err(), "parallel writer accepted a repeated id");
4836 let path =
4837 std::env::temp_dir().join(format!("readcon-repeat-id-{}.con", std::process::id()));
4838 std::fs::write(&path, b"keep").unwrap();
4839 let refs = [&frame];
4840 let err = super::write_frame_refs_to_path(&path, &refs, 4, None, false, false);
4841 assert!(err.is_err(), "path write accepted a repeated id");
4842 assert_eq!(std::fs::read(&path).unwrap(), b"keep");
4843 let _ = std::fs::remove_file(&path);
4844 }
4845
4846 #[test]
4847 fn a_failed_later_group_leaves_the_previous_file() {
4848 let good = copper_frame(2, "Cu", 63.546, |i| i as u64);
4849 let bad = copper_frame(2, "Cu", 63.546, |_| 0);
4850 let path =
4851 std::env::temp_dir().join(format!("readcon-group-fail-{}.con", std::process::id()));
4852 std::fs::write(&path, b"keep").unwrap();
4853 let refs = [&good, &bad];
4854 let err = super::write_frame_refs_to_path(&path, &refs, 4, None, false, false);
4855 assert!(err.is_err());
4856 assert_eq!(std::fs::read(&path).unwrap(), b"keep");
4857 let _ = std::fs::remove_file(&path);
4858 }
4859
4860 #[test]
4861 fn mapped_decimal_lines_match_the_serial_writer() {
4862 let frame = copper_frame(64, "Cu", 63.546, |i| i as u64);
4863 let frames = [frame];
4864 let serial = write_frames_to_vec_with_threads(&frames, 1, Some(40), false).unwrap();
4865 let path =
4866 std::env::temp_dir().join(format!("readcon-long-line-{}.con", std::process::id()));
4867 let mut writer = super::ConFrameWriter::with_precision(Vec::<u8>::new(), 40)
4868 .parallel_atoms(4)
4869 .map_rows(true);
4870 writer.stream_file = Some(
4871 std::fs::OpenOptions::new()
4872 .read(true)
4873 .write(true)
4874 .create(true)
4875 .truncate(true)
4876 .open(&path)
4877 .unwrap(),
4878 );
4879 writer.write_frame(&frames[0]).unwrap();
4880 drop(writer);
4881 let got = std::fs::read(&path).unwrap();
4882 let _ = std::fs::remove_file(&path);
4883 assert_eq!(got, serial);
4884 }
4885
4886 #[test]
4887 fn a_rejected_spread_does_not_keep_the_reserved_file() {
4888 use crate::types::ConFrameBuilder;
4889 let n = 4096usize;
4890 let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
4891 for i in 0..n {
4892 builder.add_atom(
4893 "Cu",
4894 i as f64 * 0.125,
4895 0.25,
4896 0.5,
4897 [false, false, false],
4898 i as u64,
4899 63.546,
4900 );
4901 }
4902 builder.with_spread([0.0, 0.0, 0.0]);
4903 let mut frame = builder.build().expect("frame");
4904 frame.atom_data[0].set_spread(Some([-1.0, 0.0, 0.0]));
4905 let path =
4906 std::env::temp_dir().join(format!("readcon-spread-fail-{}.con", std::process::id()));
4907 let file = std::fs::OpenOptions::new()
4908 .read(true)
4909 .write(true)
4910 .create(true)
4911 .truncate(true)
4912 .open(&path)
4913 .unwrap();
4914 let mut writer = super::ConFrameWriter::new(Vec::<u8>::new())
4915 .parallel_atoms(4)
4916 .map_rows(true)
4917 .mapped_file(file);
4918 let err = writer.write_frame(&frame);
4919 drop(writer);
4920 assert!(err.is_err());
4921 assert_eq!(std::fs::metadata(&path).unwrap().len(), 0);
4922 let _ = std::fs::remove_file(&path);
4923 }
4924
4925 #[test]
4926 fn overlapped_pwrite_does_not_keep_bytes_past_the_frame() {
4927 use std::io::{Seek, Write};
4928 let frame = copper_frame(4096, "Cu", 63.546, |i| i as u64);
4929 let serial = write_frames_to_vec(std::slice::from_ref(&frame)).unwrap();
4930 let path = std::env::current_dir()
4931 .unwrap()
4932 .join(format!("readcon-pwrite-tail-{}.con", std::process::id()));
4933 let mut file = std::fs::OpenOptions::new()
4934 .read(true)
4935 .write(true)
4936 .create(true)
4937 .truncate(true)
4938 .open(&path)
4939 .unwrap();
4940 file.write_all(&vec![b'Z'; serial.len() + 64]).unwrap();
4941 file.seek(std::io::SeekFrom::Start(0)).unwrap();
4942 let workers = super::physical_cores().max(2);
4943 let mut writer = super::ConFrameWriter::new(Vec::<u8>::new())
4944 .parallel_atoms(workers)
4945 .mapped_file(file);
4946 writer.write_frame(&frame).unwrap();
4947 drop(writer);
4948 let got = std::fs::read(&path).unwrap();
4949 let _ = std::fs::remove_file(&path);
4950 assert_eq!(got, serial);
4951 }
4952}
4953
4954#[cfg(test)]
4955mod gzip_stream_memory {
4956 use super::ConFrameWriter;
4957 use crate::iterators::{frames_from_text, read_all_frames};
4958
4959 fn parsed_frame() -> crate::types::ConFrame {
4960 let mut text = String::from(
4961 "grid\n{\"con_spec_version\":2,\"units\":{\"length\":\"angstrom\",\"energy\":\"eV\"}}\n20 20 20\n90 90 90\n0 0\n0 0\n1\n48\n63.546\nCu\nCoordinates of Component 1\n",
4962 );
4963 for i in 0..48 {
4964 text.push_str(&format!("0.{i} 1 2 0 {i}\n"));
4965 }
4966 frames_from_text(&text, Some(1)).unwrap().remove(0)
4967 }
4968
4969 #[test]
4978 fn parsed_frames_gzip_write_drops_the_metadata_maps() {
4979 let frames = vec![parsed_frame(); 4000];
4980 assert!(
4981 frames
4982 .iter()
4983 .all(|frame| !frame.header.metadata.has_cached_map())
4984 );
4985 let path = std::env::temp_dir().join(format!(
4986 "readcon-gzip-peak-{}-{}.con.gz",
4987 std::process::id(),
4988 std::time::SystemTime::now()
4989 .duration_since(std::time::UNIX_EPOCH)
4990 .unwrap()
4991 .as_nanos()
4992 ));
4993 {
4994 let mut writer = ConFrameWriter::from_path_gzip(&path).unwrap();
4995 for frame in &frames {
4996 writer.write_frame(frame).unwrap();
4997 }
4998 writer.flush().unwrap();
4999 }
5000 let cached = frames
5001 .iter()
5002 .filter(|frame| frame.header.metadata.has_cached_map())
5003 .count();
5004 assert_eq!(
5005 cached, 0,
5006 "the write retained {cached} metadata maps; each is about a kilobyte"
5007 );
5008 let back = read_all_frames(&path).unwrap();
5009 assert_eq!(back.len(), frames.len());
5010 assert_eq!(back[0].atom_data.len(), 48);
5011 assert_eq!(back[0].atom_data[3].atom_id, 3);
5012 assert_eq!(back[17].atom_data[0].symbol.as_ref(), "Cu");
5013 let _ = std::fs::remove_file(&path);
5014 }
5015
5016 #[test]
5019 fn lattice_line_overrides_box_lengths_and_can_be_reread_from_the_cache() {
5020 let text = "\
5021grid
5022{\"con_spec_version\":2,\"lattice_vectors\":[[10,0,0],[0,20,0],[0,0,30]]}
50235 5 5
502490 90 90
50250 0
50260 0
50271
50281
50291.0
5030Cu
5031Coordinates of Component 1
50320 0 0 0 1
5033";
5034 let frame = frames_from_text(text, Some(1)).unwrap().remove(0);
5035 assert!(!frame.header.metadata.has_cached_map());
5036 let vecs = frame.header.lattice_vectors().unwrap();
5037 assert_eq!(vecs, [[10.0, 0.0, 0.0], [0.0, 20.0, 0.0], [0.0, 0.0, 30.0]]);
5038 assert!(
5039 !frame.header.metadata.has_cached_map(),
5040 "reading lattice vectors cached the whole metadata map"
5041 );
5042 let path =
5043 std::env::temp_dir().join(format!("readcon-lattice-line-{}.con", std::process::id()));
5044 {
5045 let mut writer = ConFrameWriter::from_path(&path).unwrap();
5046 writer.write_frame(&frame).unwrap();
5047 writer.flush().unwrap();
5048 }
5049 let back = read_all_frames(&path).unwrap();
5050 let _ = std::fs::remove_file(&path);
5051 assert_eq!(back[0].header.boxl, [10.0, 20.0, 30.0]);
5052 assert!(frame.header.metadata.get("lattice_vectors").is_some());
5053 assert!(frame.header.metadata.has_cached_map());
5054 assert_eq!(frame.header.lattice_vectors().unwrap()[1][1], 20.0);
5055 }
5056
5057 #[test]
5060 fn repeated_lattice_lookups_share_a_bounded_cache() {
5061 let mut frames = Vec::new();
5062 for side in [10.0, 11.0, 12.0, 13.0, 14.0] {
5063 let text = format!(
5064 "grid\n{{\"con_spec_version\":2,\"lattice_vectors\":[[{side},0,0],[0,{side},0],[0,0,{side}]]}}\n5 5 5\n90 90 90\n0 0\n0 0\n1\n1\n1.0\nCu\nCoordinates of Component 1\n0 0 0 0 1\n"
5065 );
5066 frames.push(frames_from_text(&text, Some(1)).unwrap().remove(0));
5067 }
5068 for _ in 0..2 {
5069 for (index, frame) in frames.iter().enumerate() {
5070 let side = 10.0 + index as f64;
5071 let vecs = frame.header.lattice_vectors().unwrap();
5072 assert_eq!(vecs[0][0], side);
5073 assert_eq!(vecs[2][2], side);
5074 assert!(!frame.header.metadata.has_cached_map());
5075 }
5076 }
5077 let older = frames[1].header.lattice_vectors().unwrap();
5081 assert_eq!(older[0][0], 11.0);
5082 let text = "\
5083grid
5084{\"con_spec_version\":2,\"lattice_vectors\":[[14,0,0],[0,14,0],[0,0,14]]}
50855 5 5
508690 90 90
50870 0
50880 0
50891
50901
50911.0
5092Cu
5093Coordinates of Component 1
50940 0 0 0 1
5095";
5096 let twin = frames_from_text(text, Some(1)).unwrap().remove(0);
5097 assert_eq!(twin.header.lattice_vectors().unwrap()[0][0], 14.0);
5098 assert!(!twin.header.metadata.has_cached_map());
5099 }
5100
5101 #[test]
5104 fn lattice_lookups_from_several_threads_keep_their_own_lines() {
5105 let frames: Vec<_> = (1..17)
5106 .map(|side| {
5107 let text = format!(
5108 "grid\n{{\"con_spec_version\":2,\"lattice_vectors\":[[{side},0,0],[0,{side},0],[0,0,{side}]]}}\n5 5 5\n90 90 90\n0 0\n0 0\n1\n1\n1.0\nCu\nCoordinates of Component 1\n0 0 0 0 1\n"
5109 );
5110 frames_from_text(&text, Some(1)).unwrap().remove(0)
5111 })
5112 .collect();
5113 std::thread::scope(|scope| {
5114 for _ in 0..4 {
5115 scope.spawn(|| {
5116 for _ in 0..8 {
5117 for (index, frame) in frames.iter().enumerate() {
5118 let side = 1.0 + index as f64;
5119 let vecs = frame.header.lattice_vectors().unwrap();
5120 assert_eq!(vecs[0][0], side);
5121 assert_eq!(vecs[1][1], side);
5122 assert_eq!(vecs[2][2], side);
5123 }
5124 }
5125 });
5126 }
5127 });
5128 }
5129}