Skip to main content

readcon_core/
writer.rs

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::json;
7use std::fs::File;
8use std::io::{self, BufWriter, Write};
9use std::path::Path;
10
11/// Failure from [`ConFrameWriter`]. I/O stays distinct from format errors
12/// so `InvalidUnits` and `InvalidCount` remain distinguishable (CON-error.1).
13#[derive(Debug)]
14pub enum WriteError {
15    Io(io::Error),
16    Format(crate::error::ParseError),
17}
18
19impl std::fmt::Display for WriteError {
20    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        match self {
22            WriteError::Io(e) => write!(f, "{e}"),
23            WriteError::Format(e) => write!(f, "{e}"),
24        }
25    }
26}
27
28impl std::error::Error for WriteError {
29    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
30        match self {
31            WriteError::Io(e) => Some(e),
32            WriteError::Format(e) => Some(e),
33        }
34    }
35}
36
37impl From<io::Error> for WriteError {
38    fn from(e: io::Error) -> Self {
39        WriteError::Io(e)
40    }
41}
42
43impl From<WriteError> for io::Error {
44    fn from(e: WriteError) -> Self {
45        match e {
46            WriteError::Io(e) => e,
47            WriteError::Format(err) => io::Error::new(io::ErrorKind::InvalidData, err.to_string()),
48        }
49    }
50}
51
52/// Numeric text representation for CON frame fields.
53#[derive(Clone, Copy, Debug, PartialEq, Eq)]
54pub enum FloatFormat {
55    /// A fixed number of digits after the decimal point.
56    DecimalPlaces(usize),
57    /// The shortest scientific representation preserving each finite f64.
58    /// Signed zero is preserved; decimal precision does not limit this mode.
59    RoundTrip,
60}
61
62impl Default for FloatFormat {
63    fn default() -> Self {
64        // CON-num.6 (D13): the default writer round-trips binary64.
65        Self::RoundTrip
66    }
67}
68
69/// A writer that can serialize and write `ConFrame` objects to any output stream.
70///
71/// This struct encapsulates a writer (like a file) and provides a high-level API
72/// for writing simulation frames in the `.con` format.
73///
74/// # Example
75/// ```no_run
76/// # use std::fs::File;
77/// # use readcon_core::types::ConFrame;
78/// # use readcon_core::writer::ConFrameWriter;
79/// # let frames: Vec<ConFrame> = Vec::new();
80/// let mut writer = ConFrameWriter::from_path("output.con").unwrap();
81/// writer.extend(frames.iter()).unwrap();
82/// ```
83pub struct ConFrameWriter<W: Write> {
84    writer: BufWriter<W>,
85    float_format: FloatFormat,
86    /// When true: sort metadata keys in JSON and emit sections in canonical
87    /// order (velocities, forces, energies). Opt-in. Float text follows
88    /// `float_format` (round-trip unless the caller asked for decimals).
89    canonical: bool,
90    /// Cache for the JSON metadata line: when consecutive frames share
91    /// the same (spec_version, sections-set, metadata) triple the
92    /// serialized JSON object is identical, so reusing the cached
93    /// string skips the per-frame `serde_json::Map::insert` rebuild
94    /// and re-serialisation. Hot for trajectory writes where every
95    /// frame has the same `units` / `potential` / `validate` keys.
96    metadata_cache: Option<MetadataCacheEntry>,
97    /// One frame of CON text. Filled then flushed with a single `write_all`
98    /// so the sink is not entered once per atom line.
99    scratch: Vec<u8>,
100}
101
102#[derive(Debug)]
103struct MetadataCacheEntry {
104    /// Snapshot of the inputs that fully determine the serialized JSON
105    /// metadata line. Cheap to clone; cheaper than re-serialising the
106    /// whole map on every frame.
107    spec_version: u32,
108    has_velocities: bool,
109    has_forces: bool,
110    has_energies: bool,
111    has_charges: bool,
112    has_spins: bool,
113    has_magmoms: bool,
114    has_displacements: bool,
115    has_spreads: bool,
116    metadata: std::collections::BTreeMap<String, serde_json::Value>,
117    /// Cached serialised metadata line (without trailing newline).
118    serialized: String,
119}
120
121impl MetadataCacheEntry {
122    fn matches(
123        &self,
124        spec_version: u32,
125        has_velocities: bool,
126        has_forces: bool,
127        has_energies: bool,
128        has_charges: bool,
129        has_spins: bool,
130        has_magmoms: bool,
131        has_displacements: bool,
132        has_spreads: bool,
133        metadata: &std::collections::BTreeMap<String, serde_json::Value>,
134    ) -> bool {
135        self.spec_version == spec_version
136            && self.has_velocities == has_velocities
137            && self.has_forces == has_forces
138            && self.has_energies == has_energies
139            && self.has_charges == has_charges
140            && self.has_spins == has_spins
141            && self.has_magmoms == has_magmoms
142            && self.has_displacements == has_displacements
143            && self.has_spreads == has_spreads
144            && self.metadata.len() == metadata.len()
145            && self
146                .metadata
147                .iter()
148                .zip(metadata)
149                .all(|((ka, a), (kb, b))| ka == kb && metadata_value_eq(a, b))
150    }
151}
152
153// Cache equality preserves the numeric representation, including signed zero,
154// throughout nested JSON values.
155fn metadata_value_eq(a: &serde_json::Value, b: &serde_json::Value) -> bool {
156    use serde_json::Value;
157    match (a, b) {
158        (Value::Number(a), Value::Number(b)) if a.is_f64() && b.is_f64() => {
159            a.as_f64().map(f64::to_bits) == b.as_f64().map(f64::to_bits)
160        }
161        (Value::Array(a), Value::Array(b)) => {
162            a.len() == b.len() && a.iter().zip(b).all(|(a, b)| metadata_value_eq(a, b))
163        }
164        (Value::Object(a), Value::Object(b)) => {
165            a.len() == b.len()
166                && a.iter()
167                    .all(|(key, a)| b.get(key).is_some_and(|b| metadata_value_eq(a, b)))
168        }
169        _ => a == b,
170    }
171}
172
173/// CON-section.16. Minor 2 and later must list, in `section_arity`, every
174/// emitted section that minor 1 does not reserve.
175fn require_minor2_section_arity(
176    minor: u64,
177    sections: &[serde_json::Value],
178    metadata: &serde_json::Map<String, serde_json::Value>,
179) -> Result<(), WriteError> {
180    if minor < 2 {
181        return Ok(());
182    }
183    for name in sections.iter().filter_map(|value| value.as_str()) {
184        if crate::types::reserved_section_value_arity(name).is_some() {
185            continue;
186        }
187        let listed = metadata
188            .get(meta::SECTION_ARITY)
189            .and_then(|value| value.as_object())
190            .and_then(|obj| obj.get(name))
191            .and_then(|value| value.as_u64())
192            .is_some_and(|n| n >= 1);
193        if !listed {
194            return Err(WriteError::Format(
195                crate::error::ParseError::ValidationError(format!(
196                    "minor >= 2 must list {name} in section_arity"
197                )),
198            ));
199        }
200    }
201    Ok(())
202}
203
204// General implementation for any type that implements `Write`.
205impl<W: Write> ConFrameWriter<W> {
206    /// Creates a new `ConFrameWriter` that wraps a given writer.
207    ///
208    /// # Arguments
209    ///
210    /// * `writer` - Any type that implements `std::io::Write`, e.g., a `File`.
211    pub fn new(writer: W) -> Self {
212        Self::with_float_format(writer, FloatFormat::default())
213    }
214
215    /// Creates a new `ConFrameWriter` with a custom floating-point precision.
216    ///
217    /// # Arguments
218    ///
219    /// * `writer` - Any type that implements `std::io::Write`.
220    /// * `precision` - Number of decimal places for floating-point output.
221    pub fn with_precision(writer: W, precision: usize) -> Self {
222        Self::with_float_format(writer, FloatFormat::DecimalPlaces(precision))
223    }
224
225    /// Creates a writer with an explicit numeric text representation.
226    ///
227    /// `FloatFormat::RoundTrip` retains small forces and coordinates exactly.
228    /// Metadata floats use the JSON serializer independently of this setting.
229    pub fn with_float_format(writer: W, float_format: FloatFormat) -> Self {
230        Self {
231            writer: BufWriter::new(writer),
232            float_format,
233            canonical: false,
234            metadata_cache: None,
235            scratch: Vec::with_capacity(16 * 1024),
236        }
237    }
238
239    /// Opt-in **canonical** serialization: BTree-ordered metadata keys in JSON,
240    /// fixed section order. Use for corpus
241    /// materialization and content-stable hashes; not required for on-disk fidelity
242    /// of spans preserved from `next_with_raw_span`.
243    pub fn canonical(mut self, on: bool) -> Self {
244        self.set_canonical(on);
245        self
246    }
247
248    /// Set or clear canonical mode on an existing writer (C ABI / FFI).
249    pub fn set_canonical(&mut self, on: bool) {
250        self.canonical = on;
251        if on {
252            self.metadata_cache = None;
253        }
254    }
255
256    /// Whether canonical serialization is enabled.
257    pub fn is_canonical(&self) -> bool {
258        self.canonical
259    }
260
261    fn refresh_metadata_cache(&mut self, frame: &ConFrame) -> Result<(), WriteError> {
262        let spec_version = frame.header.spec_version;
263        let has_vel = frame.has_velocities();
264        let has_frc = frame.has_forces();
265        let has_eng = frame.has_energies();
266        let has_chg = frame.has_charges();
267        let has_spn = frame.has_spins();
268        let has_mm = frame.has_magmoms();
269        let has_dsp = frame.has_displacements();
270        let has_spr = frame.has_spreads();
271
272        let cache_hit = !self.canonical
273            && self.metadata_cache.as_ref().is_some_and(|c| {
274                c.matches(
275                    spec_version,
276                    has_vel,
277                    has_frc,
278                    has_eng,
279                    has_chg,
280                    has_spn,
281                    has_mm,
282                    has_dsp,
283                    has_spr,
284                    &frame.header.metadata,
285                )
286            });
287        if cache_hit {
288            return Ok(());
289        }
290
291        let mut meta_obj = serde_json::Map::new();
292        meta_obj.insert(meta::CON_SPEC_VERSION.into(), json!(spec_version));
293        let mut sections = Vec::new();
294        if has_vel {
295            sections.push(json!(SECTION_VELOCITIES));
296        }
297        if has_frc {
298            sections.push(json!(SECTION_FORCES));
299        }
300        if has_eng {
301            sections.push(json!(SECTION_ENERGIES));
302        }
303        if has_chg {
304            sections.push(json!(SECTION_CHARGES));
305        }
306        if has_spn {
307            sections.push(json!(SECTION_SPINS));
308        }
309        if has_mm {
310            sections.push(json!(SECTION_MAGMOMS));
311        }
312        if has_dsp {
313            sections.push(json!(SECTION_DISPLACEMENTS));
314        }
315        if has_spr {
316            sections.push(json!(SECTION_SPREADS));
317        }
318        let validate = frame
319            .header
320            .metadata
321            .get(meta::VALIDATE)
322            .and_then(|value| value.as_bool())
323            .unwrap_or(false);
324        if !sections.is_empty() || validate {
325            meta_obj.insert(meta::SECTIONS.into(), json!(sections));
326        }
327        for (k, v) in &frame.header.metadata {
328            if k == meta::CON_SPEC_VERSION || k == meta::SECTIONS {
329                continue;
330            }
331            meta_obj.insert(k.clone(), v.clone());
332        }
333        if let Some(u) = meta_obj.get(meta::UNITS).cloned() {
334            // D11: an invalid unit fails the write. Do not substitute a default.
335            let c = crate::units::canonicalize_units_object(&u).map_err(|e| {
336                WriteError::Format(crate::error::ParseError::InvalidUnits(e.to_string()))
337            })?;
338            meta_obj.insert(meta::UNITS.into(), c);
339        }
340        if spec_version >= 3 {
341            match meta_obj.get(meta::UNITS) {
342                None => {
343                    return Err(WriteError::Format(crate::error::ParseError::InvalidUnits(
344                        "con_spec_version >= 3 requires units.length and units.energy".into(),
345                    )));
346                }
347                Some(u) => {
348                    crate::units::validate_v3_units_metadata(u).map_err(|e| {
349                        WriteError::Format(crate::error::ParseError::InvalidUnits(e.to_string()))
350                    })?;
351                }
352            }
353        }
354        let minor = frame
355            .header
356            .metadata
357            .get(meta::CON_SPEC_MINOR)
358            .and_then(|v| v.as_u64())
359            .unwrap_or(0);
360        if spec_version >= 3 && minor >= 1 && has_vel {
361            let time_ok = meta_obj
362                .get(meta::UNITS)
363                .and_then(|u| u.get("time"))
364                .and_then(|v| v.as_str())
365                .is_some_and(|s| !s.trim().is_empty());
366            if !time_ok {
367                return Err(WriteError::Format(crate::error::ParseError::InvalidUnits(
368                    "minor >= 1 requires units.time when a velocities section is present".into(),
369                )));
370            }
371        }
372        // A minor of 2 or later lists every section minor 1 does not reserve
373        // (CON-section.16). Reserved names need no entry.
374        require_minor2_section_arity(minor, &sections, &meta_obj)?;
375        let serialized = serde_json::Value::Object(meta_obj).to_string();
376        self.metadata_cache = Some(MetadataCacheEntry {
377            spec_version,
378            has_velocities: has_vel,
379            has_forces: has_frc,
380            has_energies: has_eng,
381            has_charges: has_chg,
382            has_spins: has_spn,
383            has_magmoms: has_mm,
384            has_displacements: has_dsp,
385            has_spreads: has_spr,
386            metadata: frame.header.metadata.clone(),
387            serialized,
388        });
389        Ok(())
390    }
391
392    /// Writes a single `ConFrame` to the output stream.
393    pub fn write_frame(&mut self, frame: &ConFrame) -> Result<(), WriteError> {
394        let zero_frame = frame.header.natm_types == 0
395            && frame.atom_data.is_empty()
396            && frame.header.natms_per_type.is_empty()
397            && frame.header.masses_per_type.is_empty();
398        if zero_frame && !frame.header.sections.is_empty() {
399            return Err(WriteError::Format(crate::error::ParseError::InvalidCount(
400                "a zero-atom frame cannot declare sections".into(),
401            )));
402        }
403        if !zero_frame
404            && (frame.header.natm_types == 0
405                || frame.atom_data.is_empty()
406                || frame.header.natms_per_type.iter().any(|n| *n == 0))
407        {
408            return Err(WriteError::Format(crate::error::ParseError::InvalidCount(
409                "each per-type count must be >= 1".into(),
410            )));
411        }
412        // Writers emit ids that are unique within the frame (CON-id.3, D9).
413        // Readers accept repeats (CON-id.4) and may keep file order (CON-id.8).
414        let mut seen_ids = std::collections::HashSet::with_capacity(frame.atom_data.len());
415        for atom in &frame.atom_data {
416            if !seen_ids.insert(atom.atom_id) {
417                return Err(WriteError::Format(
418                    crate::error::ParseError::ValidationError(format!(
419                        "atom id {} is repeated in the frame",
420                        atom.atom_id
421                    )),
422                ));
423            }
424        }
425        let prec = self.float_format;
426        let meta_line = if frame.header.spec_version < 2 {
427            // Spec 1 line 2 is free text. Do not re-emit JSON major 0 or 1 (D15).
428            frame.header.prebox_header.metadata_line().to_string()
429        } else {
430            self.refresh_metadata_cache(frame)?;
431            self.metadata_cache
432                .as_ref()
433                .expect("metadata_cache populated above")
434                .serialized
435                .clone()
436        };
437        let (boxl, angles) = match frame.header.lattice_vectors() {
438            Some(vecs) => crate::types::lengths_and_angles(vecs),
439            None => (frame.header.boxl, frame.header.angles),
440        };
441        self.scratch.clear();
442        {
443            let buf = &mut self.scratch;
444
445            // --- Write the 9-line Header ---
446            let _ = writeln!(buf, "{}", frame.header.prebox_header.user);
447            let _ = writeln!(buf, "{meta_line}");
448            push_f64(buf, boxl[0], prec);
449            buf.push(b' ');
450            push_f64(buf, boxl[1], prec);
451            buf.push(b' ');
452            push_f64(buf, boxl[2], prec);
453            buf.push(b'\n');
454            push_f64(buf, angles[0], prec);
455            buf.push(b' ');
456            push_f64(buf, angles[1], prec);
457            buf.push(b' ');
458            push_f64(buf, angles[2], prec);
459            buf.push(b'\n');
460            let _ = writeln!(buf, "{}", frame.header.postbox_header[0]);
461            let _ = writeln!(buf, "{}", frame.header.postbox_header[1]);
462            let _ = writeln!(buf, "{}", frame.header.natm_types);
463
464            for (i, n) in frame.header.natms_per_type.iter().enumerate() {
465                if i > 0 {
466                    buf.push(b' ');
467                }
468                push_u64(buf, *n as u64);
469            }
470            buf.push(b'\n');
471
472            for (i, m) in frame.header.masses_per_type.iter().enumerate() {
473                if i > 0 {
474                    buf.push(b' ');
475                }
476                push_f64(buf, *m, prec);
477            }
478            buf.push(b'\n');
479
480            // --- Write the Atom Data ---
481            let mut atom_idx_offset = 0;
482            for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate() {
483                let symbol = &frame.atom_data[atom_idx_offset].symbol;
484                let _ = writeln!(buf, "{symbol}");
485                let _ = writeln!(buf, "Coordinates of Component {}", type_idx + 1);
486
487                for i in 0..num_atoms_in_type {
488                    let atom = &frame.atom_data[atom_idx_offset + i];
489                    push_xyz_line(
490                        buf,
491                        atom.x,
492                        atom.y,
493                        atom.z,
494                        prec,
495                        written_fixed_flag(atom.fixed, frame.header.spec_version),
496                        atom.atom_id,
497                    );
498                }
499                atom_idx_offset += num_atoms_in_type;
500            }
501
502            // --- Write optional velocity section ---
503            if frame.has_velocities() {
504                buf.push(b'\n');
505
506                let mut vel_idx_offset = 0;
507                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
508                {
509                    let symbol = &frame.atom_data[vel_idx_offset].symbol;
510                    let _ = writeln!(buf, "{symbol}");
511                    let _ = writeln!(buf, "Velocities of Component {}", type_idx + 1);
512
513                    for i in 0..num_atoms_in_type {
514                        let atom = &frame.atom_data[vel_idx_offset + i];
515                        let [vx, vy, vz] = atom.velocity.unwrap_or([0.0; 3]);
516                        push_xyz_line(
517                            buf,
518                            vx,
519                            vy,
520                            vz,
521                            prec,
522                            written_fixed_flag(atom.fixed, frame.header.spec_version),
523                            atom.atom_id,
524                        );
525                    }
526                    vel_idx_offset += num_atoms_in_type;
527                }
528            }
529
530            // --- Write optional force section ---
531            if frame.has_forces() {
532                buf.push(b'\n');
533
534                let mut force_idx_offset = 0;
535                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
536                {
537                    let symbol = &frame.atom_data[force_idx_offset].symbol;
538                    let _ = writeln!(buf, "{symbol}");
539                    let _ = writeln!(buf, "Forces of Component {}", type_idx + 1);
540
541                    for i in 0..num_atoms_in_type {
542                        let atom = &frame.atom_data[force_idx_offset + i];
543                        let [fx, fy, fz] = atom.force.unwrap_or([0.0; 3]);
544                        push_xyz_line(
545                            buf,
546                            fx,
547                            fy,
548                            fz,
549                            prec,
550                            written_fixed_flag(atom.fixed, frame.header.spec_version),
551                            atom.atom_id,
552                        );
553                    }
554                    force_idx_offset += num_atoms_in_type;
555                }
556            }
557
558            // --- Write optional energies section ---
559            if frame.has_energies() {
560                buf.push(b'\n');
561
562                let mut energy_idx_offset = 0;
563                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
564                {
565                    let symbol = &frame.atom_data[energy_idx_offset].symbol;
566                    let _ = writeln!(buf, "{symbol}");
567                    let _ = writeln!(buf, "Energies of Component {}", type_idx + 1);
568
569                    for i in 0..num_atoms_in_type {
570                        let atom = &frame.atom_data[energy_idx_offset + i];
571                        let e = atom.energy.unwrap_or(0.0);
572                        push_scalar_line(
573                            buf,
574                            e,
575                            prec,
576                            written_fixed_flag(atom.fixed, frame.header.spec_version),
577                            atom.atom_id,
578                        );
579                    }
580                    energy_idx_offset += num_atoms_in_type;
581                }
582            }
583
584            if frame.has_charges() {
585                buf.push(b'\n');
586                let mut off = 0;
587                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
588                {
589                    let symbol = &frame.atom_data[off].symbol;
590                    let _ = writeln!(buf, "{symbol}");
591                    let _ = writeln!(buf, "Charges of Component {}", type_idx + 1);
592                    for i in 0..num_atoms_in_type {
593                        let atom = &frame.atom_data[off + i];
594                        let q = atom.charge.unwrap_or(0.0);
595                        push_scalar_line(
596                            buf,
597                            q,
598                            prec,
599                            written_fixed_flag(atom.fixed, frame.header.spec_version),
600                            atom.atom_id,
601                        );
602                    }
603                    off += num_atoms_in_type;
604                }
605            }
606
607            if frame.has_spins() {
608                buf.push(b'\n');
609                let mut off = 0;
610                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
611                {
612                    let symbol = &frame.atom_data[off].symbol;
613                    let _ = writeln!(buf, "{symbol}");
614                    let _ = writeln!(buf, "Spins of Component {}", type_idx + 1);
615                    for i in 0..num_atoms_in_type {
616                        let atom = &frame.atom_data[off + i];
617                        let s = atom.spin.unwrap_or(0.0);
618                        push_scalar_line(
619                            buf,
620                            s,
621                            prec,
622                            written_fixed_flag(atom.fixed, frame.header.spec_version),
623                            atom.atom_id,
624                        );
625                    }
626                    off += num_atoms_in_type;
627                }
628            }
629
630            if frame.has_magmoms() {
631                buf.push(b'\n');
632                let mut off = 0;
633                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
634                {
635                    let symbol = &frame.atom_data[off].symbol;
636                    let _ = writeln!(buf, "{symbol}");
637                    let _ = writeln!(buf, "Magmoms of Component {}", type_idx + 1);
638                    for i in 0..num_atoms_in_type {
639                        let atom = &frame.atom_data[off + i];
640                        let [mx, my, mz] = atom.magmom.unwrap_or([0.0; 3]);
641                        push_xyz_line(
642                            buf,
643                            mx,
644                            my,
645                            mz,
646                            prec,
647                            written_fixed_flag(atom.fixed, frame.header.spec_version),
648                            atom.atom_id,
649                        );
650                    }
651                    off += num_atoms_in_type;
652                }
653            }
654
655            if frame.has_displacements() {
656                buf.push(b'\n');
657                let mut off = 0;
658                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
659                {
660                    let symbol = &frame.atom_data[off].symbol;
661                    let _ = writeln!(buf, "{symbol}");
662                    let _ = writeln!(buf, "Displacements of Component {}", type_idx + 1);
663                    for i in 0..num_atoms_in_type {
664                        let atom = &frame.atom_data[off + i];
665                        let [dx, dy, dz] = atom.displacement.unwrap_or([0.0; 3]);
666                        push_xyz_line(
667                            buf,
668                            dx,
669                            dy,
670                            dz,
671                            prec,
672                            written_fixed_flag(atom.fixed, frame.header.spec_version),
673                            atom.atom_id,
674                        );
675                    }
676                    off += num_atoms_in_type;
677                }
678            }
679
680            if frame.has_spreads() {
681                buf.push(b'\n');
682                let mut off = 0;
683                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
684                {
685                    let symbol = &frame.atom_data[off].symbol;
686                    let _ = writeln!(buf, "{symbol}");
687                    let _ = writeln!(buf, "Spreads of Component {}", type_idx + 1);
688                    for i in 0..num_atoms_in_type {
689                        let atom = &frame.atom_data[off + i];
690                        let [sx, sy, sz] = atom.spread.unwrap_or([0.0; 3]);
691                        if !crate::types::spread_row_ok([sx, sy, sz]) {
692                            let err = if sx.is_finite() && sy.is_finite() && sz.is_finite() {
693                                crate::types::invalid_spread(off + i)
694                            } else {
695                                crate::error::ParseError::NonFiniteValue(format!(
696                                    "spreads atom {}",
697                                    off + i
698                                ))
699                            };
700                            return Err(WriteError::Format(err));
701                        }
702                        push_xyz_line(
703                            buf,
704                            sx,
705                            sy,
706                            sz,
707                            prec,
708                            written_fixed_flag(atom.fixed, frame.header.spec_version),
709                            atom.atom_id,
710                        );
711                    }
712                    off += num_atoms_in_type;
713                }
714            }
715        }
716
717        self.writer.write_all(&self.scratch)?;
718        Ok(())
719    }
720
721    /// Flushes buffered data to the underlying writer and reports errors.
722    pub fn flush(&mut self) -> io::Result<()> {
723        self.writer.flush()
724    }
725
726    /// Writes all frames from an iterator to the output stream.
727    ///
728    /// This is the most convenient way to write a multi-frame file.
729    pub fn extend<'a>(
730        &mut self,
731        frames: impl Iterator<Item = &'a ConFrame>,
732    ) -> Result<(), WriteError> {
733        for frame in frames {
734            self.write_frame(frame)?;
735        }
736        Ok(())
737    }
738}
739
740fn written_fixed_flag(fixed: [bool; 3], spec_version: u32) -> u8 {
741    if spec_version < 2 {
742        if fixed.iter().any(|axis| *axis) { 1 } else { 0 }
743    } else {
744        encode_fixed_bitmask(fixed)
745    }
746}
747
748fn push_u64(buf: &mut Vec<u8>, n: u64) {
749    push_u128(buf, u128::from(n));
750}
751
752fn push_u128(buf: &mut Vec<u8>, mut n: u128) {
753    let mut tmp = [0u8; 40];
754    let mut i = 40;
755    if n == 0 {
756        buf.push(b'0');
757        return;
758    }
759    while n > 0 {
760        i -= 1;
761        tmp[i] = b'0' + (n % 10) as u8;
762        n /= 10;
763    }
764    buf.extend_from_slice(&tmp[i..]);
765}
766
767fn push_f64(buf: &mut Vec<u8>, value: f64, format: FloatFormat) {
768    match format {
769        FloatFormat::DecimalPlaces(precision) => push_f64_prec(buf, value, precision),
770        FloatFormat::RoundTrip => {
771            let _ = write!(buf, "{value:e}");
772        }
773    }
774}
775
776/// Fixed-point `f64` with `prec` digits after the decimal, matching `{:.prec$}`.
777/// Non-finite values and `prec > 17` fall back to `std::fmt`.
778fn push_f64_prec(buf: &mut Vec<u8>, v: f64, prec: usize) {
779    // The six-decimal path uses fixed-point arithmetic; other precisions
780    // use the standard formatter.
781    if !v.is_finite() || prec != 6 {
782        let _ = write!(buf, "{v:.prec$}");
783        return;
784    }
785    if v.is_sign_negative() {
786        buf.push(b'-');
787    }
788    let ax = v.abs();
789    if prec == 0 {
790        push_u128(buf, ax.round() as u128);
791        return;
792    }
793    let scale = 10u128.pow(prec as u32);
794    let n = (ax * scale as f64).round() as u128;
795    let int_part = n / scale;
796    let frac = n % scale;
797    push_u128(buf, int_part);
798    buf.push(b'.');
799    let mut tmp = [b'0'; 20];
800    let mut x = frac;
801    let mut i = prec;
802    while i > 0 {
803        i -= 1;
804        tmp[i] = b'0' + (x % 10) as u8;
805        x /= 10;
806    }
807    buf.extend_from_slice(&tmp[..prec]);
808}
809
810fn push_xyz_line(
811    buf: &mut Vec<u8>,
812    x: f64,
813    y: f64,
814    z: f64,
815    prec: FloatFormat,
816    fixed: u8,
817    atom_id: u64,
818) {
819    push_f64(buf, x, prec);
820    buf.push(b' ');
821    push_f64(buf, y, prec);
822    buf.push(b' ');
823    push_f64(buf, z, prec);
824    buf.push(b' ');
825    push_u64(buf, u64::from(fixed));
826    buf.push(b' ');
827    push_u64(buf, atom_id);
828    buf.push(b'\n');
829}
830
831fn push_scalar_line(buf: &mut Vec<u8>, v: f64, prec: FloatFormat, fixed: u8, atom_id: u64) {
832    push_f64(buf, v, prec);
833    buf.push(b' ');
834    push_u64(buf, u64::from(fixed));
835    buf.push(b' ');
836    push_u64(buf, atom_id);
837    buf.push(b'\n');
838}
839
840// Implementation block specifically for when the writer is a `File`.
841impl ConFrameWriter<File> {
842    /// Creates a new `ConFrameWriter` that writes to a file at the given path.
843    ///
844    /// This is a convenience function that creates the file and wraps it.
845    pub fn from_path<P: AsRef<Path>>(path: P) -> io::Result<Self> {
846        let file = File::create(path)?;
847        Ok(Self::new(file))
848    }
849
850    /// Creates a new `ConFrameWriter` that writes to a file with a custom precision.
851    pub fn from_path_with_precision<P: AsRef<Path>>(path: P, precision: usize) -> io::Result<Self> {
852        let file = File::create(path)?;
853        Ok(Self::with_precision(file, precision))
854    }
855}
856
857// Gzip-compressed writer constructors.
858impl ConFrameWriter<flate2::write::GzEncoder<File>> {
859    /// Creates a gzip-compressed writer for the given path.
860    pub fn from_path_gzip<P: AsRef<Path>>(path: P) -> io::Result<Self> {
861        let encoder = crate::compression::gzip_writer(path.as_ref())?;
862        Ok(Self::new(encoder))
863    }
864
865    /// Creates a gzip-compressed writer with custom precision.
866    pub fn from_path_gzip_with_precision<P: AsRef<Path>>(
867        path: P,
868        precision: usize,
869    ) -> io::Result<Self> {
870        let encoder = crate::compression::gzip_writer(path.as_ref())?;
871        Ok(Self::with_precision(encoder, precision))
872    }
873}
874
875// Zstd-compressed writer constructors. Available only with the `zstd`
876// Cargo feature.
877#[cfg(feature = "zstd")]
878impl ConFrameWriter<zstd::stream::write::AutoFinishEncoder<'static, File>> {
879    /// Creates a zstd-compressed writer for the given path.
880    pub fn from_path_zstd<P: AsRef<Path>>(path: P) -> io::Result<Self> {
881        let encoder = crate::compression::zstd_writer(path.as_ref())?;
882        Ok(Self::new(encoder))
883    }
884
885    /// Creates a zstd-compressed writer with custom precision.
886    pub fn from_path_zstd_with_precision<P: AsRef<Path>>(
887        path: P,
888        precision: usize,
889    ) -> io::Result<Self> {
890        let encoder = crate::compression::zstd_writer(path.as_ref())?;
891        Ok(Self::with_precision(encoder, precision))
892    }
893}
894
895#[cfg(test)]
896mod float_format_tests {
897    use super::push_f64_prec;
898
899    fn formatted(v: f64, prec: usize) -> String {
900        let mut buf = Vec::new();
901        push_f64_prec(&mut buf, v, prec);
902        String::from_utf8(buf).expect("utf8")
903    }
904
905    #[test]
906    fn matches_std_fixed_precision() {
907        let vals = [
908            0.0,
909            -0.0,
910            1.0,
911            -1.0,
912            0.9045,
913            6.975_299_999_999_995,
914            63.546,
915            1.008,
916            15.3456,
917            90.0,
918            218.0,
919            1e-6,
920            -1.23456789,
921            10.0,
922        ];
923        for prec in [0usize, 6] {
924            for v in vals {
925                let got = formatted(v, prec);
926                let exp = format!("{v:.prec$}");
927                assert_eq!(got, exp, "v={v:?} prec={prec}");
928            }
929        }
930    }
931}
932
933#[cfg(test)]
934mod section_arity_tests {
935    use super::require_minor2_section_arity;
936    use serde_json::json;
937
938    #[test]
939    fn minor2_lists_a_new_section_in_section_arity() {
940        let widgets = vec![json!("widgets")];
941        let err = require_minor2_section_arity(2, &widgets, &serde_json::Map::new()).unwrap_err();
942        let msg = err.to_string();
943        assert!(msg.contains("section_arity"), "{msg}");
944
945        let mut listed = serde_json::Map::new();
946        listed.insert("section_arity".into(), json!({"widgets": 3}));
947        require_minor2_section_arity(2, &widgets, &listed).unwrap();
948
949        let reserved = vec![json!("velocities"), json!("energies")];
950        require_minor2_section_arity(2, &reserved, &serde_json::Map::new()).unwrap();
951        require_minor2_section_arity(1, &widgets, &serde_json::Map::new()).unwrap();
952    }
953}