Skip to main content

readcon_db/
keys.rs

1/// Stable trajectory identifier (user-assigned).
2pub type TrajId = u64;
3/// Zero-based frame index within a trajectory.
4pub type FrameIdx = u32;
5
6/// xxHash3 128-bit content fingerprint (exact match / dedup).
7#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
8pub struct ContentHash(pub [u8; 16]);
9
10impl ContentHash {
11    pub fn to_bytes(self) -> [u8; 16] {
12        self.0
13    }
14    pub fn from_bytes(b: &[u8]) -> Option<Self> {
15        if b.len() != 16 {
16            return None;
17        }
18        let mut a = [0u8; 16];
19        a.copy_from_slice(b);
20        Some(Self(a))
21    }
22    pub fn to_hex(self) -> String {
23        self.0.iter().map(|b| format!("{b:02x}")).collect()
24    }
25}
26
27/// Hash canonical frame blob bytes (UTF-8 CON text as stored).
28pub fn hash_frame_bytes(blob: &[u8]) -> ContentHash {
29    use xxhash_rust::xxh3::xxh3_128;
30    let h = xxh3_128(blob);
31    ContentHash(h.to_le_bytes())
32}
33
34#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
35pub struct FrameKey {
36    pub traj_id: TrajId,
37    pub frame_idx: FrameIdx,
38}
39
40impl FrameKey {
41    pub fn to_bytes(self) -> [u8; 12] {
42        let mut out = [0u8; 12];
43        out[..8].copy_from_slice(&self.traj_id.to_be_bytes());
44        out[8..].copy_from_slice(&self.frame_idx.to_be_bytes());
45        out
46    }
47
48    pub fn from_bytes(b: &[u8]) -> Option<Self> {
49        if b.len() != 12 {
50            return None;
51        }
52        let mut t = [0u8; 8];
53        let mut f = [0u8; 4];
54        t.copy_from_slice(&b[..8]);
55        f.copy_from_slice(&b[8..]);
56        Some(Self {
57            traj_id: u64::from_be_bytes(t),
58            frame_idx: u32::from_be_bytes(f),
59        })
60    }
61}
62
63/// Secondary index key: n_atoms (BE u32) || FrameKey bytes
64pub(crate) fn natoms_key(n_atoms: u32, fk: FrameKey) -> [u8; 16] {
65    let mut out = [0u8; 16];
66    out[..4].copy_from_slice(&n_atoms.to_be_bytes());
67    out[4..].copy_from_slice(&fk.to_bytes());
68    out
69}
70
71/// symbol (utf-8) || 0xff || FrameKey
72pub(crate) fn symbol_key(symbol: &str, fk: FrameKey) -> Vec<u8> {
73    let mut v = symbol.as_bytes().to_vec();
74    v.push(0xff);
75    v.extend_from_slice(&fk.to_bytes());
76    v
77}
78
79pub(crate) fn symbol_prefix(symbol: &str) -> Vec<u8> {
80    let mut v = symbol.as_bytes().to_vec();
81    v.push(0xff);
82    v
83}
84
85/// Order-preserving map of finite f64 → u64 for BE key prefixes.
86pub(crate) fn ordered_f64_bits(x: f64) -> Option<u64> {
87    if !x.is_finite() {
88        return None;
89    }
90    let bits = x.to_bits();
91    Some(if x.is_sign_negative() {
92        !bits
93    } else {
94        bits ^ (1u64 << 63)
95    })
96}
97
98/// Quantize finite energy for ordered range scans (f64 bits, BE; NaN/Inf skipped at index time).
99pub(crate) fn energy_bin_key(energy: f64, fk: FrameKey) -> Option<[u8; 20]> {
100    let ordered = ordered_f64_bits(energy)?;
101    let mut out = [0u8; 20];
102    out[..8].copy_from_slice(&ordered.to_be_bytes());
103    out[8..].copy_from_slice(&fk.to_bytes());
104    Some(out)
105}
106
107/// Max force magnitude bin (only when forces exist and fmax is finite).
108pub(crate) fn fmax_bin_key(fmax: f64, fk: FrameKey) -> Option<[u8; 20]> {
109    energy_bin_key(fmax, fk)
110}
111
112/// Section / capability flag key: flag_id (u8) || FrameKey
113pub(crate) fn flag_key(flag: u8, fk: FrameKey) -> [u8; 13] {
114    let mut out = [0u8; 13];
115    out[0] = flag;
116    out[1..].copy_from_slice(&fk.to_bytes());
117    out
118}
119
120pub(crate) const FLAG_HAS_FORCES: u8 = 1;
121pub(crate) const FLAG_HAS_VELOCITIES: u8 = 2;
122pub(crate) const FLAG_HAS_ENERGY: u8 = 3;
123
124/// Per-element count: symbol || 0xff || BE u32 count || FrameKey
125pub(crate) fn elem_count_key(symbol: &str, count: u32, fk: FrameKey) -> Vec<u8> {
126    let mut v = symbol.as_bytes().to_vec();
127    v.push(0xff);
128    v.extend_from_slice(&count.to_be_bytes());
129    v.extend_from_slice(&fk.to_bytes());
130    v
131}
132
133pub(crate) fn elem_count_symbol_prefix(symbol: &str) -> Vec<u8> {
134    let mut v = symbol.as_bytes().to_vec();
135    v.push(0xff);
136    v
137}
138
139/// Parse count and FrameKey from an `elem_count_key` byte key (after verifying prefix).
140pub(crate) fn parse_elem_count_key(k: &[u8], symbol: &str) -> Option<(u32, FrameKey)> {
141    let pref = elem_count_symbol_prefix(symbol);
142    if !k.starts_with(&pref) || k.len() < pref.len() + 4 + 12 {
143        return None;
144    }
145    let mut cb = [0u8; 4];
146    cb.copy_from_slice(&k[pref.len()..pref.len() + 4]);
147    let count = u32::from_be_bytes(cb);
148    let fk = FrameKey::from_bytes(&k[pref.len() + 4..pref.len() + 4 + 12])?;
149    Some((count, fk))
150}
151
152/// Canonical formula string — delegated to [`readcon_core::index_proj`] (single encoding).
153pub use readcon_core::index_proj::{composition_formula, species_counts_from_symbols};
154
155/// formula || 0xff || FrameKey
156pub(crate) fn formula_key(formula: &str, fk: FrameKey) -> Vec<u8> {
157    let mut v = formula.as_bytes().to_vec();
158    v.push(0xff);
159    v.extend_from_slice(&fk.to_bytes());
160    v
161}
162
163pub(crate) fn formula_prefix(formula: &str) -> Vec<u8> {
164    let mut v = formula.as_bytes().to_vec();
165    v.push(0xff);
166    v
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    #[test]
174    fn formula_canonical_order() {
175        let f1 = composition_formula(&[("H".into(), 2), ("Cu".into(), 2)]);
176        let f2 = composition_formula(&[("Cu".into(), 2), ("H".into(), 2)]);
177        assert_eq!(f1, "Cu:2|H:2");
178        assert_eq!(f1, f2);
179    }
180
181    #[test]
182    fn ordered_bits_monotone() {
183        let a = ordered_f64_bits(-1.0).unwrap();
184        let b = ordered_f64_bits(0.0).unwrap();
185        let c = ordered_f64_bits(1.0).unwrap();
186        assert!(a < b && b < c);
187    }
188
189    #[test]
190    fn elem_count_roundtrip_parse() {
191        let fk = FrameKey {
192            traj_id: 3,
193            frame_idx: 7,
194        };
195        let k = elem_count_key("Cu", 2, fk);
196        let (c, fk2) = parse_elem_count_key(&k, "Cu").unwrap();
197        assert_eq!(c, 2);
198        assert_eq!(fk2, fk);
199    }
200}
201
202/// Meta scalar channel ids for `idx_meta` (u8 prefix).
203pub(crate) const META_TIME: u8 = 1;
204pub(crate) const META_TIMESTEP: u8 = 2;
205pub(crate) const META_FRAME_INDEX: u8 = 3;
206pub(crate) const META_NEB_BEAD: u8 = 4;
207pub(crate) const META_NEB_BAND: u8 = 5;
208pub(crate) const META_CHARGE: u8 = 6;
209pub(crate) const META_MAGMOM: u8 = 7;
210
211/// Ordered scalar bin: channel_id || ord(f64) BE || FrameKey (21 bytes)
212pub(crate) fn meta_scalar_key(channel: u8, value: f64, fk: FrameKey) -> Option<[u8; 21]> {
213    let ordered = ordered_f64_bits(value)?;
214    let mut out = [0u8; 21];
215    out[0] = channel;
216    out[1..9].copy_from_slice(&ordered.to_be_bytes());
217    out[9..].copy_from_slice(&fk.to_bytes());
218    Some(out)
219}
220
221pub(crate) fn meta_channel_prefix(channel: u8) -> [u8; 1] {
222    [channel]
223}
224
225/// PBC mask key: bit0=x, bit1=y, bit2=z (true=1); only written when metadata has pbc.
226pub(crate) fn pbc_key(mask: u8, fk: FrameKey) -> [u8; 13] {
227    let mut out = [0u8; 13];
228    out[0] = mask & 0x07;
229    out[1..].copy_from_slice(&fk.to_bytes());
230    out
231}
232
233pub(crate) fn pbc_mask_from_bools(p: [bool; 3]) -> u8 {
234    (p[0] as u8) | ((p[1] as u8) << 1) | ((p[2] as u8) << 2)
235}
236
237/// Mass / volume use same layout as energy (20 bytes, no channel).
238pub(crate) fn mass_bin_key(mass: f64, fk: FrameKey) -> Option<[u8; 20]> {
239    energy_bin_key(mass, fk)
240}
241
242pub(crate) fn volume_bin_key(vol: f64, fk: FrameKey) -> Option<[u8; 20]> {
243    energy_bin_key(vol, fk)
244}