readcon_db/
cooked_soa.rs

1//! Optional **cooked SoA** payload: derived binary numerics beside authoritative CON text.
2//!
3//! # Why CON text is still required (RCSO is **not** fully equivalent)
4//!
5//! RCSO stores only POD numerics (positions and optional forces/velocities). It does
6//! **not** carry element symbols, masses, cell/angles, constraint masks, JSON metadata,
7//! section labels, or exact on-disk CON bytes. Therefore it cannot replace
8//! `frames` for xxHash3 dedup, join/split fidelity, `reindex`, formula/symbol indexes,
9//! or CON export. The improvement is skipping **CON parse on numeric hot paths** when a
10//! valid cooked blob exists—not omitting the text tier from storage.
11//!
12//! Layout (little-endian, version 1):
13//! - magic `RCSO` (4 bytes)
14//! - version `u32` (=1)
15//! - natoms `u32`
16//! - flags `u32` (bit0 = forces block present, bit1 = velocities block present)
17//! - pos_dtype `u8` (0 = f64 row-major N×3); pad 3 bytes
18//! - reserved `u32` (=0)
19//! - positions: `natoms * 3` f64 LE
20//! - forces (if flag): `natoms * 3` f64 LE
21//! - velocities (if flag): `natoms * 3` f64 LE
22//!
23//! Never used for xxHash dedup, join-split fidelity, or secondary index rebuild.
24//! Rebuild by re-parsing CON text → [`CookedSoa::encode_frame`].
25
26use readcon_core::types::ConFrame;
27
28use crate::error::{Error, Result};
29
30pub const COOKED_MAGIC: &[u8; 4] = b"RCSO";
31pub const COOKED_VERSION: u32 = 1;
32pub const DTYPE_F64: u8 = 0;
33
34pub const FLAG_FORCES: u32 = 1 << 0;
35pub const FLAG_VELOCITIES: u32 = 1 << 1;
36
37const HEADER_LEN: usize = 4 + 4 + 4 + 4 + 1 + 3 + 4; // 24
38
39/// Decoded cooked numerics (always f64 for v1).
40#[derive(Clone, Debug, PartialEq)]
41pub struct CookedSoa {
42    pub natoms: u32,
43    pub positions: Vec<[f64; 3]>,
44    pub forces: Option<Vec<[f64; 3]>>,
45    pub velocities: Option<Vec<[f64; 3]>>,
46}
47
48impl CookedSoa {
49    pub fn encode_frame(frame: &ConFrame) -> Result<Vec<u8>> {
50        let n = frame.atom_data.len();
51        if n > u32::MAX as usize {
52            return Err(Error::Message("too many atoms for cooked SoA".into()));
53        }
54        let natoms = n as u32;
55        let mut flags = 0u32;
56        let has_f = frame.atom_data.iter().any(|a| a.force.is_some());
57        let has_v = frame.atom_data.iter().any(|a| a.velocity.is_some());
58        if has_f {
59            flags |= FLAG_FORCES;
60        }
61        if has_v {
62            flags |= FLAG_VELOCITIES;
63        }
64
65        let mut out = Vec::with_capacity(HEADER_LEN + n * 3 * 8 * (1 + has_f as usize + has_v as usize));
66        out.extend_from_slice(COOKED_MAGIC);
67        out.extend_from_slice(&COOKED_VERSION.to_le_bytes());
68        out.extend_from_slice(&natoms.to_le_bytes());
69        out.extend_from_slice(&flags.to_le_bytes());
70        out.push(DTYPE_F64);
71        out.extend_from_slice(&[0u8; 3]);
72        out.extend_from_slice(&0u32.to_le_bytes());
73
74        for a in &frame.atom_data {
75            for c in [a.x, a.y, a.z] {
76                out.extend_from_slice(&c.to_le_bytes());
77            }
78        }
79        if has_f {
80            for a in &frame.atom_data {
81                let f = a.force.unwrap_or([0.0; 3]);
82                for c in f {
83                    out.extend_from_slice(&c.to_le_bytes());
84                }
85            }
86        }
87        if has_v {
88            for a in &frame.atom_data {
89                let v = a.velocity.unwrap_or([0.0; 3]);
90                for c in v {
91                    out.extend_from_slice(&c.to_le_bytes());
92                }
93            }
94        }
95        Ok(out)
96    }
97
98    pub fn decode(bytes: &[u8]) -> Result<Self> {
99        if bytes.len() < HEADER_LEN {
100            return Err(Error::Message("cooked SoA truncated header".into()));
101        }
102        if &bytes[0..4] != COOKED_MAGIC {
103            return Err(Error::Message("cooked SoA bad magic".into()));
104        }
105        let version = u32::from_le_bytes(bytes[4..8].try_into().unwrap());
106        if version != COOKED_VERSION {
107            return Err(Error::Message(format!(
108                "cooked SoA unsupported version {version}"
109            )));
110        }
111        let natoms = u32::from_le_bytes(bytes[8..12].try_into().unwrap());
112        let flags = u32::from_le_bytes(bytes[12..16].try_into().unwrap());
113        let dtype = bytes[16];
114        if dtype != DTYPE_F64 {
115            return Err(Error::Message(format!(
116                "cooked SoA unsupported dtype {dtype}"
117            )));
118        }
119        let n = natoms as usize;
120        let block = n.checked_mul(3).ok_or_else(|| Error::Message("overflow".into()))?;
121        let block_bytes = block
122            .checked_mul(8)
123            .ok_or_else(|| Error::Message("overflow".into()))?;
124        let mut need = HEADER_LEN + block_bytes;
125        let has_f = flags & FLAG_FORCES != 0;
126        let has_v = flags & FLAG_VELOCITIES != 0;
127        if has_f {
128            need = need
129                .checked_add(block_bytes)
130                .ok_or_else(|| Error::Message("overflow".into()))?;
131        }
132        if has_v {
133            need = need
134                .checked_add(block_bytes)
135                .ok_or_else(|| Error::Message("overflow".into()))?;
136        }
137        if bytes.len() < need {
138            return Err(Error::Message("cooked SoA truncated body".into()));
139        }
140
141        let mut off = HEADER_LEN;
142        let positions = read_vec3_block(&bytes[off..off + block_bytes], n)?;
143        off += block_bytes;
144        let forces = if has_f {
145            let f = read_vec3_block(&bytes[off..off + block_bytes], n)?;
146            off += block_bytes;
147            Some(f)
148        } else {
149            None
150        };
151        let velocities = if has_v {
152            let v = read_vec3_block(&bytes[off..off + block_bytes], n)?;
153            Some(v)
154        } else {
155            None
156        };
157        let _ = off;
158        Ok(Self {
159            natoms,
160            positions,
161            forces,
162            velocities,
163        })
164    }
165
166    /// Prefer cooked bytes; return None if missing/invalid (caller parses CON).
167    pub fn try_decode(bytes: &[u8]) -> Option<Self> {
168        Self::decode(bytes).ok()
169    }
170}
171
172fn read_vec3_block(bytes: &[u8], n: usize) -> Result<Vec<[f64; 3]>> {
173    let mut out = Vec::with_capacity(n);
174    let mut i = 0;
175    for _ in 0..n {
176        let mut row = [0.0f64; 3];
177        for c in 0..3 {
178            let start = i;
179            let end = start + 8;
180            if end > bytes.len() {
181                return Err(Error::Message("cooked SoA short block".into()));
182            }
183            row[c] = f64::from_le_bytes(bytes[start..end].try_into().unwrap());
184            i = end;
185        }
186        out.push(row);
187    }
188    Ok(out)
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194    use readcon_core::iterators::ConFrameIterator;
195    use std::path::PathBuf;
196
197    fn fixture(name: &str) -> String {
198        let p = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
199            .join("../readcon-core/resources/test")
200            .join(name);
201        std::fs::read_to_string(p).unwrap()
202    }
203
204    #[test]
205    fn encode_decode_positions_fixture() {
206        let text = fixture("tiny_cuh2.con");
207        let fr = ConFrameIterator::new(&text).next().unwrap().unwrap();
208        let bytes = CookedSoa::encode_frame(&fr).unwrap();
209        assert!(bytes.len() > HEADER_LEN);
210        let cooked = CookedSoa::decode(&bytes).unwrap();
211        assert_eq!(cooked.natoms as usize, fr.atom_data.len());
212        for (i, a) in fr.atom_data.iter().enumerate() {
213            assert_eq!(cooked.positions[i], [a.x, a.y, a.z]);
214        }
215        assert!(cooked.forces.is_none());
216    }
217
218    #[test]
219    fn encode_decode_forces_fixture() {
220        let text = fixture("tiny_cuh2_forces.con");
221        let fr = ConFrameIterator::new(&text).next().unwrap().unwrap();
222        let bytes = CookedSoa::encode_frame(&fr).unwrap();
223        let cooked = CookedSoa::decode(&bytes).unwrap();
224        assert!(cooked.forces.is_some());
225        let forces = cooked.forces.as_ref().unwrap();
226        for (i, a) in fr.atom_data.iter().enumerate() {
227            assert_eq!(cooked.positions[i], [a.x, a.y, a.z]);
228            if let Some(f) = a.force {
229                assert_eq!(forces[i], f);
230            }
231        }
232    }
233
234    #[test]
235    fn bad_magic_rejected() {
236        assert!(CookedSoa::decode(b"XXXX").is_err());
237        assert!(CookedSoa::try_decode(b"XXXX").is_none());
238    }
239}
readcon_db/cooked_soa.rs

readcon_db/
cooked_soa.rs
