Skip to main content

readcon_db/
corpus.rs

1use std::collections::BTreeSet;
2use std::fs;
3use std::io::Cursor;
4use std::path::{Path, PathBuf};
5use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
6use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
7
8use heed::types::{Bytes, Str, Unit};
9use heed::{CompactionOption, Database, Env, EnvFlags, EnvOpenOptions, RwTxn};
10use readcon_core::iterators::ConFrameIterator;
11use readcon_core::types::ConFrame;
12use readcon_core::writer::{ConFrameWriter, FloatFormat};
13use serde::{Deserialize, Serialize};
14
15use crate::error::{Error, Result};
16use crate::export_xyz::write_frame_extxyz;
17use crate::frame_scalars::{
18    frame_cell_volume, frame_charge, frame_frame_index, frame_magmom, frame_neb_band,
19    frame_neb_bead, frame_pbc_mask, frame_time, frame_timestep, frame_total_mass,
20};
21use crate::keys::{
22    composition_formula, elem_count_key, elem_count_symbol_prefix, energy_bin_key, flag_key,
23    fmax_bin_key, formula_key, formula_prefix, hash_frame_bytes, mass_bin_key, meta_scalar_key,
24    natoms_key, ordered_f64_bits, parse_elem_count_key, pbc_key, pbc_mask_from_bools, symbol_key,
25    symbol_prefix, topo_key, topo_prefix, volume_bin_key, ContentHash, FrameKey, TrajId,
26    FLAG_HAS_ENERGY, FLAG_HAS_FORCES, FLAG_HAS_VELOCITIES, META_CHARGE, META_FRAME_INDEX,
27    META_MAGMOM, META_NEB_BAND, META_NEB_BEAD, META_TIME, META_TIMESTEP,
28};
29use crate::select::Select;
30use crate::topology::{
31    fingerprint_con_bytes, mixed_topo_error, normalize_topo_hex, resolve_seams_binary,
32    run_seams_fingerprint, AnnotateTopologyOpts, TopologyParams,
33};
34
35/// Default LMDB map size (2 GiB) when neither the caller nor
36/// [`MAP_SIZE_ENV`] asks for more. The map is address space, not disk:
37/// `data.mdb` only grows as pages are written.
38pub const DEFAULT_MAP_SIZE: usize = 2 * 1024 * 1024 * 1024;
39/// Environment variable that overrides [`DEFAULT_MAP_SIZE`]. Accepts a byte
40/// count with an optional `K`, `M`, `G` or `T` suffix (powers of 1024).
41pub const MAP_SIZE_ENV: &str = "READCON_DB_MAP_SIZE";
42/// Map sizes are rounded up to a multiple of this (1 MiB, a multiple of
43/// every LMDB page size in use).
44const MAP_ALIGN: usize = 1 << 20;
45/// How many times a write transaction may grow the map and start over
46/// after `MDB_MAP_FULL`. Each attempt at least doubles the map.
47const MAP_GROW_RETRIES: u32 = 8;
48
49/// Parse a map size such as `8G`, `512M` or `1073741824`.
50pub fn parse_map_size(s: &str) -> Option<usize> {
51    let s = s.trim();
52    let (num, mul) = match s.as_bytes().last()? {
53        b'k' | b'K' => (&s[..s.len() - 1], 1usize << 10),
54        b'm' | b'M' => (&s[..s.len() - 1], 1 << 20),
55        b'g' | b'G' => (&s[..s.len() - 1], 1 << 30),
56        b't' | b'T' => (&s[..s.len() - 1], 1 << 40),
57        _ => (s, 1),
58    };
59    num.trim()
60        .parse::<usize>()
61        .ok()?
62        .checked_mul(mul)
63        .filter(|&n| n > 0)
64}
65
66fn align_map(n: usize) -> usize {
67    n.div_ceil(MAP_ALIGN).saturating_mul(MAP_ALIGN)
68}
69
70/// Map size for opening `path`: the request, else [`MAP_SIZE_ENV`], else
71/// [`DEFAULT_MAP_SIZE`]; never less than twice an existing `data.mdb`, so a
72/// corpus that already outgrew the default keeps room to append.
73fn initial_map_size(path: &Path, requested: Option<usize>) -> Result<usize> {
74    let base = match requested {
75        Some(n) => n,
76        None => match std::env::var(MAP_SIZE_ENV) {
77            Ok(v) => parse_map_size(&v).ok_or_else(|| {
78                Error::Message(format!("{MAP_SIZE_ENV}={v:?} is not a byte count"))
79            })?,
80            Err(_) => DEFAULT_MAP_SIZE,
81        },
82    };
83    let used = fs::metadata(path.join("data.mdb"))
84        .map(|m| m.len() as usize)
85        .unwrap_or(0);
86    Ok(align_map(base.max(used.saturating_mul(2))))
87}
88const MAX_DBS: u32 = 48;
89/// LMDB multi-process reader slots (analysis farms / concurrent CLI processes).
90/// Classic embedded SOTA: one writer, many readers across OS processes sharing the mmap
91/// (Gray/Reuter MVCC; LMDB design notes — not a cluster consensus protocol).
92const MAX_READERS: u32 = 512;
93
94#[derive(Clone, Debug, Default, Serialize, Deserialize)]
95pub struct TrajMeta {
96    pub n_frames: u32,
97    pub source: String,
98    #[serde(default)]
99    pub topo_cutoff: Option<f64>,
100    #[serde(default)]
101    pub topo_graph: Option<String>,
102    #[serde(default)]
103    pub topo_hops: Option<u32>,
104    #[serde(default)]
105    pub topo_method: Option<String>,
106}
107
108impl TrajMeta {
109    fn topo_params(&self) -> TopologyParams {
110        TopologyParams {
111            cutoff: self.topo_cutoff,
112            graph: self.topo_graph.clone(),
113            hops: self.topo_hops,
114            method: self.topo_method.clone(),
115        }
116    }
117
118    fn has_topo(&self) -> bool {
119        !self.topo_params().is_empty()
120    }
121
122    fn stamp_topo(&mut self, cutoff: f64, graph: &str, hops: u32, method: &str) {
123        self.topo_cutoff = Some(cutoff);
124        self.topo_graph = Some(graph.to_owned());
125        self.topo_hops = Some(hops);
126        self.topo_method = Some(method.to_owned());
127    }
128}
129
130/// All secondary index keys precomputed **outside** exclusive write_txn (commit = puts only).
131struct PreparedIndexPuts {
132    fk: FrameKey,
133    /// Bytes written to `frames` (tagged zstd, or the UTF-8 when compression is off or does not shrink).
134    stored: Vec<u8>,
135    /// xxh3-128 of the CON UTF-8, computed before compression.
136    hash: [u8; 16],
137    natoms_k: [u8; 16],
138    symbol_keys: Vec<Vec<u8>>,
139    elem_count_keys: Vec<Vec<u8>>,
140    formula_k: Option<Vec<u8>>,
141    energy_k: Option<[u8; 20]>,
142    fmax_k: Option<[u8; 20]>,
143    mass_k: Option<[u8; 20]>,
144    volume_k: Option<[u8; 20]>,
145    pbc_k: Option<[u8; 13]>,
146    flag_keys: Vec<[u8; 13]>,
147    meta_keys: Vec<[u8; 21]>,
148}
149
150pub struct ConCorpus {
151    path: PathBuf,
152    env: Env,
153    /// Shared by read transactions, exclusive for a write transaction or a
154    /// resize. LMDB requires that this process hold no transaction while the
155    /// map grows, so the write guard stays held across the abort and the
156    /// `resize` call.
157    txn_gate: RwLock<()>,
158    frames: Database<Bytes, Bytes>,
159    /// When false, new frames are stored as raw UTF-8. Existing tagged values still decode.
160    compress_frames: AtomicBool,
161    traj_meta: Database<Bytes, Str>,
162    idx_natoms: Database<Bytes, Unit>,
163    idx_symbol: Database<Bytes, Unit>,
164    idx_energy: Database<Bytes, Unit>,
165    idx_flags: Database<Bytes, Unit>,
166    /// Per-element atom counts: symbol ‖ 0xff ‖ BE count ‖ FrameKey
167    idx_elem_count: Database<Bytes, Unit>,
168    /// Canonical formula ‖ 0xff ‖ FrameKey
169    idx_formula: Database<Bytes, Unit>,
170    /// Ordered max force magnitude (forces present only)
171    idx_fmax: Database<Bytes, Unit>,
172    idx_mass: Database<Bytes, Unit>,
173    idx_volume: Database<Bytes, Unit>,
174    /// PBC mask (metadata present only)
175    idx_pbc: Database<Bytes, Unit>,
176    /// Meta scalars: channel || ord(value) || FrameKey
177    idx_meta: Database<Bytes, Unit>,
178    frame_by_hash: Database<Bytes, Bytes>,
179    hash_by_frame: Database<Bytes, Bytes>,
180    /// Optional derived SoA numerics (`FrameKey` → cooked bytes). Not authoritative.
181    frames_soa: Database<Bytes, Bytes>,
182    /// Optional topology HEX ‖ 0xff ‖ FrameKey. Missing on readonly of older corpora.
183    idx_topo: Option<Database<Bytes, Unit>>,
184    /// Optional FrameKey → HEX. Authority for a matching `reindex` rebuild.
185    topo_by_frame: Option<Database<Bytes, Bytes>>,
186}
187
188fn frame_has_forces(frame: &ConFrame) -> bool {
189    readcon_core::index_proj::frame_has_forces(frame)
190}
191
192fn frame_has_velocities(frame: &ConFrame) -> bool {
193    readcon_core::index_proj::frame_has_velocities(frame)
194}
195
196fn frame_energy(frame: &ConFrame) -> Option<f64> {
197    readcon_core::index_proj::finite_energy(frame)
198}
199
200/// Euclidean max ||F_i|| over atoms with force data; None if no forces.
201pub fn frame_fmax(frame: &ConFrame) -> Option<f64> {
202    readcon_core::index_proj::frame_fmax(frame)
203}
204
205fn frame_species(frame: &ConFrame) -> Vec<(String, u32)> {
206    readcon_core::index_proj::frame_species_counts(frame)
207}
208
209/// Chunk bounds are a function of the input length only (D10 wave thread rule).
210fn ingest_chunk_bounds(n: usize) -> Vec<(usize, usize)> {
211    if n == 0 {
212        return Vec::new();
213    }
214    let chunk = (n / 8).max(1);
215    let mut bounds = Vec::new();
216    let mut i = 0;
217    while i < n {
218        let j = (i + chunk).min(n);
219        bounds.push((i, j));
220        i = j;
221    }
222    bounds
223}
224
225fn ingest_thread_count() -> usize {
226    match std::env::var("RAYON_NUM_THREADS") {
227        Ok(raw) => raw.parse::<usize>().unwrap_or(1).max(1),
228        Err(_) => std::thread::available_parallelism()
229            .map(|n| n.get())
230            .unwrap_or(1),
231    }
232}
233
234/// A read transaction plus the shared lock that keeps a resize from starting
235/// while it is open. `txn` is declared first so it is dropped first.
236struct HeldRead<'a> {
237    txn: heed::RoTxn<'a>,
238    _guard: RwLockReadGuard<'a, ()>,
239}
240
241impl<'a> std::ops::Deref for HeldRead<'a> {
242    type Target = heed::RoTxn<'a>;
243
244    fn deref(&self) -> &Self::Target {
245        &self.txn
246    }
247}
248
249fn heed_map_full(err: &heed::Error) -> bool {
250    matches!(err, heed::Error::Mdb(heed::MdbError::MapFull))
251}
252
253fn is_map_full(err: &Error) -> bool {
254    matches!(err, Error::Heed(e) if heed_map_full(e))
255}
256
257fn prepared_bytes(prepared: &[PreparedIndexPuts]) -> usize {
258    prepared
259        .iter()
260        .fold(0usize, |n, p| n.saturating_add(p.stored.len()))
261}
262
263impl ConCorpus {
264    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
265        Self::open_with(path, false, None)
266    }
267
268    /// Write-open an existing corpus. Does not mkdir or mint `data.mdb`.
269    pub fn open_existing(path: impl AsRef<Path>) -> Result<Self> {
270        let path = path.as_ref();
271        if !path.join("data.mdb").is_file() {
272            return Err(Error::Message(format!(
273                "not a corpus directory (missing data.mdb): {}",
274                path.display()
275            )));
276        }
277        Self::open_with(path, false, None)
278    }
279
280    /// Open or create a writable corpus with an explicit LMDB map size in
281    /// bytes, overriding [`MAP_SIZE_ENV`] and [`DEFAULT_MAP_SIZE`]. The map
282    /// is still at least twice an existing `data.mdb`.
283    pub fn open_with_map_size(path: impl AsRef<Path>, map_size: usize) -> Result<Self> {
284        Self::open_with(path, false, Some(map_size))
285    }
286
287    /// Current LMDB map size in bytes.
288    pub fn map_size(&self) -> usize {
289        self.env.info().map_size
290    }
291
292    fn read_guard(&self) -> RwLockReadGuard<'_, ()> {
293        self.txn_gate.read().unwrap_or_else(|e| e.into_inner())
294    }
295
296    fn write_guard(&self) -> RwLockWriteGuard<'_, ()> {
297        self.txn_gate.write().unwrap_or_else(|e| e.into_inner())
298    }
299
300    /// Read transaction. The guard is held until the transaction is dropped,
301    /// so a resize (which takes the write side) waits for it.
302    fn read_txn(&self) -> Result<HeldRead<'_>> {
303        let guard = self.read_guard();
304        let txn = self.env.read_txn()?;
305        Ok(HeldRead { txn, _guard: guard })
306    }
307
308    fn used_bytes(&self) -> usize {
309        fs::metadata(self.path.join("data.mdb"))
310            .map(|m| m.len() as usize)
311            .unwrap_or(0)
312    }
313
314    /// Grow after `MDB_MAP_FULL`. Caller holds the write side of [`Self::txn_gate`]
315    /// and has dropped every transaction.
316    fn grow_after_map_full(&self, input_bytes: usize) -> Result<()> {
317        let cur = self.map_size();
318        let used_plus = self.used_bytes().saturating_add(input_bytes);
319        let mut want = align_map(cur.saturating_mul(2).max(used_plus));
320        if want <= cur {
321            want = align_map(cur.saturating_add(MAP_ALIGN));
322        }
323        // SAFETY: the write guard is held and the failed transaction was dropped,
324        // so this process has no open transaction.
325        unsafe { self.env.resize(want)? };
326        Ok(())
327    }
328
329    /// Run `body` in one write transaction. On `MDB_MAP_FULL`, abort it, grow
330    /// the map, and retry the whole transaction up to [`MAP_GROW_RETRIES`]
331    /// times. `input_bytes` is the payload size used when choosing the new map.
332    fn with_write<T>(
333        &self,
334        input_bytes: usize,
335        mut body: impl FnMut(&mut RwTxn<'_>) -> Result<T>,
336    ) -> Result<T> {
337        let _guard = self.write_guard();
338        // Grow before the transaction when the payload clearly cannot fit,
339        // so a large serial write does not fill the map and redo its work.
340        // Pages written also carry index keys and B-tree overhead, hence 2x.
341        let need = self
342            .used_bytes()
343            .saturating_add(input_bytes.saturating_mul(2));
344        let cur = self.map_size();
345        if need > cur {
346            let want = align_map(need.max(cur.saturating_mul(2)));
347            // SAFETY: the write guard is held and no transaction is open.
348            unsafe { self.env.resize(want)? };
349        }
350        let mut last = Error::Message("LMDB map stayed full after growing it".into());
351        for _ in 0..MAP_GROW_RETRIES {
352            let mut wtxn = match self.env.write_txn() {
353                Ok(txn) => txn,
354                Err(e) if heed_map_full(&e) => {
355                    last = Error::Heed(e);
356                    self.grow_after_map_full(input_bytes)?;
357                    continue;
358                }
359                Err(e) => return Err(e.into()),
360            };
361            match body(&mut wtxn) {
362                Ok(value) => match wtxn.commit() {
363                    Ok(()) => return Ok(value),
364                    // `commit` consumes the transaction. LMDB frees it even
365                    // when the commit fails, so the map can grow immediately.
366                    Err(e) if heed_map_full(&e) => {
367                        last = Error::Heed(e);
368                        self.grow_after_map_full(input_bytes)?;
369                    }
370                    Err(e) => return Err(e.into()),
371                },
372                Err(e) if is_map_full(&e) => {
373                    drop(wtxn);
374                    last = e;
375                    self.grow_after_map_full(input_bytes)?;
376                }
377                Err(e) => return Err(e),
378            }
379        }
380        Err(last)
381    }
382
383    /// Grow the map so at least `extra_bytes` fit beyond the pages already
384    /// used. Never shrinks. Returns the resulting map size.
385    ///
386    /// The write side of the transaction lock is held across the resize, so
387    /// no transaction from this process is open. The MPI ingest calls it on
388    /// each writer rank with the input size before its first batch.
389    pub fn reserve_map(&self, extra_bytes: usize) -> Result<usize> {
390        let _guard = self.write_guard();
391        let want = align_map(self.used_bytes().saturating_add(extra_bytes));
392        if want > self.map_size() {
393            // SAFETY: write guard is held and this method opens no transaction.
394            unsafe { self.env.resize(want)? };
395        }
396        Ok(self.map_size())
397    }
398
399    /// Open an existing corpus with `MDB_RDONLY`. No mkdir, no write txn.
400    /// MPI workers that only *read* should use this — or, better, rank 0
401    /// of the *caller* communicator [`Self::pack_frame`] and `MPI_Bcast`
402    /// on that same handle so other ranks never open LMDB. Never
403    /// initialize MPI here; the host (LAMMPS, mpi4py) already did.
404    pub fn open_readonly(path: impl AsRef<Path>) -> Result<Self> {
405        Self::open_with(path, true, None)
406    }
407
408    fn open_with(path: impl AsRef<Path>, readonly: bool, map_size: Option<usize>) -> Result<Self> {
409        let path = path.as_ref().to_path_buf();
410        if path.join("shards.json").is_file() {
411            return Err(Error::Message(format!(
412                "refusing to open a sharded root as a single-env corpus: {}",
413                path.display()
414            )));
415        }
416        if readonly {
417            if !path.join("data.mdb").is_file() {
418                return Err(Error::Message(format!(
419                    "readonly open: {} is not a corpus directory",
420                    path.display()
421                )));
422            }
423        } else {
424            fs::create_dir_all(&path)?;
425        }
426        let mut opts = EnvOpenOptions::new();
427        opts.map_size(initial_map_size(&path, map_size)?);
428        opts.max_dbs(MAX_DBS);
429        opts.max_readers(MAX_READERS);
430        if readonly {
431            // heed 0.21: MDB_RDONLY via EnvFlags, not a builder method.
432            unsafe {
433                opts.flags(EnvFlags::READ_ONLY);
434            }
435        }
436        let env = unsafe { opts.open(&path)? };
437
438        if readonly {
439            let rtxn = env.read_txn()?;
440            let missing =
441                |name: &str| Error::Message(format!("readonly open: missing database {name}"));
442            let frames = env
443                .open_database(&rtxn, Some("frames"))?
444                .ok_or_else(|| missing("frames"))?;
445            let traj_meta = env
446                .open_database(&rtxn, Some("traj_meta"))?
447                .ok_or_else(|| missing("traj_meta"))?;
448            let idx_natoms = env
449                .open_database(&rtxn, Some("idx_natoms"))?
450                .ok_or_else(|| missing("idx_natoms"))?;
451            let idx_symbol = env
452                .open_database(&rtxn, Some("idx_symbol"))?
453                .ok_or_else(|| missing("idx_symbol"))?;
454            let idx_energy = env
455                .open_database(&rtxn, Some("idx_energy"))?
456                .ok_or_else(|| missing("idx_energy"))?;
457            let idx_flags = env
458                .open_database(&rtxn, Some("idx_flags"))?
459                .ok_or_else(|| missing("idx_flags"))?;
460            let idx_elem_count = env
461                .open_database(&rtxn, Some("idx_elem_count"))?
462                .ok_or_else(|| missing("idx_elem_count"))?;
463            let idx_formula = env
464                .open_database(&rtxn, Some("idx_formula"))?
465                .ok_or_else(|| missing("idx_formula"))?;
466            let idx_fmax = env
467                .open_database(&rtxn, Some("idx_fmax"))?
468                .ok_or_else(|| missing("idx_fmax"))?;
469            let idx_mass = env
470                .open_database(&rtxn, Some("idx_mass"))?
471                .ok_or_else(|| missing("idx_mass"))?;
472            let idx_volume = env
473                .open_database(&rtxn, Some("idx_volume"))?
474                .ok_or_else(|| missing("idx_volume"))?;
475            let idx_pbc = env
476                .open_database(&rtxn, Some("idx_pbc"))?
477                .ok_or_else(|| missing("idx_pbc"))?;
478            let idx_meta = env
479                .open_database(&rtxn, Some("idx_meta"))?
480                .ok_or_else(|| missing("idx_meta"))?;
481            let frame_by_hash = env
482                .open_database(&rtxn, Some("frame_by_hash"))?
483                .ok_or_else(|| missing("frame_by_hash"))?;
484            let hash_by_frame = env
485                .open_database(&rtxn, Some("hash_by_frame"))?
486                .ok_or_else(|| missing("hash_by_frame"))?;
487            let frames_soa = env
488                .open_database(&rtxn, Some("frames_soa"))?
489                .ok_or_else(|| missing("frames_soa"))?;
490            // Older corpora lack these; treat a missing index as empty.
491            let idx_topo = env.open_database(&rtxn, Some("idx_topo"))?;
492            let topo_by_frame = env.open_database(&rtxn, Some("topo_by_frame"))?;
493            // heed: uncommitted open_database txn → EINVAL on the next RoTxn.
494            rtxn.commit()?;
495            return Ok(Self {
496                path,
497                env,
498                txn_gate: RwLock::new(()),
499                frames,
500                compress_frames: AtomicBool::new(true),
501                traj_meta,
502                idx_natoms,
503                idx_symbol,
504                idx_energy,
505                idx_flags,
506                idx_elem_count,
507                idx_formula,
508                idx_fmax,
509                idx_mass,
510                idx_volume,
511                idx_pbc,
512                idx_meta,
513                frame_by_hash,
514                hash_by_frame,
515                frames_soa,
516                idx_topo,
517                topo_by_frame,
518            });
519        }
520
521        let mut wtxn = env.write_txn()?;
522        let frames = env.create_database(&mut wtxn, Some("frames"))?;
523        let traj_meta = env.create_database(&mut wtxn, Some("traj_meta"))?;
524        let idx_natoms = env.create_database(&mut wtxn, Some("idx_natoms"))?;
525        let idx_symbol = env.create_database(&mut wtxn, Some("idx_symbol"))?;
526        let idx_energy = env.create_database(&mut wtxn, Some("idx_energy"))?;
527        let idx_flags = env.create_database(&mut wtxn, Some("idx_flags"))?;
528        let idx_elem_count = env.create_database(&mut wtxn, Some("idx_elem_count"))?;
529        let idx_formula = env.create_database(&mut wtxn, Some("idx_formula"))?;
530        let idx_fmax = env.create_database(&mut wtxn, Some("idx_fmax"))?;
531        let idx_mass = env.create_database(&mut wtxn, Some("idx_mass"))?;
532        let idx_volume = env.create_database(&mut wtxn, Some("idx_volume"))?;
533        let idx_pbc = env.create_database(&mut wtxn, Some("idx_pbc"))?;
534        let idx_meta = env.create_database(&mut wtxn, Some("idx_meta"))?;
535        let frame_by_hash = env.create_database(&mut wtxn, Some("frame_by_hash"))?;
536        let hash_by_frame = env.create_database(&mut wtxn, Some("hash_by_frame"))?;
537        let frames_soa = env.create_database(&mut wtxn, Some("frames_soa"))?;
538        let idx_topo = env.create_database(&mut wtxn, Some("idx_topo"))?;
539        let topo_by_frame = env.create_database(&mut wtxn, Some("topo_by_frame"))?;
540        wtxn.commit()?;
541
542        Ok(Self {
543            path,
544            env,
545            txn_gate: RwLock::new(()),
546            frames,
547            compress_frames: AtomicBool::new(true),
548            traj_meta,
549            idx_natoms,
550            idx_symbol,
551            idx_energy,
552            idx_flags,
553            idx_elem_count,
554            idx_formula,
555            idx_fmax,
556            idx_mass,
557            idx_volume,
558            idx_pbc,
559            idx_meta,
560            frame_by_hash,
561            hash_by_frame,
562            frames_soa,
563            idx_topo: Some(idx_topo),
564            topo_by_frame: Some(topo_by_frame),
565        })
566    }
567
568    /// RCSO bytes for a unidirectional broadcast on the caller communicator
569    /// (rank 0 of that comm packs, others unpack). Prefers a valid cooked
570    /// blob; otherwise encodes from the CON text.
571    pub fn pack_frame(&self, key: FrameKey) -> Result<Vec<u8>> {
572        if let Some(b) = self.get_cooked_soa_bytes(key)? {
573            if crate::cooked_soa::CookedSoa::try_decode(&b).is_some() {
574                return Ok(b);
575            }
576        }
577        let fr = self.get_frame(key)?;
578        crate::cooked_soa::CookedSoa::encode_frame(&fr)
579    }
580
581    /// Many RCSO blobs in one RCSB envelope for a single Bcast on the caller comm.
582    pub fn pack_frames(&self, keys: &[FrameKey]) -> Result<Vec<u8>> {
583        let mut blobs = Vec::with_capacity(keys.len());
584        for k in keys {
585            blobs.push(self.pack_frame(*k)?);
586        }
587        crate::cooked_soa::encode_batch(&blobs)
588    }
589
590    pub fn path(&self) -> &Path {
591        &self.path
592    }
593
594    /// Wait until this process's LMDB env is gone. Required before
595    /// `open_readonly` on the same path in the same process.
596    pub fn close(self) {
597        self.env.prepare_for_closing().wait();
598    }
599
600    /// Compact snapshot of `data.mdb` only. No lockfile. Creates `dest_dir`
601    /// if needed; refuses if `dest_dir/data.mdb` already exists.
602    pub fn snapshot_to(&self, dest_dir: impl AsRef<Path>) -> Result<()> {
603        let dest = dest_dir.as_ref();
604        fs::create_dir_all(dest)?;
605        let data = dest.join("data.mdb");
606        if data.exists() {
607            return Err(Error::Message(format!(
608                "snapshot dest exists: {}",
609                data.display()
610            )));
611        }
612        self.env.copy_to_file(&data, CompactionOption::Enabled)?;
613        Ok(())
614    }
615
616    /// All frame keys in this env (for compaction / export pipelines).
617    pub fn list_frame_keys(&self) -> Result<Vec<FrameKey>> {
618        let rtxn = self.read_txn()?;
619        let mut out = Vec::new();
620        let mut iter = self.frames.iter(&rtxn)?;
621        while let Some(Ok((k, _))) = iter.next() {
622            if let Some(fk) = FrameKey::from_bytes(k) {
623                out.push(fk);
624            }
625        }
626        out.sort();
627        Ok(out)
628    }
629
630    /// Copy one frame blob+indexes from this corpus into `dst` (same FrameKey).
631    /// Used by join compaction; prepare keys outside dst exclusive section via append path on text.
632    pub fn export_frame_blob(&self, key: FrameKey) -> Result<String> {
633        self.get_frame_text(key)
634    }
635
636    /// When false, later ingests store raw UTF-8. Reads of either form are unchanged.
637    pub fn set_compress_frames(&self, compress: bool) {
638        self.compress_frames.store(compress, Ordering::Relaxed);
639    }
640
641    fn compress_on(&self) -> bool {
642        self.compress_frames.load(Ordering::Relaxed)
643    }
644
645    /// CPU-only: derive every secondary key for one frame (call **outside** write_txn).
646    /// Screening scalars come from [`readcon_core::index_proj::FrameIndexProjection`].
647    /// The content hash is the UTF-8. Compression of that UTF-8 happens here, still
648    /// outside the write transaction.
649    fn build_index_puts(
650        fk: FrameKey,
651        frame: &ConFrame,
652        text: String,
653        compress: bool,
654    ) -> Result<PreparedIndexPuts> {
655        use readcon_core::index_proj::FrameIndexProjection;
656        let hash = hash_frame_bytes(text.as_bytes()).to_bytes();
657        let stored = crate::frame_codec::store_frame(&text, compress)?;
658        let proj = FrameIndexProjection::from_frame(frame);
659        let mut symbol_keys = Vec::new();
660        let mut elem_count_keys = Vec::new();
661        let mut seen = BTreeSet::new();
662        for (sym, cnt) in &proj.species_counts {
663            elem_count_keys.push(elem_count_key(sym, *cnt, fk));
664            if seen.insert(sym.clone()) {
665                symbol_keys.push(symbol_key(sym, fk));
666            }
667        }
668        let formula_k = if proj.formula.is_empty() {
669            None
670        } else {
671            Some(formula_key(&proj.formula, fk))
672        };
673        let mut flag_keys = Vec::new();
674        let energy_k = proj.energy.and_then(|e| {
675            flag_keys.push(flag_key(FLAG_HAS_ENERGY, fk));
676            energy_bin_key(e, fk)
677        });
678        let fmax_k = if proj.has_forces {
679            flag_keys.push(flag_key(FLAG_HAS_FORCES, fk));
680            proj.fmax.and_then(|fm| fmax_bin_key(fm, fk))
681        } else {
682            None
683        };
684        if proj.has_velocities {
685            flag_keys.push(flag_key(FLAG_HAS_VELOCITIES, fk));
686        }
687        let mass_k = proj.total_mass.and_then(|m| mass_bin_key(m, fk));
688        let volume_k = proj.cell_volume.and_then(|v| volume_bin_key(v, fk));
689        let pbc_k = proj
690            .pbc
691            .map(|p| pbc_key(crate::keys::pbc_mask_from_bools(p), fk));
692        let mut meta_keys = Vec::new();
693        for (ch, val) in [
694            (META_TIME, proj.time),
695            (META_TIMESTEP, proj.timestep),
696            (META_FRAME_INDEX, proj.frame_index),
697            (META_NEB_BEAD, proj.neb_bead),
698            (META_NEB_BAND, proj.neb_band),
699            (META_CHARGE, proj.charge),
700            (META_MAGMOM, proj.magmom),
701        ] {
702            if let Some(x) = val {
703                if let Some(k) = meta_scalar_key(ch, x, fk) {
704                    meta_keys.push(k);
705                }
706            }
707        }
708        Ok(PreparedIndexPuts {
709            fk,
710            stored,
711            hash,
712            natoms_k: natoms_key(proj.n_atoms, fk),
713            symbol_keys,
714            elem_count_keys,
715            formula_k,
716            energy_k,
717            fmax_k,
718            mass_k,
719            volume_k,
720            pbc_k,
721            flag_keys,
722            meta_keys,
723        })
724    }
725
726    /// Exclusive section: LMDB puts only (no hash/species/bin derivation).
727    fn put_index_puts(&self, wtxn: &mut RwTxn, p: &PreparedIndexPuts) -> Result<()> {
728        let fk_b = p.fk.to_bytes();
729        self.frames.put(wtxn, &fk_b[..], &p.stored[..])?;
730        self.hash_by_frame.put(wtxn, &fk_b[..], &p.hash[..])?;
731        if self.frame_by_hash.get(wtxn, &p.hash[..])?.is_none() {
732            self.frame_by_hash.put(wtxn, &p.hash[..], &fk_b[..])?;
733        }
734        self.idx_natoms.put(wtxn, &p.natoms_k[..], &())?;
735        for sk in &p.symbol_keys {
736            self.idx_symbol.put(wtxn, &sk[..], &())?;
737        }
738        for ek in &p.elem_count_keys {
739            self.idx_elem_count.put(wtxn, &ek[..], &())?;
740        }
741        if let Some(ref fk_form) = p.formula_k {
742            self.idx_formula.put(wtxn, &fk_form[..], &())?;
743        }
744        if let Some(ek) = p.energy_k {
745            self.idx_energy.put(wtxn, &ek[..], &())?;
746        }
747        if let Some(fk_fm) = p.fmax_k {
748            self.idx_fmax.put(wtxn, &fk_fm[..], &())?;
749        }
750        if let Some(k) = p.mass_k {
751            self.idx_mass.put(wtxn, &k[..], &())?;
752        }
753        if let Some(k) = p.volume_k {
754            self.idx_volume.put(wtxn, &k[..], &())?;
755        }
756        if let Some(pk) = p.pbc_k {
757            self.idx_pbc.put(wtxn, &pk[..], &())?;
758        }
759        for fk_flag in &p.flag_keys {
760            self.idx_flags.put(wtxn, &fk_flag[..], &())?;
761        }
762        for mk in &p.meta_keys {
763            self.idx_meta.put(wtxn, &mk[..], &())?;
764        }
765        Ok(())
766    }
767
768    fn delete_index_puts(&self, wtxn: &mut RwTxn, p: &PreparedIndexPuts) -> Result<()> {
769        let fk_b = p.fk.to_bytes();
770        let _ = self.frames.delete(wtxn, &fk_b[..])?;
771        let _ = self.hash_by_frame.delete(wtxn, &fk_b[..])?;
772        let _ = self.frame_by_hash.delete(wtxn, &p.hash[..])?;
773        let _ = self.frames_soa.delete(wtxn, &fk_b[..])?;
774        let _ = self.idx_natoms.delete(wtxn, &p.natoms_k[..])?;
775        for sk in &p.symbol_keys {
776            let _ = self.idx_symbol.delete(wtxn, &sk[..])?;
777        }
778        for ek in &p.elem_count_keys {
779            let _ = self.idx_elem_count.delete(wtxn, &ek[..])?;
780        }
781        if let Some(ref fk_form) = p.formula_k {
782            let _ = self.idx_formula.delete(wtxn, &fk_form[..])?;
783        }
784        if let Some(ek) = p.energy_k {
785            let _ = self.idx_energy.delete(wtxn, &ek[..])?;
786        }
787        if let Some(fk_fm) = p.fmax_k {
788            let _ = self.idx_fmax.delete(wtxn, &fk_fm[..])?;
789        }
790        if let Some(k) = p.mass_k {
791            let _ = self.idx_mass.delete(wtxn, &k[..])?;
792        }
793        if let Some(k) = p.volume_k {
794            let _ = self.idx_volume.delete(wtxn, &k[..])?;
795        }
796        if let Some(pk) = p.pbc_k {
797            let _ = self.idx_pbc.delete(wtxn, &pk[..])?;
798        }
799        for fk_flag in &p.flag_keys {
800            let _ = self.idx_flags.delete(wtxn, &fk_flag[..])?;
801        }
802        for mk in &p.meta_keys {
803            let _ = self.idx_meta.delete(wtxn, &mk[..])?;
804        }
805        Ok(())
806    }
807
808    /// Reindex path: build puts then put (derivation not under a multi-frame exclusive prepare loop).
809    fn index_frame(
810        &self,
811        wtxn: &mut RwTxn,
812        fk: FrameKey,
813        frame: &ConFrame,
814        blob: &str,
815    ) -> Result<()> {
816        let puts = Self::build_index_puts(fk, frame, blob.to_owned(), self.compress_on())?;
817        self.put_index_puts(wtxn, &puts)
818    }
819
820    /// Parse path/string CON into owned blobs + frames (no exclusive LMDB write section).
821    fn prepare_trajectory_str(
822        traj_id: TrajId,
823        file_contents: &str,
824        compress: bool,
825    ) -> Result<Vec<PreparedIndexPuts>> {
826        let mut out = Vec::new();
827        let mut frame_idx: u32 = 0;
828        let mut iter = ConFrameIterator::new(file_contents);
829        loop {
830            let (frame, blob) = match iter.next_with_raw_span(file_contents) {
831                None => break,
832                Some(Err(e)) => return Err(Error::Parse(e.to_string())),
833                Some(Ok(x)) => x,
834            };
835            let fk = FrameKey { traj_id, frame_idx };
836            out.push(Self::build_index_puts(
837                fk,
838                &frame,
839                blob.to_owned(),
840                compress,
841            )?);
842            frame_idx += 1;
843        }
844        Ok(out)
845    }
846
847    /// Serialize ConFrames to CON text **outside** write_txn (in-memory ingress path).
848    fn prepare_trajectory_frames(
849        traj_id: TrajId,
850        frames: &[ConFrame],
851        start_idx: u32,
852        compress: bool,
853    ) -> Result<Vec<PreparedIndexPuts>> {
854        Self::prepare_trajectory_frames_formatted(
855            traj_id,
856            frames,
857            start_idx,
858            FloatFormat::default(),
859            compress,
860        )
861    }
862
863    fn prepare_trajectory_frames_formatted(
864        traj_id: TrajId,
865        frames: &[ConFrame],
866        start_idx: u32,
867        float_format: FloatFormat,
868        compress: bool,
869    ) -> Result<Vec<PreparedIndexPuts>> {
870        let mut out = Vec::with_capacity(frames.len());
871        let mut frame_idx = start_idx;
872        for fr in frames {
873            let mut buf = Cursor::new(Vec::new());
874            {
875                let mut w = ConFrameWriter::with_float_format(&mut buf, float_format);
876                w.write_frame(fr)
877                    .map_err(|e| Error::Parse(format!("serialize: {e}")))?;
878            }
879            let blob = String::from_utf8(buf.into_inner())
880                .map_err(|e| Error::Message(format!("utf8: {e}")))?;
881            let fk = FrameKey { traj_id, frame_idx };
882            out.push(Self::build_index_puts(fk, fr, blob, compress)?);
883            frame_idx += 1;
884        }
885        Ok(out)
886    }
887
888    /// Fail with `TrajExists` if `traj_id` already has metadata.
889    #[cfg(feature = "mpi")]
890    pub(crate) fn check_traj_free(&self, traj_id: TrajId) -> Result<()> {
891        let rtxn = self.read_txn()?;
892        let tid_key = traj_id.to_be_bytes();
893        if let Some(existing) = self.traj_meta.get(&rtxn, &tid_key[..])? {
894            let meta: TrajMeta = serde_json::from_str(existing)?;
895            return Err(Error::TrajExists(traj_id, meta.n_frames));
896        }
897        Ok(())
898    }
899
900    /// Derive one frame's index keys on any rank and encode them without the
901    /// blob. Returns the 16-byte content hash and the encoded keys.
902    #[cfg(feature = "mpi")]
903    pub(crate) fn encode_frame_keys(
904        fk: FrameKey,
905        frame: &ConFrame,
906        blob: &str,
907    ) -> Result<([u8; 16], Vec<u8>, Vec<u8>)> {
908        let p = Self::build_index_puts(fk, frame, blob.to_owned(), true)?;
909        let mut v = Vec::with_capacity(160);
910        fn put(v: &mut Vec<u8>, b: &[u8]) {
911            v.extend_from_slice(&(b.len() as u32).to_le_bytes());
912            v.extend_from_slice(b);
913        }
914        fn list<T: AsRef<[u8]>>(v: &mut Vec<u8>, xs: &[T]) {
915            v.extend_from_slice(&(xs.len() as u32).to_le_bytes());
916            for x in xs {
917                put(v, x.as_ref());
918            }
919        }
920        v.extend_from_slice(&p.hash);
921        v.extend_from_slice(&p.natoms_k);
922        list(&mut v, &p.symbol_keys);
923        list(&mut v, &p.elem_count_keys);
924        list(&mut v, p.formula_k.as_slice());
925        list(&mut v, p.energy_k.as_slice());
926        list(&mut v, p.fmax_k.as_slice());
927        list(&mut v, p.mass_k.as_slice());
928        list(&mut v, p.volume_k.as_slice());
929        list(&mut v, p.pbc_k.as_slice());
930        list(&mut v, &p.flag_keys);
931        list(&mut v, &p.meta_keys);
932        Ok((p.hash, v, p.stored))
933    }
934
935    /// Insert frames from encoded keys plus blobs, in slice order, in one
936    /// write transaction. Same puts as the serial path.
937    #[cfg(feature = "mpi")]
938    pub(crate) fn put_encoded(&self, rows: Vec<(FrameKey, Vec<u8>, Vec<u8>)>) -> Result<()> {
939        struct Rd<'a>(&'a [u8]);
940        impl<'a> Rd<'a> {
941            fn take(&mut self, n: usize) -> Result<&'a [u8]> {
942                if self.0.len() < n {
943                    return Err(Error::Message("truncated MPI frame record".into()));
944                }
945                let (a, b) = self.0.split_at(n);
946                self.0 = b;
947                Ok(a)
948            }
949            fn u32(&mut self) -> Result<usize> {
950                Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()) as usize)
951            }
952            fn list(&mut self) -> Result<Vec<Vec<u8>>> {
953                let n = self.u32()?;
954                (0..n)
955                    .map(|_| {
956                        let l = self.u32()?;
957                        Ok(self.take(l)?.to_vec())
958                    })
959                    .collect()
960            }
961            fn fixed<const N: usize>(&mut self) -> Result<Vec<[u8; N]>> {
962                self.list()?
963                    .into_iter()
964                    .map(|x| {
965                        x.try_into()
966                            .map_err(|_| Error::Message("bad MPI key width".into()))
967                    })
968                    .collect()
969            }
970        }
971        // Decode once and move each blob in, so a MapFull retry neither
972        // re-parses nor copies the payload.
973        let prepared = rows
974            .into_iter()
975            .map(|(fk, enc, stored)| {
976                let mut r = Rd(&enc);
977                let hash: [u8; 16] = r.take(16)?.try_into().unwrap();
978                let natoms_k: [u8; 16] = r.take(16)?.try_into().unwrap();
979                Ok(PreparedIndexPuts {
980                    fk,
981                    stored,
982                    hash,
983                    natoms_k,
984                    symbol_keys: r.list()?,
985                    elem_count_keys: r.list()?,
986                    formula_k: r.list()?.pop(),
987                    energy_k: r.fixed::<20>()?.pop(),
988                    fmax_k: r.fixed::<20>()?.pop(),
989                    mass_k: r.fixed::<20>()?.pop(),
990                    volume_k: r.fixed::<20>()?.pop(),
991                    pbc_k: r.fixed::<13>()?.pop(),
992                    flag_keys: r.fixed::<13>()?,
993                    meta_keys: r.fixed::<21>()?,
994                })
995            })
996            .collect::<Result<Vec<_>>>()?;
997        self.with_write(prepared_bytes(&prepared), |wtxn| {
998            for p in &prepared {
999                self.put_index_puts(wtxn, p)?;
1000            }
1001            Ok(())
1002        })
1003    }
1004
1005    /// Write trajectory metadata after an MPI ingest (serial `commit_prepared` tail).
1006    #[cfg(feature = "mpi")]
1007    pub(crate) fn put_traj_meta(
1008        &self,
1009        traj_id: TrajId,
1010        n_frames: u32,
1011        source: String,
1012    ) -> Result<()> {
1013        let input_bytes = source.len().saturating_add(64);
1014        self.with_write(input_bytes, |wtxn| {
1015            let meta = TrajMeta {
1016                n_frames,
1017                source: source.clone(),
1018                ..Default::default()
1019            };
1020            self.traj_meta.put(
1021                wtxn,
1022                &traj_id.to_be_bytes()[..],
1023                &serde_json::to_string(&meta)?,
1024            )?;
1025            Ok(())
1026        })
1027    }
1028
1029    /// Every table in LMDB key order as length-prefixed (name, key, value)
1030    /// records. LMDB page layout can differ between equal corpora; this dump
1031    /// cannot, so its digest is the determinism check for parallel ingest.
1032    pub fn canonical_dump(&self) -> Result<Vec<u8>> {
1033        let rtxn = self.read_txn()?;
1034        let mut out = Vec::new();
1035        let mut tables: Vec<(&str, Database<Bytes, Bytes>)> = vec![
1036            ("frames", self.frames.remap_types()),
1037            ("traj_meta", self.traj_meta.remap_types()),
1038            ("idx_natoms", self.idx_natoms.remap_types()),
1039            ("idx_symbol", self.idx_symbol.remap_types()),
1040            ("idx_energy", self.idx_energy.remap_types()),
1041            ("idx_flags", self.idx_flags.remap_types()),
1042            ("idx_elem_count", self.idx_elem_count.remap_types()),
1043            ("idx_formula", self.idx_formula.remap_types()),
1044            ("idx_fmax", self.idx_fmax.remap_types()),
1045            ("idx_mass", self.idx_mass.remap_types()),
1046            ("idx_volume", self.idx_volume.remap_types()),
1047            ("idx_pbc", self.idx_pbc.remap_types()),
1048            ("idx_meta", self.idx_meta.remap_types()),
1049            ("frame_by_hash", self.frame_by_hash),
1050            ("hash_by_frame", self.hash_by_frame),
1051            ("frames_soa", self.frames_soa),
1052        ];
1053        if let Some(d) = self.idx_topo {
1054            tables.push(("idx_topo", d.remap_types()));
1055        }
1056        if let Some(d) = self.topo_by_frame {
1057            tables.push(("topo_by_frame", d));
1058        }
1059        for (name, db) in tables {
1060            for part in [name.as_bytes()] {
1061                out.extend_from_slice(&(part.len() as u64).to_le_bytes());
1062                out.extend_from_slice(part);
1063            }
1064            for kv in db.iter(&rtxn)? {
1065                let (k, v) = kv?;
1066                for part in [k, v] {
1067                    out.extend_from_slice(&(part.len() as u64).to_le_bytes());
1068                    out.extend_from_slice(part);
1069                }
1070            }
1071        }
1072        Ok(out)
1073    }
1074
1075    /// Share one zstd frame across an aligned run when that is shorter than
1076    /// the per-frame values. `prepared` must still be per-frame bytes from
1077    /// [`Self::build_index_puts`]. Runs outside the write transaction.
1078    fn group_prepared(prepared: &mut [PreparedIndexPuts]) -> Result<()> {
1079        if prepared.len() < 2 {
1080            return Ok(());
1081        }
1082        let texts: Vec<(u32, String)> = prepared
1083            .iter()
1084            .map(|p| Ok((p.fk.frame_idx, crate::frame_codec::load_frame(&p.stored)?)))
1085            .collect::<Result<_>>()?;
1086        let view: Vec<(u32, &str)> = texts.iter().map(|(i, s)| (*i, s.as_str())).collect();
1087        let packed = crate::frame_codec::pack_aligned(&view)?;
1088        for (p, stored) in prepared.iter_mut().zip(packed) {
1089            p.stored = stored;
1090        }
1091        Ok(())
1092    }
1093
1094    /// Short exclusive commit: traj_meta guard + puts only (no parse / no ConFrameWriter).
1095    fn commit_prepared(
1096        &self,
1097        traj_id: TrajId,
1098        prepared: &mut [PreparedIndexPuts],
1099        source: String,
1100        replace_meta: bool,
1101        compress: bool,
1102    ) -> Result<u32> {
1103        if compress {
1104            Self::group_prepared(prepared)?;
1105        }
1106        let input_bytes = prepared_bytes(prepared).saturating_add(source.len());
1107        self.with_write(input_bytes, |wtxn| {
1108            let tid_key = traj_id.to_be_bytes();
1109            if !replace_meta {
1110                if let Some(existing) = self.traj_meta.get(wtxn, &tid_key[..])? {
1111                    let meta: TrajMeta = serde_json::from_str(existing)?;
1112                    return Err(Error::TrajExists(traj_id, meta.n_frames));
1113                }
1114            }
1115            for p in prepared.iter() {
1116                self.put_index_puts(wtxn, p)?;
1117            }
1118            let n_frames = prepared.last().map(|p| p.fk.frame_idx + 1).unwrap_or(0);
1119            let existing: Option<TrajMeta> = if replace_meta {
1120                self.traj_meta
1121                    .get(wtxn, &tid_key[..])?
1122                    .map(serde_json::from_str)
1123                    .transpose()?
1124            } else {
1125                None
1126            };
1127            let mut meta = existing.unwrap_or_default();
1128            meta.n_frames = n_frames;
1129            meta.source = source.clone();
1130            self.traj_meta
1131                .put(wtxn, &tid_key[..], &serde_json::to_string(&meta)?)?;
1132            Ok(n_frames)
1133        })
1134    }
1135
1136    pub fn append_trajectory_path(&self, traj_id: TrajId, file: impl AsRef<Path>) -> Result<u32> {
1137        self.append_trajectory_path_units(traj_id, file, None)
1138    }
1139
1140    /// Ingest a CON path and stamp canonical `units` into every frame's metadata.
1141    pub fn append_trajectory_path_units(
1142        &self,
1143        traj_id: TrajId,
1144        file: impl AsRef<Path>,
1145        units: Option<serde_json::Value>,
1146    ) -> Result<u32> {
1147        let text = fs::read_to_string(file.as_ref())?;
1148        let source = file.as_ref().display().to_string();
1149        let Some(units) = units else {
1150            return self.append_trajectory_str(traj_id, &text, source);
1151        };
1152        let units = crate::units_canon::canonicalize_units_object(units)?;
1153        let mut frames = Vec::new();
1154        for item in ConFrameIterator::new(&text) {
1155            let mut fr = item.map_err(|e| Error::Parse(e.to_string()))?;
1156            fr.header.set_units(units.clone());
1157            frames.push(fr);
1158        }
1159        self.append_trajectory_frames(traj_id, &frames, source)
1160    }
1161
1162    /// Rewrite every frame of `traj_id` into `units`. Stored numbers are
1163    /// converted so the new label matches the values (not a silent relabel).
1164    pub fn set_trajectory_units(&self, traj_id: TrajId, units: serde_json::Value) -> Result<u32> {
1165        let keys = self.select(&Select::new().trajectory(traj_id))?;
1166        if keys.is_empty() {
1167            return Err(Error::Message("no frames for traj".into()));
1168        }
1169        let mut frames = Vec::with_capacity(keys.len());
1170        let mut old_puts = Vec::with_capacity(keys.len());
1171        for k in &keys {
1172            let blob = self.get_frame_text(*k)?;
1173            let mut fr = self.get_frame(*k)?;
1174            old_puts.push(Self::build_index_puts(*k, &fr, blob, self.compress_on())?);
1175            let merged = crate::units_canon::merge_units(fr.header.units(), units.clone())?;
1176            crate::units_canon::rescale_frame_units(&mut fr, &merged)?;
1177            fr.header.set_units(merged);
1178            frames.push(fr);
1179        }
1180        // Unit conversion retains the binary64 values of every stored field.
1181        let mut prepared = Self::prepare_trajectory_frames_formatted(
1182            traj_id,
1183            &frames,
1184            0,
1185            FloatFormat::RoundTrip,
1186            self.compress_on(),
1187        )?;
1188        if self.compress_on() {
1189            Self::group_prepared(&mut prepared)?;
1190        }
1191        let input_bytes = prepared_bytes(&prepared);
1192        self.with_write(input_bytes, |wtxn| {
1193            for p in &old_puts {
1194                self.delete_index_puts(wtxn, p)?;
1195            }
1196            for p in &prepared {
1197                self.put_index_puts(wtxn, p)?;
1198            }
1199            Ok(keys.len() as u32)
1200        })
1201    }
1202
1203    /// Units object on one frame (after canonicalize on write).
1204    pub fn frame_units(&self, key: FrameKey) -> Result<Option<serde_json::Value>> {
1205        let fr = self.get_frame(key)?;
1206        Ok(fr.header.metadata.get("units").cloned())
1207    }
1208
1209    pub fn append_trajectory_str(
1210        &self,
1211        traj_id: TrajId,
1212        file_contents: &str,
1213        source: impl Into<String>,
1214    ) -> Result<u32> {
1215        let source = source.into();
1216        // Prepare outside exclusive write section (parallel writers can parse concurrently).
1217        let mut prepared =
1218            Self::prepare_trajectory_str(traj_id, file_contents, self.compress_on())?;
1219        self.commit_prepared(traj_id, &mut prepared, source, false, self.compress_on())
1220    }
1221
1222    /// Ingest already-parsed frames (chemfiles / builder path). Serialize in prepare, not under write_txn.
1223    pub fn append_trajectory_frames(
1224        &self,
1225        traj_id: TrajId,
1226        frames: &[ConFrame],
1227        source: impl Into<String>,
1228    ) -> Result<u32> {
1229        let source = source.into();
1230        let mut prepared = Self::prepare_trajectory_frames(traj_id, frames, 0, self.compress_on())?;
1231        self.commit_prepared(traj_id, &mut prepared, source, false, self.compress_on())
1232    }
1233
1234    /// [`Self::append_trajectory_frames`] with an explicit number of decimal places.
1235    /// Use [`Self::append_trajectory_frames_with_float_format`] and
1236    /// [`FloatFormat::RoundTrip`] for exact binary64 observation storage.
1237    pub fn append_trajectory_frames_with_precision(
1238        &self,
1239        traj_id: TrajId,
1240        frames: &[ConFrame],
1241        source: impl Into<String>,
1242        precision: usize,
1243    ) -> Result<u32> {
1244        self.append_trajectory_frames_with_float_format(
1245            traj_id,
1246            frames,
1247            source,
1248            FloatFormat::DecimalPlaces(precision),
1249        )
1250    }
1251
1252    /// Append frames using the selected CON numeric representation.
1253    pub fn append_trajectory_frames_with_float_format(
1254        &self,
1255        traj_id: TrajId,
1256        frames: &[ConFrame],
1257        source: impl Into<String>,
1258        float_format: FloatFormat,
1259    ) -> Result<u32> {
1260        let source = source.into();
1261        let mut prepared = Self::prepare_trajectory_frames_formatted(
1262            traj_id,
1263            frames,
1264            0,
1265            float_format,
1266            self.compress_on(),
1267        )?;
1268        self.commit_prepared(traj_id, &mut prepared, source, false, self.compress_on())
1269    }
1270
1271    /// Append frames to an existing trajectory (or create it). Returns new total frame count.
1272    /// Concurrent extend of the **same** `traj_id` is rejected if `n_frames` moved (TOCTOU guard).
1273    pub fn extend_trajectory_frames(
1274        &self,
1275        traj_id: TrajId,
1276        frames: &[ConFrame],
1277        source_hint: impl Into<String>,
1278    ) -> Result<u32> {
1279        let (start_idx, prior) = {
1280            let rtxn = self.read_txn()?;
1281            let tid_key = traj_id.to_be_bytes();
1282            if let Some(existing) = self.traj_meta.get(&rtxn, &tid_key[..])? {
1283                let meta: TrajMeta = serde_json::from_str(existing)?;
1284                (meta.n_frames, Some(meta))
1285            } else {
1286                (0u32, None)
1287            }
1288        };
1289        // Key derivation fully outside exclusive section.
1290        let mut prepared =
1291            Self::prepare_trajectory_frames(traj_id, frames, start_idx, self.compress_on())?;
1292        if self.compress_on() {
1293            Self::group_prepared(&mut prepared)?;
1294        }
1295        let hint = source_hint.into();
1296        let input_bytes = prepared_bytes(&prepared).saturating_add(hint.len());
1297        self.with_write(input_bytes, |wtxn| {
1298            let tid_key = traj_id.to_be_bytes();
1299            let live: Option<TrajMeta> =
1300                if let Some(existing) = self.traj_meta.get(wtxn, &tid_key[..])? {
1301                    Some(serde_json::from_str(existing)?)
1302                } else {
1303                    None
1304                };
1305            let live_n = live.as_ref().map(|m| m.n_frames).unwrap_or(0);
1306            if live_n != start_idx {
1307                return Err(Error::Message(format!(
1308                    "concurrent extend race on traj {traj_id}: expected start {start_idx}, live {live_n}"
1309                )));
1310            }
1311            for p in &prepared {
1312                self.put_index_puts(wtxn, p)?;
1313            }
1314            let n_frames = prepared
1315                .last()
1316                .map(|p| p.fk.frame_idx + 1)
1317                .unwrap_or(start_idx);
1318            let mut meta = live.or_else(|| prior.clone()).unwrap_or_else(|| TrajMeta {
1319                source: hint.clone(),
1320                ..Default::default()
1321            });
1322            meta.n_frames = n_frames;
1323            self.traj_meta
1324                .put(wtxn, &tid_key[..], &serde_json::to_string(&meta)?)?;
1325            Ok(n_frames)
1326        })
1327    }
1328
1329    /// [`Self::extend_trajectory_frames`] from a CON path: create the
1330    /// trajectory or append frames after the live count.
1331    pub fn extend_trajectory_path(&self, traj_id: TrajId, file: impl AsRef<Path>) -> Result<u32> {
1332        self.extend_trajectory_path_units(traj_id, file, None)
1333    }
1334
1335    /// Append a CON path, optionally stamping canonical `units` on the new frames.
1336    pub fn extend_trajectory_path_units(
1337        &self,
1338        traj_id: TrajId,
1339        file: impl AsRef<Path>,
1340        units: Option<serde_json::Value>,
1341    ) -> Result<u32> {
1342        let text = fs::read_to_string(file.as_ref())?;
1343        let mut frames = Vec::new();
1344        let units = match units {
1345            Some(u) => Some(crate::units_canon::canonicalize_units_object(u)?),
1346            None => None,
1347        };
1348        for item in ConFrameIterator::new(&text) {
1349            let mut fr = item.map_err(|e| Error::Parse(e.to_string()))?;
1350            if let Some(ref u) = units {
1351                fr.header.set_units(u.clone());
1352            }
1353            frames.push(fr);
1354        }
1355        self.extend_trajectory_frames(traj_id, &frames, file.as_ref().display().to_string())
1356    }
1357
1358    fn clear_secondary(&self, wtxn: &mut RwTxn) -> Result<()> {
1359        self.idx_natoms.clear(wtxn)?;
1360        self.idx_symbol.clear(wtxn)?;
1361        self.idx_energy.clear(wtxn)?;
1362        self.idx_flags.clear(wtxn)?;
1363        self.idx_elem_count.clear(wtxn)?;
1364        self.idx_formula.clear(wtxn)?;
1365        self.idx_fmax.clear(wtxn)?;
1366        self.idx_mass.clear(wtxn)?;
1367        self.idx_volume.clear(wtxn)?;
1368        self.idx_pbc.clear(wtxn)?;
1369        self.idx_meta.clear(wtxn)?;
1370        self.frame_by_hash.clear(wtxn)?;
1371        self.hash_by_frame.clear(wtxn)?;
1372        Ok(())
1373    }
1374
1375    /// Rebuild all secondary indexes from authoritative `frames` blobs (schema upgrade path).
1376    ///
1377    /// `idx_topo` is rebuilt from `topo_by_frame` only when every annotated
1378    /// trajectory records the same cutoff/graph/hops/method. Mixed parameters
1379    /// return an error without writing. Unannotated corpora leave `idx_topo`
1380    /// empty. The engine binary is not required for a matching rebuild.
1381    pub fn reindex(&self) -> Result<u32> {
1382        let recorded = self.collected_topo_params()?;
1383        self.with_write(0, |wtxn| {
1384            self.clear_secondary(wtxn)?;
1385            if let Some(idx) = self.idx_topo {
1386                idx.clear(wtxn)?;
1387            }
1388
1389            let mut n = 0u32;
1390            // Collect keys first (heed iterator + puts on same txn is safer with snapshot list)
1391            let mut keys: Vec<Vec<u8>> = Vec::new();
1392            {
1393                let mut iter = self.frames.iter(wtxn)?;
1394                while let Some(Ok((k, _))) = iter.next() {
1395                    keys.push(k.to_vec());
1396                }
1397            }
1398            for k in keys {
1399                let Some(fk) = FrameKey::from_bytes(&k) else {
1400                    continue;
1401                };
1402                let Some(stored) = self.frames.get(wtxn, &k)? else {
1403                    continue;
1404                };
1405                let stored_owned = stored.to_vec();
1406                let blob_owned = crate::frame_codec::decode_frame(&stored_owned, &mut |idx| {
1407                    let hk = FrameKey {
1408                        traj_id: fk.traj_id,
1409                        frame_idx: idx,
1410                    }
1411                    .to_bytes();
1412                    match self.frames.get(&*wtxn, &hk[..])? {
1413                        Some(b) => Ok(b.to_vec()),
1414                        None => Err(Error::MissingFrame(FrameKey {
1415                            traj_id: fk.traj_id,
1416                            frame_idx: idx,
1417                        })),
1418                    }
1419                })?;
1420                let mut it = ConFrameIterator::new(&blob_owned);
1421                let frame = it
1422                    .next()
1423                    .ok_or_else(|| Error::Parse("empty blob on reindex".into()))?
1424                    .map_err(|e| Error::Parse(e.to_string()))?;
1425                // index_frame would put frames again — only indexes + hashes
1426                let hash = hash_frame_bytes(blob_owned.as_bytes());
1427                let hb = hash.to_bytes();
1428                let fk_b = fk.to_bytes();
1429                self.hash_by_frame.put(wtxn, &fk_b[..], &hb[..])?;
1430                if self.frame_by_hash.get(wtxn, &hb[..])?.is_none() {
1431                    self.frame_by_hash.put(wtxn, &hb[..], &fk_b[..])?;
1432                }
1433                let n_atoms = frame.atom_data.len() as u32;
1434                self.idx_natoms
1435                    .put(wtxn, &natoms_key(n_atoms, fk)[..], &())?;
1436                let counts = frame_species(&frame);
1437                for (sym, cnt) in &counts {
1438                    self.idx_elem_count
1439                        .put(wtxn, &elem_count_key(sym, *cnt, fk)[..], &())?;
1440                    self.idx_symbol.put(wtxn, &symbol_key(sym, fk)[..], &())?;
1441                }
1442                let formula = composition_formula(&counts);
1443                if !formula.is_empty() {
1444                    self.idx_formula
1445                        .put(wtxn, &formula_key(&formula, fk)[..], &())?;
1446                }
1447                if let Some(e) = frame_energy(&frame) {
1448                    if let Some(ek) = energy_bin_key(e, fk) {
1449                        self.idx_energy.put(wtxn, &ek[..], &())?;
1450                    }
1451                    self.idx_flags
1452                        .put(wtxn, &flag_key(FLAG_HAS_ENERGY, fk)[..], &())?;
1453                }
1454                if frame_has_forces(&frame) {
1455                    self.idx_flags
1456                        .put(wtxn, &flag_key(FLAG_HAS_FORCES, fk)[..], &())?;
1457                    if let Some(fm) = frame_fmax(&frame) {
1458                        if let Some(fk_fm) = fmax_bin_key(fm, fk) {
1459                            self.idx_fmax.put(wtxn, &fk_fm[..], &())?;
1460                        }
1461                    }
1462                }
1463                if frame_has_velocities(&frame) {
1464                    self.idx_flags
1465                        .put(wtxn, &flag_key(FLAG_HAS_VELOCITIES, fk)[..], &())?;
1466                }
1467                if let Some(m) = frame_total_mass(&frame) {
1468                    if let Some(k) = mass_bin_key(m, fk) {
1469                        self.idx_mass.put(wtxn, &k[..], &())?;
1470                    }
1471                }
1472                if let Some(v) = frame_cell_volume(&frame) {
1473                    if let Some(k) = volume_bin_key(v, fk) {
1474                        self.idx_volume.put(wtxn, &k[..], &())?;
1475                    }
1476                }
1477                if let Some(mask) = frame_pbc_mask(&frame) {
1478                    self.idx_pbc.put(wtxn, &pbc_key(mask, fk)[..], &())?;
1479                }
1480                let meta_puts: [(u8, Option<f64>); 7] = [
1481                    (META_TIME, frame_time(&frame)),
1482                    (META_TIMESTEP, frame_timestep(&frame)),
1483                    (META_FRAME_INDEX, frame_frame_index(&frame)),
1484                    (META_NEB_BEAD, frame_neb_bead(&frame)),
1485                    (META_NEB_BAND, frame_neb_band(&frame)),
1486                    (META_CHARGE, frame_charge(&frame)),
1487                    (META_MAGMOM, frame_magmom(&frame)),
1488                ];
1489                for (ch, val) in meta_puts {
1490                    if let Some(x) = val {
1491                        if let Some(k) = meta_scalar_key(ch, x, fk) {
1492                            self.idx_meta.put(wtxn, &k[..], &())?;
1493                        }
1494                    }
1495                }
1496                n += 1;
1497            }
1498            if recorded.is_some() {
1499                self.rebuild_idx_topo_from_stored(wtxn)?;
1500            }
1501            Ok(n)
1502        })
1503    }
1504
1505    pub fn traj_meta(&self, traj_id: TrajId) -> Result<Option<TrajMeta>> {
1506        let rtxn = self.read_txn()?;
1507        let tid_key = traj_id.to_be_bytes();
1508        match self.traj_meta.get(&rtxn, &tid_key[..])? {
1509            None => Ok(None),
1510            Some(s) => Ok(Some(serde_json::from_str(s)?)),
1511        }
1512    }
1513
1514    fn list_traj_metas(&self) -> Result<Vec<(TrajId, TrajMeta)>> {
1515        let rtxn = self.read_txn()?;
1516        let mut out = Vec::new();
1517        let mut iter = self.traj_meta.iter(&rtxn)?;
1518        while let Some(Ok((k, v))) = iter.next() {
1519            if k.len() != 8 {
1520                continue;
1521            }
1522            let mut tb = [0u8; 8];
1523            tb.copy_from_slice(k);
1524            out.push((u64::from_be_bytes(tb), serde_json::from_str(v)?));
1525        }
1526        Ok(out)
1527    }
1528
1529    fn collected_topo_params(&self) -> Result<Option<TopologyParams>> {
1530        let mut seen: Option<TopologyParams> = None;
1531        for (_, meta) in self.list_traj_metas()? {
1532            if !meta.has_topo() {
1533                continue;
1534            }
1535            let rec = meta.topo_params();
1536            match &seen {
1537                None => seen = Some(rec),
1538                Some(prev) if prev.agrees(&rec) => {}
1539                Some(prev) => return Err(mixed_topo_error(prev, &rec)),
1540            }
1541        }
1542        Ok(seen)
1543    }
1544
1545    fn require_topo_dbs(&self) -> Result<(Database<Bytes, Unit>, Database<Bytes, Bytes>)> {
1546        match (self.idx_topo, self.topo_by_frame) {
1547            (Some(idx), Some(by)) => Ok((idx, by)),
1548            _ => Err(Error::Message(
1549                "topology indexes missing; open the corpus for write to create idx_topo".into(),
1550            )),
1551        }
1552    }
1553
1554    fn rebuild_idx_topo_from_stored(&self, wtxn: &mut RwTxn) -> Result<u32> {
1555        let (idx_topo, topo_by_frame) = self.require_topo_dbs()?;
1556        let mut pairs: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
1557        {
1558            let mut iter = topo_by_frame.iter(wtxn)?;
1559            while let Some(Ok((k, v))) = iter.next() {
1560                pairs.push((k.to_vec(), v.to_vec()));
1561            }
1562        }
1563        let mut n = 0u32;
1564        for (fk_b, hex_b) in pairs {
1565            let Some(fk) = FrameKey::from_bytes(&fk_b) else {
1566                continue;
1567            };
1568            let hex = std::str::from_utf8(&hex_b)?;
1569            let hex = normalize_topo_hex(hex)?;
1570            idx_topo.put(wtxn, &topo_key(&hex, fk)[..], &())?;
1571            n += 1;
1572        }
1573        Ok(n)
1574    }
1575
1576    fn put_topo_entry(
1577        &self,
1578        wtxn: &mut RwTxn,
1579        idx_topo: Database<Bytes, Unit>,
1580        topo_by_frame: Database<Bytes, Bytes>,
1581        fk: FrameKey,
1582        hex: &str,
1583    ) -> Result<()> {
1584        let fk_b = fk.to_bytes();
1585        if let Some(old) = topo_by_frame.get(wtxn, &fk_b[..])? {
1586            let old_hex = std::str::from_utf8(old)?;
1587            let _ = idx_topo.delete(wtxn, &topo_key(old_hex, fk)[..])?;
1588        }
1589        topo_by_frame.put(wtxn, &fk_b[..], hex.as_bytes())?;
1590        idx_topo.put(wtxn, &topo_key(hex, fk)[..], &())?;
1591        Ok(())
1592    }
1593
1594    /// Stored topology HEX for one frame, if annotated.
1595    pub fn frame_topo_key(&self, key: FrameKey) -> Result<Option<String>> {
1596        let Some(by) = self.topo_by_frame else {
1597            return Ok(None);
1598        };
1599        let rtxn = self.read_txn()?;
1600        let fk_b = key.to_bytes();
1601        match by.get(&rtxn, &fk_b[..])? {
1602            None => Ok(None),
1603            Some(b) => Ok(Some(std::str::from_utf8(b)?.to_owned())),
1604        }
1605    }
1606
1607    /// Populate `idx_topo` / `topo_by_frame` by shelling out to
1608    /// `seams fingerprint FILE --format json`.
1609    ///
1610    /// Cutoff is required. Default graph is `cutoff`, hops is 2. The binary
1611    /// comes from `opts.seams`, `SEAMS`, or `PATH`. A recorded method that
1612    /// disagrees with this run is refused so methods are never mixed.
1613    pub fn annotate_topology(&self, opts: AnnotateTopologyOpts) -> Result<u32> {
1614        let seams = resolve_seams_binary(opts.seams.as_deref())?;
1615        let (idx_topo, topo_by_frame) = self.require_topo_dbs()?;
1616        let recorded = self.collected_topo_params()?;
1617        if let Some(ref rec) = recorded {
1618            let cutoff_diff = rec.cutoff.is_some() && rec.cutoff != Some(opts.cutoff);
1619            let graph_diff = rec.graph.as_deref().is_some_and(|g| g != opts.graph);
1620            let hops_diff = rec.hops.is_some_and(|h| h != opts.hops);
1621            if cutoff_diff || graph_diff || hops_diff {
1622                let incoming = TopologyParams {
1623                    cutoff: Some(opts.cutoff),
1624                    graph: Some(opts.graph.clone()),
1625                    hops: Some(opts.hops),
1626                    method: rec.method.clone(),
1627                };
1628                return Err(mixed_topo_error(rec, &incoming));
1629            }
1630        }
1631        let metas = self.list_traj_metas()?;
1632        let mut pending: Vec<(FrameKey, String)> = Vec::new();
1633        let mut run_method: Option<String> = None;
1634        for (tid, meta) in &metas {
1635            if meta.n_frames == 0 {
1636                continue;
1637            }
1638            let mut blob = String::new();
1639            for frame_idx in 0..meta.n_frames {
1640                blob.push_str(&self.get_frame_text(FrameKey {
1641                    traj_id: *tid,
1642                    frame_idx,
1643                })?);
1644            }
1645            let recs = fingerprint_con_bytes(
1646                &seams,
1647                blob.as_bytes(),
1648                opts.cutoff,
1649                &opts.graph,
1650                opts.hops,
1651            )?;
1652            if recs.len() != meta.n_frames as usize {
1653                return Err(Error::Message(format!(
1654                    "seams fingerprint returned {} objects for traj {tid} with {} frames",
1655                    recs.len(),
1656                    meta.n_frames
1657                )));
1658            }
1659            for (i, rec) in recs.into_iter().enumerate() {
1660                match &run_method {
1661                    None => run_method = Some(rec.method.clone()),
1662                    Some(m) if m == &rec.method => {}
1663                    Some(m) => {
1664                        return Err(Error::Message(format!(
1665                            "seams fingerprint mixed methods in one run ({m} vs {})",
1666                            rec.method
1667                        )));
1668                    }
1669                }
1670                pending.push((
1671                    FrameKey {
1672                        traj_id: *tid,
1673                        frame_idx: i as u32,
1674                    },
1675                    rec.key,
1676                ));
1677            }
1678        }
1679        if let (Some(rec), Some(method)) = (recorded.as_ref(), run_method.as_deref()) {
1680            if rec.method.as_deref().is_some_and(|m| m != method) {
1681                let incoming = TopologyParams {
1682                    cutoff: Some(opts.cutoff),
1683                    graph: Some(opts.graph.clone()),
1684                    hops: Some(opts.hops),
1685                    method: Some(method.to_owned()),
1686                };
1687                return Err(mixed_topo_error(rec, &incoming));
1688            }
1689        }
1690        let Some(method) = run_method else {
1691            return Ok(0);
1692        };
1693        let input_bytes = pending
1694            .iter()
1695            .fold(0usize, |n, (_, hex)| n.saturating_add(hex.len()));
1696        self.with_write(input_bytes, |wtxn| {
1697            for (fk, hex) in &pending {
1698                self.put_topo_entry(wtxn, idx_topo, topo_by_frame, *fk, hex)?;
1699            }
1700            for (tid, meta) in &metas {
1701                let mut meta = meta.clone();
1702                meta.stamp_topo(opts.cutoff, &opts.graph, opts.hops, &method);
1703                let tid_key = tid.to_be_bytes();
1704                self.traj_meta
1705                    .put(wtxn, &tid_key[..], &serde_json::to_string(&meta)?)?;
1706            }
1707            Ok(pending.len() as u32)
1708        })
1709    }
1710
1711    /// Fingerprint `file` with the recorded corpus parameters and look up
1712    /// `idx_topo`. Errors if nothing is annotated or if seams is missing.
1713    pub fn find_by_topology_path(&self, file: impl AsRef<Path>) -> Result<Vec<FrameKey>> {
1714        let rec = self.collected_topo_params()?.ok_or_else(|| {
1715            Error::Message(
1716                "no recorded topology parameters on this corpus; run annotate-topology first"
1717                    .into(),
1718            )
1719        })?;
1720        let cutoff = rec.cutoff_or_err()?;
1721        let graph = rec.graph_or_default();
1722        let hops = rec.hops_or_default();
1723        let seams = resolve_seams_binary(None)?;
1724        let recs = run_seams_fingerprint(&seams, file.as_ref(), cutoff, &graph, hops)?;
1725        self.lookup_topo_keys(recs.into_iter().map(|r| r.key))
1726    }
1727
1728    /// Fingerprint CON text with the recorded corpus parameters.
1729    pub fn find_by_topology_text(&self, text: &str) -> Result<Vec<FrameKey>> {
1730        let rec = self.collected_topo_params()?.ok_or_else(|| {
1731            Error::Message(
1732                "no recorded topology parameters on this corpus; run annotate-topology first"
1733                    .into(),
1734            )
1735        })?;
1736        let cutoff = rec.cutoff_or_err()?;
1737        let graph = rec.graph_or_default();
1738        let hops = rec.hops_or_default();
1739        let seams = resolve_seams_binary(None)?;
1740        let recs = fingerprint_con_bytes(&seams, text.as_bytes(), cutoff, &graph, hops)?;
1741        self.lookup_topo_keys(recs.into_iter().map(|r| r.key))
1742    }
1743
1744    fn lookup_topo_keys<I>(&self, hexes: I) -> Result<Vec<FrameKey>>
1745    where
1746        I: IntoIterator<Item = String>,
1747    {
1748        let mut out = BTreeSet::new();
1749        let Some(idx) = self.idx_topo else {
1750            return Ok(Vec::new());
1751        };
1752        let rtxn = self.read_txn()?;
1753        for hex in hexes {
1754            let pref = topo_prefix(&normalize_topo_hex(&hex)?);
1755            let mut iter = idx.prefix_iter(&rtxn, &pref)?;
1756            while let Some(Ok((k, _))) = iter.next() {
1757                if k.len() < 12 {
1758                    continue;
1759                }
1760                if let Some(fk) = FrameKey::from_bytes(&k[k.len() - 12..]) {
1761                    out.insert(fk);
1762                }
1763            }
1764        }
1765        Ok(out.into_iter().collect())
1766    }
1767
1768    pub fn get_frame_text(&self, key: FrameKey) -> Result<String> {
1769        let rtxn = self.read_txn()?;
1770        self.get_frame_text_txn(&rtxn, key)
1771    }
1772
1773    fn stored_in(&self, rtxn: &heed::RoTxn<'_>, key: FrameKey) -> Result<Vec<u8>> {
1774        let fk_b = key.to_bytes();
1775        match self.frames.get(rtxn, &fk_b[..])? {
1776            Some(b) => Ok(b.to_vec()),
1777            None => Err(Error::MissingFrame(key)),
1778        }
1779    }
1780
1781    fn get_frame_text_txn(&self, rtxn: &heed::RoTxn<'_>, key: FrameKey) -> Result<String> {
1782        let stored = self.stored_in(rtxn, key)?;
1783        crate::frame_codec::decode_frame(&stored, &mut |idx| {
1784            self.stored_in(
1785                rtxn,
1786                FrameKey {
1787                    traj_id: key.traj_id,
1788                    frame_idx: idx,
1789                },
1790            )
1791        })
1792    }
1793
1794    /// Batch point-gets under **one** read transaction (extract / multi-reader hot path).
1795    pub fn get_frame_texts(&self, keys: &[FrameKey]) -> Result<Vec<String>> {
1796        let rtxn = self.read_txn()?;
1797        let mut out = Vec::with_capacity(keys.len());
1798        for &key in keys {
1799            out.push(self.get_frame_text_txn(&rtxn, key)?);
1800        }
1801        Ok(out)
1802    }
1803
1804    /// Materialize every frame blob for `traj_id` in **one** read txn (full extract hot path).
1805    ///
1806    /// Copies each frame's CON text (decompressing a tagged value) into an owned
1807    /// `String` and folds those bytes so the campaign pays for the payload.
1808    /// Returns `(sum of lengths, byte checksum)` for observability.
1809    pub fn touch_trajectory_blobs(&self, traj_id: TrajId, n_frames: u32) -> Result<(u64, u64)> {
1810        let rtxn = self.read_txn()?;
1811        let mut total = 0u64;
1812        let mut checksum = 0u64;
1813        for frame_idx in 0..n_frames {
1814            let key = FrameKey { traj_id, frame_idx };
1815            // Owned copy — same cost class as get_frame_text, batched under one txn.
1816            let owned = self.get_frame_text_txn(&rtxn, key)?;
1817            total += owned.len() as u64;
1818            for (i, b) in owned.as_bytes().iter().enumerate() {
1819                checksum =
1820                    checksum.wrapping_add((*b as u64).wrapping_mul((i as u64).wrapping_add(1)));
1821            }
1822        }
1823        // Prevent dead-code elimination of the fold in optimized builds used by benchmarks.
1824        std::hint::black_box(checksum);
1825        Ok((total, checksum))
1826    }
1827
1828    pub fn get_frame(&self, key: FrameKey) -> Result<ConFrame> {
1829        let text = self.get_frame_text(key)?;
1830        let mut it = ConFrameIterator::new(&text);
1831        let fr = it
1832            .next()
1833            .ok_or_else(|| Error::Parse("empty blob".into()))?
1834            .map_err(|e| Error::Parse(e.to_string()))?;
1835        Ok(fr)
1836    }
1837
1838    /// Raw optional cooked SoA bytes for `key` (None if never cooked).
1839    /// Does **not** read `frames` (CON text).
1840    pub fn get_cooked_soa_bytes(&self, key: FrameKey) -> Result<Option<Vec<u8>>> {
1841        let rtxn = self.read_txn()?;
1842        let fk_b = key.to_bytes();
1843        Ok(self.frames_soa.get(&rtxn, &fk_b[..])?.map(|b| b.to_vec()))
1844    }
1845
1846    /// True if a **valid** RCSO blob is stored (decode succeeds). Missing/corrupt → false.
1847    /// Authority APIs must not use this flag to skip CON text.
1848    pub fn has_valid_cooked_soa(&self, key: FrameKey) -> Result<bool> {
1849        Ok(self.get_cooked_soa(key)?.is_some())
1850    }
1851
1852    /// Decode cooked SoA if present and valid; `None` if missing or corrupt (CON text unchanged).
1853    /// Reads **only** `frames_soa`—no CON parse.
1854    pub fn get_cooked_soa(&self, key: FrameKey) -> Result<Option<crate::cooked_soa::CookedSoa>> {
1855        Ok(self
1856            .get_cooked_soa_bytes(key)?
1857            .and_then(|b| crate::cooked_soa::CookedSoa::try_decode(&b)))
1858    }
1859
1860    /// Prefer valid cooked SoA (**no CON parse** on hit); otherwise parse CON text in `frames`.
1861    /// RCSO is not fully equivalent to CON (no symbols/metadata/exact bytes)—do not omit `frames`.
1862    pub fn get_positions(&self, key: FrameKey) -> Result<Vec<[f64; 3]>> {
1863        // Fast path: frames_soa only.
1864        if let Some(c) = self.get_cooked_soa(key)? {
1865            return Ok(c.positions);
1866        }
1867        // Fallback: authoritative CON text → parse (never invent from indexes alone).
1868        let fr = self.get_frame(key)?;
1869        Ok(fr.atom_data.iter().map(|a| [a.x, a.y, a.z]).collect())
1870    }
1871
1872    /// Prefer cooked forces when the RCSO flag block is present (**no CON parse** on hit);
1873    /// else parse CON (`None` if no forces on frame).
1874    pub fn get_forces(&self, key: FrameKey) -> Result<Option<Vec<[f64; 3]>>> {
1875        if let Some(c) = self.get_cooked_soa(key)? {
1876            if c.forces.is_some() {
1877                return Ok(c.forces);
1878            }
1879            // Cooked present but no forces block — fall through to CON for fidelity.
1880        }
1881        let fr = self.get_frame(key)?;
1882        if !fr.atom_data.iter().any(|a| a.force().is_some()) {
1883            return Ok(None);
1884        }
1885        Ok(Some(
1886            fr.atom_data
1887                .iter()
1888                .map(|a| a.force().unwrap_or([0.0; 3]))
1889                .collect(),
1890        ))
1891    }
1892
1893    /// Prefer cooked velocities when present; else parse CON.
1894    pub fn get_velocities(&self, key: FrameKey) -> Result<Option<Vec<[f64; 3]>>> {
1895        if let Some(c) = self.get_cooked_soa(key)? {
1896            if c.velocities.is_some() {
1897                return Ok(c.velocities);
1898            }
1899        }
1900        let fr = self.get_frame(key)?;
1901        if !fr.atom_data.iter().any(|a| a.velocity().is_some()) {
1902            return Ok(None);
1903        }
1904        Ok(Some(
1905            fr.atom_data
1906                .iter()
1907                .map(|a| a.velocity().unwrap_or([0.0; 3]))
1908                .collect(),
1909        ))
1910    }
1911
1912    /// Cook one frame from authoritative CON text (overwrite existing cooked entry).
1913    pub fn cook_frame(&self, key: FrameKey) -> Result<usize> {
1914        let fr = self.get_frame(key)?;
1915        let bytes = crate::cooked_soa::CookedSoa::encode_frame(&fr)?;
1916        let n = bytes.len();
1917        self.with_write(n, |wtxn| {
1918            let fk_b = key.to_bytes();
1919            self.frames_soa.put(wtxn, &fk_b[..], &bytes[..])?;
1920            Ok(n)
1921        })
1922    }
1923
1924    /// On-disk `frames` value, tagged or legacy. Tests and `migrate` use this.
1925    fn stored_frame(&self, key: FrameKey) -> Result<Vec<u8>> {
1926        let rtxn = self.read_txn()?;
1927        let fk_b = key.to_bytes();
1928        match self.frames.get(&rtxn, &fk_b[..])? {
1929            Some(b) => Ok(b.to_vec()),
1930            None => Err(Error::MissingFrame(key)),
1931        }
1932    }
1933
1934    /// Rewrite each trajectory's frame values to the current codec.
1935    ///
1936    /// Legacy UTF-8 and per-frame zstd become one shared zstd group when that
1937    /// group is shorter. A second call sees those bytes and writes nothing, so
1938    /// a group is not wrapped in another zstd frame. Hashes stay on the UTF-8.
1939    /// This does not cook `frames_soa`.
1940    pub fn migrate_frames(&self) -> Result<u32> {
1941        let keys = self.list_frame_keys()?;
1942        let mut trajs: Vec<TrajId> = keys.iter().map(|k| k.traj_id).collect();
1943        trajs.sort_unstable();
1944        trajs.dedup();
1945        let mut n = 0u32;
1946        for traj in trajs {
1947            n += self.repack_traj(traj)?;
1948        }
1949        Ok(n)
1950    }
1951
1952    fn repack_traj(&self, traj_id: TrajId) -> Result<u32> {
1953        let keys: Vec<FrameKey> = self
1954            .list_frame_keys()?
1955            .into_iter()
1956            .filter(|k| k.traj_id == traj_id)
1957            .collect();
1958        if keys.is_empty() {
1959            return Ok(0);
1960        }
1961        let mut texts = Vec::with_capacity(keys.len());
1962        let mut old = Vec::with_capacity(keys.len());
1963        for key in &keys {
1964            texts.push((key.frame_idx, self.get_frame_text(*key)?));
1965            old.push(self.stored_frame(*key)?);
1966        }
1967        let view: Vec<(u32, &str)> = texts.iter().map(|(i, s)| (*i, s.as_str())).collect();
1968        let packed = crate::frame_codec::pack_aligned(&view)?;
1969        let changed: Vec<(FrameKey, Vec<u8>)> = keys
1970            .into_iter()
1971            .zip(packed)
1972            .enumerate()
1973            .filter_map(|(i, (key, next))| (next != old[i]).then_some((key, next)))
1974            .collect();
1975        let n = changed.len() as u32;
1976        if changed.is_empty() {
1977            return Ok(0);
1978        }
1979        let bytes = changed.iter().map(|(_, b)| b.len()).sum();
1980        self.with_write(bytes, |wtxn| {
1981            for (key, next) in &changed {
1982                let fk_b = key.to_bytes();
1983                self.frames.put(wtxn, &fk_b[..], next)?;
1984            }
1985            Ok(())
1986        })?;
1987        Ok(n)
1988    }
1989
1990    /// Cook all frames that have CON text; returns number of frames cooked.
1991    /// Not called by ingest or [`Self::migrate_frames`].
1992    pub fn recook_all(&self) -> Result<u32> {
1993        let keys = self.list_frame_keys()?;
1994        let mut n = 0u32;
1995        for k in keys {
1996            self.cook_frame(k)?;
1997            n += 1;
1998        }
1999        Ok(n)
2000    }
2001
2002    /// Drop cooked tier for one key (CON text and indexes untouched).
2003    pub fn delete_cooked_soa(&self, key: FrameKey) -> Result<()> {
2004        self.with_write(0, |wtxn| {
2005            let fk_b = key.to_bytes();
2006            let _ = self.frames_soa.delete(wtxn, &fk_b[..])?;
2007            Ok(())
2008        })
2009    }
2010
2011    /// Ingest path with optional cook-on-write (still stores CON text as authority).
2012    pub fn append_trajectory_path_cook(
2013        &self,
2014        traj_id: TrajId,
2015        file: impl AsRef<Path>,
2016        cook: bool,
2017    ) -> Result<u32> {
2018        let n = self.append_trajectory_path(traj_id, file)?;
2019        if cook {
2020            for i in 0..n {
2021                self.cook_frame(FrameKey {
2022                    traj_id,
2023                    frame_idx: i,
2024                })?;
2025            }
2026        }
2027        Ok(n)
2028    }
2029
2030    pub fn frame_hash(&self, key: FrameKey) -> Result<ContentHash> {
2031        let rtxn = self.read_txn()?;
2032        let fk_b = key.to_bytes();
2033        match self.hash_by_frame.get(&rtxn, &fk_b[..])? {
2034            Some(b) => ContentHash::from_bytes(b).ok_or_else(|| Error::Message("bad hash".into())),
2035            None => {
2036                let text = self.get_frame_text_txn(&rtxn, key)?;
2037                Ok(hash_frame_bytes(text.as_bytes()))
2038            }
2039        }
2040    }
2041
2042    pub fn find_by_hash(&self, hash: ContentHash) -> Result<Option<FrameKey>> {
2043        let rtxn = self.read_txn()?;
2044        let hb = hash.to_bytes();
2045        match self.frame_by_hash.get(&rtxn, &hb[..])? {
2046            Some(b) => Ok(FrameKey::from_bytes(b)),
2047            None => Ok(None),
2048        }
2049    }
2050
2051    pub fn hash_con_text(text: &str) -> Result<ContentHash> {
2052        let mut it = ConFrameIterator::new(text);
2053        let frame = it
2054            .next()
2055            .ok_or_else(|| Error::Parse("no frame".into()))?
2056            .map_err(|e| Error::Parse(e.to_string()))?;
2057        let mut buf = Cursor::new(Vec::new());
2058        {
2059            let mut w = ConFrameWriter::new(&mut buf);
2060            w.write_frame(&frame)
2061                .map_err(|e| Error::Parse(format!("re-serialize: {e}")))?;
2062        }
2063        Ok(hash_frame_bytes(&buf.into_inner()))
2064    }
2065
2066    /// Canonical formula for a stored frame (decode once).
2067    pub fn frame_formula(&self, key: FrameKey) -> Result<String> {
2068        let fr = self.get_frame(key)?;
2069        Ok(composition_formula(&frame_species(&fr)))
2070    }
2071
2072    pub fn select(&self, sel: &Select) -> Result<Vec<FrameKey>> {
2073        let rtxn = self.read_txn()?;
2074        let mut sets: Vec<BTreeSet<FrameKey>> = Vec::new();
2075
2076        if let Some(h) = sel.content_hash {
2077            let mut s = BTreeSet::new();
2078            if let Some(b) = self.frame_by_hash.get(&rtxn, &h[..])? {
2079                if let Some(fk) = FrameKey::from_bytes(b) {
2080                    if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2081                        s.insert(fk);
2082                    }
2083                }
2084            }
2085            sets.push(s);
2086        }
2087
2088        if sel.natoms_min.is_some() || sel.natoms_max.is_some() {
2089            let lo = sel.natoms_min.unwrap_or(0);
2090            let hi = sel.natoms_max.unwrap_or(u32::MAX);
2091            let mut s = BTreeSet::new();
2092            let mut iter = self.idx_natoms.iter(&rtxn)?;
2093            while let Some(Ok((k, _))) = iter.next() {
2094                if k.len() < 16 {
2095                    continue;
2096                }
2097                let mut nb = [0u8; 4];
2098                nb.copy_from_slice(&k[..4]);
2099                let n = u32::from_be_bytes(nb);
2100                if n > hi {
2101                    break;
2102                }
2103                if n >= lo {
2104                    if let Some(fk) = FrameKey::from_bytes(&k[4..16]) {
2105                        if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2106                            s.insert(fk);
2107                        }
2108                    }
2109                }
2110            }
2111            sets.push(s);
2112        }
2113
2114        for sym in &sel.symbols_all {
2115            let mut s = BTreeSet::new();
2116            let pref = symbol_prefix(sym);
2117            let mut iter = self.idx_symbol.prefix_iter(&rtxn, &pref)?;
2118            while let Some(Ok((k, _))) = iter.next() {
2119                if k.len() < 12 {
2120                    continue;
2121                }
2122                let fk_bytes = &k[k.len() - 12..];
2123                if let Some(fk) = FrameKey::from_bytes(fk_bytes) {
2124                    if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2125                        s.insert(fk);
2126                    }
2127                }
2128            }
2129            sets.push(s);
2130        }
2131
2132        if let Some(ref formula) = sel.exact_formula {
2133            let mut s = BTreeSet::new();
2134            let pref = formula_prefix(formula);
2135            let mut iter = self.idx_formula.prefix_iter(&rtxn, &pref)?;
2136            while let Some(Ok((k, _))) = iter.next() {
2137                if k.len() < 12 {
2138                    continue;
2139                }
2140                let fk_bytes = &k[k.len() - 12..];
2141                if let Some(fk) = FrameKey::from_bytes(fk_bytes) {
2142                    if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2143                        s.insert(fk);
2144                    }
2145                }
2146            }
2147            sets.push(s);
2148        }
2149
2150        if let Some(ref hex) = sel.topo_key {
2151            let mut s = BTreeSet::new();
2152            if let Some(idx) = self.idx_topo {
2153                let pref = topo_prefix(&normalize_topo_hex(hex)?);
2154                let mut iter = idx.prefix_iter(&rtxn, &pref)?;
2155                while let Some(Ok((k, _))) = iter.next() {
2156                    if k.len() < 12 {
2157                        continue;
2158                    }
2159                    let fk_bytes = &k[k.len() - 12..];
2160                    if let Some(fk) = FrameKey::from_bytes(fk_bytes) {
2161                        if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2162                            s.insert(fk);
2163                        }
2164                    }
2165                }
2166            }
2167            sets.push(s);
2168        }
2169
2170        for (sym, exact_c) in &sel.element_count_exact {
2171            let mut s = BTreeSet::new();
2172            let pref = elem_count_symbol_prefix(sym);
2173            let mut iter = self.idx_elem_count.prefix_iter(&rtxn, &pref)?;
2174            while let Some(Ok((k, _))) = iter.next() {
2175                if let Some((c, fk)) = parse_elem_count_key(k, sym) {
2176                    if c == *exact_c && sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2177                        s.insert(fk);
2178                    }
2179                }
2180            }
2181            sets.push(s);
2182        }
2183
2184        for (sym, min_c) in &sel.element_count_min {
2185            let mut s = BTreeSet::new();
2186            let pref = elem_count_symbol_prefix(sym);
2187            let mut iter = self.idx_elem_count.prefix_iter(&rtxn, &pref)?;
2188            while let Some(Ok((k, _))) = iter.next() {
2189                if let Some((c, fk)) = parse_elem_count_key(k, sym) {
2190                    if c >= *min_c && sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2191                        s.insert(fk);
2192                    }
2193                }
2194            }
2195            sets.push(s);
2196        }
2197
2198        if sel.energy_min.is_some() || sel.energy_max.is_some() {
2199            let lo_e = sel.energy_min.unwrap_or(f64::NEG_INFINITY);
2200            let hi_e = sel.energy_max.unwrap_or(f64::INFINITY);
2201            let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
2202            let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
2203            let mut s = BTreeSet::new();
2204            let mut iter = self.idx_energy.iter(&rtxn)?;
2205            while let Some(Ok((k, _))) = iter.next() {
2206                if k.len() < 20 {
2207                    continue;
2208                }
2209                let mut eb = [0u8; 8];
2210                eb.copy_from_slice(&k[..8]);
2211                let bits = u64::from_be_bytes(eb);
2212                if bits > hi_bits {
2213                    break;
2214                }
2215                if bits >= lo_bits {
2216                    if let Some(fk) = FrameKey::from_bytes(&k[8..20]) {
2217                        if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2218                            s.insert(fk);
2219                        }
2220                    }
2221                }
2222            }
2223            sets.push(s);
2224        }
2225
2226        if sel.fmax_min.is_some() || sel.fmax_max.is_some() {
2227            let lo_e = sel.fmax_min.unwrap_or(0.0);
2228            let hi_e = sel.fmax_max.unwrap_or(f64::INFINITY);
2229            let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
2230            let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
2231            let mut s = BTreeSet::new();
2232            let mut iter = self.idx_fmax.iter(&rtxn)?;
2233            while let Some(Ok((k, _))) = iter.next() {
2234                if k.len() < 20 {
2235                    continue;
2236                }
2237                let mut eb = [0u8; 8];
2238                eb.copy_from_slice(&k[..8]);
2239                let bits = u64::from_be_bytes(eb);
2240                if bits > hi_bits {
2241                    break;
2242                }
2243                if bits >= lo_bits {
2244                    if let Some(fk) = FrameKey::from_bytes(&k[8..20]) {
2245                        if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2246                            s.insert(fk);
2247                        }
2248                    }
2249                }
2250            }
2251            sets.push(s);
2252        }
2253
2254        let scan_ord = |db: Database<Bytes, Unit>,
2255                        lo: Option<f64>,
2256                        hi: Option<f64>|
2257         -> Result<BTreeSet<FrameKey>> {
2258            let lo_e = lo.unwrap_or(f64::NEG_INFINITY);
2259            let hi_e = hi.unwrap_or(f64::INFINITY);
2260            let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
2261            let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
2262            let mut s = BTreeSet::new();
2263            let mut iter = db.iter(&rtxn)?;
2264            while let Some(Ok((k, _))) = iter.next() {
2265                if k.len() < 20 {
2266                    continue;
2267                }
2268                let mut eb = [0u8; 8];
2269                eb.copy_from_slice(&k[..8]);
2270                let bits = u64::from_be_bytes(eb);
2271                if bits > hi_bits {
2272                    break;
2273                }
2274                if bits >= lo_bits {
2275                    if let Some(fk) = FrameKey::from_bytes(&k[8..20]) {
2276                        if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2277                            s.insert(fk);
2278                        }
2279                    }
2280                }
2281            }
2282            Ok(s)
2283        };
2284
2285        if sel.mass_min.is_some() || sel.mass_max.is_some() {
2286            sets.push(scan_ord(self.idx_mass, sel.mass_min, sel.mass_max)?);
2287        }
2288        if sel.volume_min.is_some() || sel.volume_max.is_some() {
2289            sets.push(scan_ord(self.idx_volume, sel.volume_min, sel.volume_max)?);
2290        }
2291
2292        if let Some(pbc) = sel.pbc_exact {
2293            let mask = pbc_mask_from_bools(pbc);
2294            let mut s = BTreeSet::new();
2295            let pref = [mask];
2296            let mut iter = self.idx_pbc.prefix_iter(&rtxn, &pref)?;
2297            while let Some(Ok((k, _))) = iter.next() {
2298                if k.len() < 13 {
2299                    continue;
2300                }
2301                if let Some(fk) = FrameKey::from_bytes(&k[1..13]) {
2302                    if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2303                        s.insert(fk);
2304                    }
2305                }
2306            }
2307            sets.push(s);
2308        }
2309
2310        let scan_meta = |ch: u8, lo: Option<f64>, hi: Option<f64>| -> Result<BTreeSet<FrameKey>> {
2311            let lo_e = lo.unwrap_or(f64::NEG_INFINITY);
2312            let hi_e = hi.unwrap_or(f64::INFINITY);
2313            let lo_bits = ordered_f64_bits(lo_e).unwrap_or(0);
2314            let hi_bits = ordered_f64_bits(hi_e).unwrap_or(u64::MAX);
2315            let mut s = BTreeSet::new();
2316            let pref = [ch];
2317            let mut iter = self.idx_meta.prefix_iter(&rtxn, &pref)?;
2318            while let Some(Ok((k, _))) = iter.next() {
2319                if k.len() < 21 {
2320                    continue;
2321                }
2322                let mut eb = [0u8; 8];
2323                eb.copy_from_slice(&k[1..9]);
2324                let bits = u64::from_be_bytes(eb);
2325                if bits > hi_bits {
2326                    continue; // prefix iter not fully ordered across all, but within channel it is
2327                }
2328                if bits >= lo_bits {
2329                    if let Some(fk) = FrameKey::from_bytes(&k[9..21]) {
2330                        if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2331                            s.insert(fk);
2332                        }
2333                    }
2334                }
2335            }
2336            Ok(s)
2337        };
2338
2339        if sel.time_min.is_some() || sel.time_max.is_some() {
2340            sets.push(scan_meta(META_TIME, sel.time_min, sel.time_max)?);
2341        }
2342        if sel.timestep_min.is_some() || sel.timestep_max.is_some() {
2343            sets.push(scan_meta(
2344                META_TIMESTEP,
2345                sel.timestep_min,
2346                sel.timestep_max,
2347            )?);
2348        }
2349        if sel.frame_index_min.is_some() || sel.frame_index_max.is_some() {
2350            sets.push(scan_meta(
2351                META_FRAME_INDEX,
2352                sel.frame_index_min,
2353                sel.frame_index_max,
2354            )?);
2355        }
2356        if sel.neb_bead_min.is_some() || sel.neb_bead_max.is_some() {
2357            sets.push(scan_meta(
2358                META_NEB_BEAD,
2359                sel.neb_bead_min,
2360                sel.neb_bead_max,
2361            )?);
2362        }
2363        if sel.neb_band_min.is_some() || sel.neb_band_max.is_some() {
2364            sets.push(scan_meta(
2365                META_NEB_BAND,
2366                sel.neb_band_min,
2367                sel.neb_band_max,
2368            )?);
2369        }
2370        if sel.charge_min.is_some() || sel.charge_max.is_some() {
2371            sets.push(scan_meta(META_CHARGE, sel.charge_min, sel.charge_max)?);
2372        }
2373        if sel.magmom_min.is_some() || sel.magmom_max.is_some() {
2374            sets.push(scan_meta(META_MAGMOM, sel.magmom_min, sel.magmom_max)?);
2375        }
2376
2377        let mut push_flag = |flag: u8| -> Result<()> {
2378            let mut s = BTreeSet::new();
2379            let pref = [flag];
2380            let mut iter = self.idx_flags.prefix_iter(&rtxn, &pref)?;
2381            while let Some(Ok((k, _))) = iter.next() {
2382                if k.len() < 13 {
2383                    continue;
2384                }
2385                if let Some(fk) = FrameKey::from_bytes(&k[1..13]) {
2386                    if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2387                        s.insert(fk);
2388                    }
2389                }
2390            }
2391            sets.push(s);
2392            Ok(())
2393        };
2394        if sel.require_forces {
2395            push_flag(FLAG_HAS_FORCES)?;
2396        }
2397        if sel.require_velocities {
2398            push_flag(FLAG_HAS_VELOCITIES)?;
2399        }
2400        if sel.require_energy {
2401            push_flag(FLAG_HAS_ENERGY)?;
2402        }
2403
2404        if sets.is_empty() {
2405            let mut s = BTreeSet::new();
2406            let mut iter = self.frames.iter(&rtxn)?;
2407            while let Some(Ok((k, _))) = iter.next() {
2408                if let Some(fk) = FrameKey::from_bytes(k) {
2409                    if sel.traj_id.is_none_or(|t| t == fk.traj_id) {
2410                        s.insert(fk);
2411                    }
2412                }
2413            }
2414            sets.push(s);
2415        } else if let Some(tid) = sel.traj_id {
2416            for set in &mut sets {
2417                set.retain(|fk| fk.traj_id == tid);
2418            }
2419        }
2420
2421        // Intersect smallest-first to minimize temporary allocations.
2422        sets.sort_by_key(|s| s.len());
2423        let mut acc = sets.remove(0);
2424        for s in sets {
2425            if acc.is_empty() {
2426                break;
2427            }
2428            if s.len() < acc.len() {
2429                acc = s.intersection(&acc).copied().collect();
2430            } else {
2431                acc = acc.intersection(&s).copied().collect();
2432            }
2433        }
2434
2435        let mut out: Vec<FrameKey> = acc.into_iter().collect();
2436        out.sort_unstable();
2437        if let Some(lim) = sel.limit {
2438            out.truncate(lim);
2439        }
2440        Ok(out)
2441    }
2442
2443    pub fn export_extxyz(
2444        &self,
2445        keys: &[FrameKey],
2446        path: impl AsRef<Path>,
2447        energy_key: &str,
2448    ) -> Result<usize> {
2449        use std::fs::OpenOptions;
2450        use std::io::BufWriter;
2451        let path = path.as_ref();
2452        if path.exists() {
2453            return Err(Error::Message(format!(
2454                "export_extxyz dest exists: {}",
2455                path.display()
2456            )));
2457        }
2458        let file = OpenOptions::new().write(true).create_new(true).open(path)?;
2459        let written = (|| -> Result<usize> {
2460            use std::io::Write;
2461            let mut w = BufWriter::new(file);
2462            let mut n = 0usize;
2463            for k in keys {
2464                let fr = self.get_frame(*k)?;
2465                write_frame_extxyz(&mut w, &fr, energy_key)?;
2466                n += 1;
2467            }
2468            w.flush()?;
2469            Ok(n)
2470        })();
2471        if written.is_err() {
2472            let _ = std::fs::remove_file(path);
2473        }
2474        written
2475    }
2476
2477    pub fn ingest_directory(
2478        &self,
2479        dir: impl AsRef<Path>,
2480        start_traj_id: TrajId,
2481    ) -> Result<Vec<(TrajId, u32, String)>> {
2482        self.ingest_directory_units(dir, start_traj_id, None)
2483    }
2484
2485    /// Ingest every `.con`/`.convel` in `dir`, optionally stamping `units`.
2486    ///
2487    /// Files are sorted by path. Parse runs in chunks whose bounds depend
2488    /// only on the file count. Commits run in that same order.
2489    /// `RAYON_NUM_THREADS` sets the worker count (default: available parallelism).
2490    pub fn ingest_directory_units(
2491        &self,
2492        dir: impl AsRef<Path>,
2493        start_traj_id: TrajId,
2494        units: Option<serde_json::Value>,
2495    ) -> Result<Vec<(TrajId, u32, String)>> {
2496        self.ingest_directory_units_nthreads(dir, start_traj_id, units, ingest_thread_count())
2497    }
2498
2499    /// [`ingest_directory_units`] with an explicit worker count.
2500    pub fn ingest_directory_units_nthreads(
2501        &self,
2502        dir: impl AsRef<Path>,
2503        start_traj_id: TrajId,
2504        units: Option<serde_json::Value>,
2505        threads: usize,
2506    ) -> Result<Vec<(TrajId, u32, String)>> {
2507        let mut paths: Vec<_> = fs::read_dir(dir)?
2508            .filter_map(|e| e.ok())
2509            .map(|e| e.path())
2510            .filter(|p| {
2511                let name = p.file_name().and_then(|s| s.to_str()).unwrap_or("");
2512                name.ends_with(".con") || name.ends_with(".convel")
2513            })
2514            .collect();
2515        paths.sort();
2516        let n = paths.len();
2517        let bounds = ingest_chunk_bounds(n);
2518        let threads = threads.max(1);
2519        let compress = self.compress_on();
2520        let mut prepared: Vec<Result<Vec<PreparedIndexPuts>>> =
2521            (0..n).map(|_| Ok(Vec::new())).collect();
2522        if threads == 1 || bounds.len() <= 1 {
2523            for (lo, hi) in bounds {
2524                for i in lo..hi {
2525                    let tid = start_traj_id + i as u64;
2526                    prepared[i] = Self::prepare_path(tid, &paths[i], units.as_ref(), compress);
2527                }
2528            }
2529        } else {
2530            let next = AtomicUsize::new(0);
2531            let prepared_slots = std::sync::Mutex::new(prepared);
2532            std::thread::scope(|scope| {
2533                for _ in 0..threads.min(bounds.len()) {
2534                    let next = &next;
2535                    let bounds = &bounds;
2536                    let paths = &paths;
2537                    let units = &units;
2538                    let prepared_slots = &prepared_slots;
2539                    scope.spawn(move || loop {
2540                        let chunk = next.fetch_add(1, Ordering::Relaxed);
2541                        if chunk >= bounds.len() {
2542                            break;
2543                        }
2544                        let (lo, hi) = bounds[chunk];
2545                        for (i, path) in paths.iter().enumerate().take(hi).skip(lo) {
2546                            let tid = start_traj_id + i as u64;
2547                            let row = Self::prepare_path(tid, path, units.as_ref(), compress);
2548                            prepared_slots.lock().unwrap()[i] = row;
2549                        }
2550                    });
2551                }
2552            });
2553            prepared = prepared_slots.into_inner().unwrap();
2554        }
2555        let mut out = Vec::with_capacity(n);
2556        for (i, row) in prepared.into_iter().enumerate() {
2557            let mut puts = row?;
2558            let tid = start_traj_id + i as u64;
2559            let source = paths[i].display().to_string();
2560            let n_frames = self.commit_prepared(tid, &mut puts, source.clone(), false, compress)?;
2561            out.push((tid, n_frames, source));
2562        }
2563        Ok(out)
2564    }
2565
2566    fn prepare_path(
2567        traj_id: TrajId,
2568        path: &Path,
2569        units: Option<&serde_json::Value>,
2570        compress: bool,
2571    ) -> Result<Vec<PreparedIndexPuts>> {
2572        let text = fs::read_to_string(path)?;
2573        let Some(units) = units else {
2574            return Self::prepare_trajectory_str(traj_id, &text, compress);
2575        };
2576        let units = crate::units_canon::canonicalize_units_object(units.clone())?;
2577        let mut frames = Vec::new();
2578        for item in ConFrameIterator::new(&text) {
2579            let mut fr = item.map_err(|e| Error::Parse(e.to_string()))?;
2580            fr.header.set_units(units.clone());
2581            frames.push(fr);
2582        }
2583        Self::prepare_trajectory_frames(traj_id, &frames, 0, compress)
2584    }
2585
2586    /// Sorted `database key value` lines for every logical entry.
2587    /// Values are hex so the dump does not depend on LMDB page layout.
2588    pub fn sorted_logical_dump(&self) -> Result<Vec<String>> {
2589        let rtxn = self.read_txn()?;
2590        let mut lines = Vec::new();
2591        fn hex(bytes: &[u8]) -> String {
2592            const DIGITS: &[u8] = b"0123456789abcdef";
2593            let mut s = String::with_capacity(bytes.len() * 2);
2594            for b in bytes {
2595                s.push(DIGITS[(b >> 4) as usize] as char);
2596                s.push(DIGITS[(b & 0xf) as usize] as char);
2597            }
2598            s
2599        }
2600        {
2601            let mut iter = self.frames.iter(&rtxn)?;
2602            while let Some(Ok((k, v))) = iter.next() {
2603                lines.push(format!("frames\t{}\t{}", hex(k), hex(v)));
2604            }
2605        }
2606        {
2607            let mut iter = self.traj_meta.iter(&rtxn)?;
2608            while let Some(Ok((k, v))) = iter.next() {
2609                lines.push(format!("traj_meta\t{}\t{}", hex(k), hex(v.as_bytes())));
2610            }
2611        }
2612        {
2613            let mut iter = self.frame_by_hash.iter(&rtxn)?;
2614            while let Some(Ok((k, v))) = iter.next() {
2615                lines.push(format!("frame_by_hash\t{}\t{}", hex(k), hex(v)));
2616            }
2617        }
2618        {
2619            let mut iter = self.hash_by_frame.iter(&rtxn)?;
2620            while let Some(Ok((k, v))) = iter.next() {
2621                lines.push(format!("hash_by_frame\t{}\t{}", hex(k), hex(v)));
2622            }
2623        }
2624        {
2625            let mut iter = self.frames_soa.iter(&rtxn)?;
2626            while let Some(Ok((k, v))) = iter.next() {
2627                lines.push(format!("frames_soa\t{}\t{}", hex(k), hex(v)));
2628            }
2629        }
2630        let unit_dbs: [(&str, Database<Bytes, Unit>); 11] = [
2631            ("idx_natoms", self.idx_natoms),
2632            ("idx_symbol", self.idx_symbol),
2633            ("idx_energy", self.idx_energy),
2634            ("idx_flags", self.idx_flags),
2635            ("idx_elem_count", self.idx_elem_count),
2636            ("idx_formula", self.idx_formula),
2637            ("idx_fmax", self.idx_fmax),
2638            ("idx_mass", self.idx_mass),
2639            ("idx_volume", self.idx_volume),
2640            ("idx_pbc", self.idx_pbc),
2641            ("idx_meta", self.idx_meta),
2642        ];
2643        for (name, db) in unit_dbs {
2644            let mut iter = db.iter(&rtxn)?;
2645            while let Some(Ok((k, ()))) = iter.next() {
2646                lines.push(format!("{name}\t{}", hex(k)));
2647            }
2648        }
2649        if let Some(db) = self.idx_topo {
2650            let mut iter = db.iter(&rtxn)?;
2651            while let Some(Ok((k, ()))) = iter.next() {
2652                lines.push(format!("idx_topo\t{}", hex(k)));
2653            }
2654        }
2655        if let Some(db) = self.topo_by_frame {
2656            let mut iter = db.iter(&rtxn)?;
2657            while let Some(Ok((k, v))) = iter.next() {
2658                lines.push(format!("topo_by_frame\t{}\t{}", hex(k), hex(v)));
2659            }
2660        }
2661        lines.sort();
2662        Ok(lines)
2663    }
2664
2665    pub fn unique_frame_keys(&self, sel: &Select) -> Result<Vec<FrameKey>> {
2666        let keys = self.select(sel)?;
2667        let mut uniq = Vec::new();
2668        for k in keys {
2669            let h = self.frame_hash(k)?;
2670            if self.find_by_hash(h)? == Some(k) {
2671                uniq.push(k);
2672            }
2673        }
2674        Ok(uniq)
2675    }
2676}
2677
2678#[cfg(test)]
2679mod tests {
2680    use super::*;
2681    use std::path::PathBuf;
2682
2683    fn fixture(name: &str) -> PathBuf {
2684        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
2685            .join("resources/test")
2686            .join(name)
2687    }
2688
2689    fn fixtures_dir() -> PathBuf {
2690        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("resources/test")
2691    }
2692
2693    #[test]
2694    fn parse_map_size_suffixes() {
2695        assert_eq!(parse_map_size("1048576"), Some(1 << 20));
2696        assert_eq!(parse_map_size("512M"), Some(512 << 20));
2697        assert_eq!(parse_map_size(" 8g "), Some(8usize << 30));
2698        assert_eq!(parse_map_size("2T"), Some(2usize << 40));
2699        assert_eq!(parse_map_size("0"), None);
2700        assert_eq!(parse_map_size("big"), None);
2701        assert_eq!(parse_map_size(""), None);
2702    }
2703
2704    #[test]
2705    fn map_size_env_overrides_default_and_request_overrides_env() {
2706        let dir = tempfile::tempdir().unwrap();
2707        // Larger than the default, so concurrent tests that open corpora are unaffected.
2708        std::env::set_var(MAP_SIZE_ENV, "3G");
2709        let from_env = initial_map_size(dir.path(), None).unwrap();
2710        let requested = initial_map_size(dir.path(), Some(5 << 20)).unwrap();
2711        std::env::remove_var(MAP_SIZE_ENV);
2712        assert_eq!(from_env, 3usize << 30);
2713        assert_eq!(requested, 5 << 20);
2714        assert_eq!(
2715            initial_map_size(dir.path(), None).unwrap(),
2716            DEFAULT_MAP_SIZE
2717        );
2718    }
2719
2720    #[test]
2721    fn small_map_fills_then_reserve_map_grows_it() {
2722        let dir = tempfile::tempdir().unwrap();
2723        let db = ConCorpus::open_with_map_size(dir.path(), 1 << 20).unwrap();
2724        assert_eq!(db.map_size(), 1 << 20);
2725        let grown = db.reserve_map(3 << 20).unwrap();
2726        assert!(grown >= 3 << 20, "{grown}");
2727        assert!(grown >= db.map_size());
2728    }
2729
2730    /// A 1 MiB map ingests 4 MiB through `append_trajectory_path` with no
2731    /// `reserve_map`. The write hits `MDB_MAP_FULL`, grows, and retries.
2732    #[test]
2733    fn serial_path_grows_1mib_map_for_4mib_input() {
2734        let one = std::fs::read_to_string(fixture("sulfolene.con")).unwrap();
2735        let text = one.repeat((4 << 20) / one.len() + 1);
2736        assert!(text.len() >= 4 << 20, "{}", text.len());
2737        let src = tempfile::tempdir().unwrap();
2738        let file = src.path().join("big.con");
2739        std::fs::write(&file, &text).unwrap();
2740        let dir = tempfile::tempdir().unwrap();
2741        let db = ConCorpus::open_with_map_size(dir.path(), 1 << 20).unwrap();
2742        assert_eq!(db.map_size(), 1 << 20);
2743        let n = db.append_trajectory_path(1, &file).unwrap();
2744        assert!(n > 1, "{n}");
2745        assert!(db.map_size() > 1 << 20, "{}", db.map_size());
2746        drop(db);
2747        let used = std::fs::metadata(dir.path().join("data.mdb"))
2748            .unwrap()
2749            .len() as usize;
2750        let db = ConCorpus::open_with_map_size(dir.path(), 1 << 20).unwrap();
2751        assert!(db.map_size() >= 2 * used, "{} vs {used}", db.map_size());
2752        assert_eq!(db.traj_meta(1).unwrap().unwrap().n_frames, n);
2753    }
2754
2755    #[test]
2756    fn ingest_dump_identical_at_1_2_3_8_threads() {
2757        let src = tempfile::tempdir().unwrap();
2758        for name in ["a.con", "b.convel", "c.con", "d.con"] {
2759            let from = if name.ends_with(".convel") {
2760                fixture("tiny_cuh2.convel")
2761            } else if name.starts_with('c') {
2762                fixture("tiny_cuh2_forces.con")
2763            } else if name.starts_with('d') {
2764                fixture("tiny_multi_cuh2.con")
2765            } else {
2766                fixture("tiny_cuh2.con")
2767            };
2768            std::fs::copy(from, src.path().join(name)).unwrap();
2769        }
2770        let mut dumps = Vec::new();
2771        for threads in [1usize, 2, 3, 8] {
2772            let dir = tempfile::tempdir().unwrap();
2773            let db = ConCorpus::open(dir.path()).unwrap();
2774            let rows = db
2775                .ingest_directory_units_nthreads(src.path(), 1, None, threads)
2776                .unwrap();
2777            assert_eq!(rows.len(), 4);
2778            let ids: Vec<u64> = rows.iter().map(|r| r.0).collect();
2779            assert_eq!(ids, vec![1, 2, 3, 4]);
2780            let dump = db.sorted_logical_dump().unwrap();
2781            assert!(dump.iter().any(|l| l.starts_with("frames\t")));
2782            assert!(dump.iter().any(|l| l.starts_with("traj_meta\t")));
2783            dumps.push(dump);
2784        }
2785        for dump in &dumps[1..] {
2786            assert_eq!(dump, &dumps[0]);
2787        }
2788    }
2789
2790    #[test]
2791    fn ingest_first_error_independent_of_thread_count() {
2792        let src = tempfile::tempdir().unwrap();
2793        std::fs::copy(fixture("tiny_cuh2.con"), src.path().join("a.con")).unwrap();
2794        std::fs::write(src.path().join("b.con"), "not a con frame\n").unwrap();
2795        std::fs::write(src.path().join("c.con"), "also not a frame\n").unwrap();
2796        let mut messages = Vec::new();
2797        for threads in [1usize, 2, 3, 8] {
2798            let dir = tempfile::tempdir().unwrap();
2799            let db = ConCorpus::open(dir.path()).unwrap();
2800            let err = db
2801                .ingest_directory_units_nthreads(src.path(), 1, None, threads)
2802                .unwrap_err();
2803            messages.push(err.to_string());
2804        }
2805        for msg in &messages[1..] {
2806            assert_eq!(msg, &messages[0], "{messages:?}");
2807        }
2808        assert!(!messages[0].is_empty(), "{messages:?}");
2809        let lower = messages[0].to_lowercase();
2810        assert!(
2811            lower.contains("parse")
2812                || lower.contains("invalid")
2813                || lower.contains("incomplete")
2814                || lower.contains("header"),
2815            "{messages:?}"
2816        );
2817    }
2818
2819    #[test]
2820    fn ingest_chunk_bounds_depend_on_file_count_only() {
2821        assert!(ingest_chunk_bounds(0).is_empty());
2822        assert_eq!(
2823            ingest_chunk_bounds(8),
2824            (0..8).map(|i| (i, i + 1)).collect::<Vec<_>>()
2825        );
2826        let sixteen = ingest_chunk_bounds(16);
2827        assert_eq!(sixteen.len(), 8);
2828        assert_eq!(sixteen[0], (0, 2));
2829        assert_eq!(sixteen[7], (14, 16));
2830        assert_eq!(ingest_chunk_bounds(10), ingest_chunk_bounds(10));
2831    }
2832
2833    #[test]
2834    fn ingest_hash_dedup_and_select() {
2835        let dir = tempfile::tempdir().unwrap();
2836        let db = ConCorpus::open(dir.path()).unwrap();
2837        let n = db
2838            .append_trajectory_path(1, fixture("tiny_cuh2.con"))
2839            .unwrap();
2840        assert!(n >= 1);
2841        let k0 = FrameKey {
2842            traj_id: 1,
2843            frame_idx: 0,
2844        };
2845        let h = db.frame_hash(k0).unwrap();
2846        assert_eq!(db.find_by_hash(h).unwrap(), Some(k0));
2847
2848        let by_hash = db.select(&Select::new().exact_hash(h.to_bytes())).unwrap();
2849        assert_eq!(by_hash, vec![k0]);
2850
2851        let cu = db
2852            .select(&Select::new().require_symbol("Cu").trajectory(1))
2853            .unwrap();
2854        assert!(!cu.is_empty());
2855
2856        let n2 = db
2857            .append_trajectory_path(2, fixture("tiny_multi_cuh2.con"))
2858            .unwrap();
2859        assert!(n2 >= 2);
2860    }
2861
2862    #[test]
2863    fn get_velocities_returns_native_convel() {
2864        let dir = tempfile::tempdir().unwrap();
2865        let db = ConCorpus::open(dir.path()).unwrap();
2866        db.append_trajectory_path(1, fixture("tiny_cuh2.convel"))
2867            .unwrap();
2868        let v = db
2869            .get_velocities(FrameKey {
2870                traj_id: 1,
2871                frame_idx: 0,
2872            })
2873            .unwrap()
2874            .expect("velocities");
2875        assert!((v[0][0] - 0.001234).abs() < 1e-12, "got {}", v[0][0]);
2876        assert!(db
2877            .get_velocities(FrameKey {
2878                traj_id: 1,
2879                frame_idx: 0,
2880            })
2881            .unwrap()
2882            .is_some());
2883        db.append_trajectory_path(2, fixture("tiny_cuh2.con"))
2884            .unwrap();
2885        assert!(db
2886            .get_velocities(FrameKey {
2887                traj_id: 2,
2888                frame_idx: 0,
2889            })
2890            .unwrap()
2891            .is_none());
2892    }
2893
2894    #[test]
2895    fn caller_units_canonical_in_con_metadata() {
2896        let dir = tempfile::tempdir().unwrap();
2897        let db = ConCorpus::open(dir.path()).unwrap();
2898        db.append_trajectory_path_units(
2899            1,
2900            fixture("tiny_cuh2.con"),
2901            Some(serde_json::json!({
2902                "length": "A",
2903                "energy": "ev",
2904                "time": "femtosecond"
2905            })),
2906        )
2907        .unwrap();
2908        let k = FrameKey {
2909            traj_id: 1,
2910            frame_idx: 0,
2911        };
2912        let u = db.frame_units(k).unwrap().expect("units");
2913        assert_eq!(u["length"], "angstrom");
2914        assert_eq!(u["energy"], "eV");
2915        assert_eq!(u["time"], "fs");
2916        let text = db.get_frame_text(k).unwrap();
2917        assert!(text.contains("angstrom"), "{text}");
2918        assert!(text.contains("fs"), "{text}");
2919        db.set_trajectory_units(
2920            1,
2921            serde_json::json!({"length": "nm", "energy": "hartree", "time": "ps"}),
2922        )
2923        .unwrap();
2924        let u2 = db.frame_units(k).unwrap().expect("units");
2925        assert_eq!(u2["length"], "nm");
2926        assert_eq!(u2["energy"], "hartree");
2927        assert_eq!(u2["time"], "ps");
2928        let fr0 = db.get_frame(k).unwrap();
2929        let vol_nm = crate::frame_scalars::frame_cell_volume(&fr0).expect("vol");
2930        let old_ang_vol = vol_nm * 1000.0;
2931        let stale = db
2932            .select(&Select::new().volume_range(old_ang_vol * 0.5, old_ang_vol * 1.5))
2933            .unwrap();
2934        assert!(
2935            stale.is_empty(),
2936            "old angstrom volume key must be gone after set_units"
2937        );
2938        let fresh = db
2939            .select(&Select::new().volume_range(vol_nm * 0.5, vol_nm * 1.5))
2940            .unwrap();
2941        assert_eq!(fresh, vec![k]);
2942    }
2943
2944    #[test]
2945    fn ingest_directory_units_stamps_canonical() {
2946        let dir = tempfile::tempdir().unwrap();
2947        let src = dir.path().join("cons");
2948        std::fs::create_dir(&src).unwrap();
2949        std::fs::copy(fixture("tiny_cuh2.con"), src.join("a.con")).unwrap();
2950        let db = ConCorpus::open(dir.path().join("db")).unwrap();
2951        let n = db
2952            .ingest_directory_units(
2953                &src,
2954                1,
2955                Some(serde_json::json!({"length":"A","energy":"ev"})),
2956            )
2957            .unwrap();
2958        assert_eq!(n.len(), 1);
2959        let u = db
2960            .frame_units(FrameKey {
2961                traj_id: 1,
2962                frame_idx: 0,
2963            })
2964            .unwrap()
2965            .unwrap();
2966        assert_eq!(u["length"], "angstrom");
2967        assert_eq!(u["energy"], "eV");
2968    }
2969
2970    #[test]
2971    fn readonly_open_and_pack_roundtrip() {
2972        let dir = tempfile::tempdir().unwrap();
2973        let key = FrameKey {
2974            traj_id: 1,
2975            frame_idx: 0,
2976        };
2977        let key1 = FrameKey {
2978            traj_id: 1,
2979            frame_idx: 1,
2980        };
2981        let packed = {
2982            let db = ConCorpus::open(dir.path()).unwrap();
2983            db.append_trajectory_path(1, fixture("tiny_multi_cuh2.con"))
2984                .unwrap();
2985            let packed = db.pack_frame(key).unwrap();
2986            db.close();
2987            packed
2988        };
2989        let decoded = crate::cooked_soa::CookedSoa::decode(&packed).unwrap();
2990        assert!(decoded.natoms >= 1);
2991
2992        let ro = ConCorpus::open_readonly(dir.path()).unwrap();
2993        let packed2 = ro.pack_frame(key).unwrap();
2994        assert_eq!(packed, packed2);
2995        let batch = ro.pack_frames(&[key, key1]).unwrap();
2996        let parts = crate::cooked_soa::decode_batch(&batch).unwrap();
2997        assert_eq!(parts.len(), 2);
2998        assert_eq!(parts[0], packed);
2999        let d1 = crate::cooked_soa::CookedSoa::decode(&parts[1]).unwrap();
3000        assert!((d1.positions[2][0] - 8.8549).abs() < 1e-4);
3001        assert!(ConCorpus::open_readonly("/no/such/corpus-dir-readonly").is_err());
3002        assert!(ro
3003            .append_trajectory_path(9, fixture("tiny_cuh2.con"))
3004            .is_err());
3005    }
3006
3007    #[test]
3008    fn workflow_metatrain_extxyz_export() {
3009        let dir = tempfile::tempdir().unwrap();
3010        let db = ConCorpus::open(dir.path().join("corpus")).unwrap();
3011        let ingested = db.ingest_directory(fixtures_dir(), 1).unwrap();
3012        assert!(ingested.len() >= 3);
3013
3014        let keys = db
3015            .select(
3016                &Select::new()
3017                    .require_symbol("Cu")
3018                    .natoms_range(1, 10_000)
3019                    .limit(50),
3020            )
3021            .unwrap();
3022        assert!(!keys.is_empty());
3023
3024        let uniq = db
3025            .unique_frame_keys(&Select::new().require_symbol("Cu"))
3026            .unwrap();
3027        assert!(!uniq.is_empty());
3028
3029        let xyz = dir.path().join("train_subset.xyz");
3030        let n = db.export_extxyz(&uniq, &xyz, "energy").unwrap();
3031        assert_eq!(n, uniq.len());
3032        let text = std::fs::read_to_string(&xyz).unwrap();
3033        assert!(text.contains("Lattice="));
3034        assert!(text.contains("Properties="));
3035        assert!(text.lines().any(|l| l.trim_start().starts_with("Cu ")));
3036        let err = db.export_extxyz(&uniq, &xyz, "energy").unwrap_err();
3037        assert!(err.to_string().contains("dest exists"), "{err}");
3038        let missing = dir.path().join("missing_key.xyz");
3039        let bad = db.export_extxyz(
3040            &[FrameKey {
3041                traj_id: 99,
3042                frame_idx: 0,
3043            }],
3044            &missing,
3045            "energy",
3046        );
3047        assert!(bad.is_err());
3048        assert!(!missing.exists());
3049    }
3050
3051    #[test]
3052    fn open_refuses_sharded_root() {
3053        let dir = tempfile::tempdir().unwrap();
3054        let root = dir.path().join("sharded");
3055        crate::ShardedConCorpus::open(&root, 2).unwrap();
3056        assert!(root.join("shards.json").is_file());
3057        match ConCorpus::open(&root) {
3058            Ok(_) => panic!("expected sharded-root refuse"),
3059            Err(err) => assert!(err.to_string().contains("sharded root"), "{err}"),
3060        }
3061        assert!(!root.join("data.mdb").exists());
3062    }
3063
3064    #[test]
3065    fn workflow_dedup_identical_ingest() {
3066        let dir = tempfile::tempdir().unwrap();
3067        let db = ConCorpus::open(dir.path()).unwrap();
3068        let f = fixture("tiny_cuh2.con");
3069        db.append_trajectory_path(1, &f).unwrap();
3070        db.append_trajectory_path(2, &f).unwrap();
3071        let k1 = FrameKey {
3072            traj_id: 1,
3073            frame_idx: 0,
3074        };
3075        let k2 = FrameKey {
3076            traj_id: 2,
3077            frame_idx: 0,
3078        };
3079        let h1 = db.frame_hash(k1).unwrap();
3080        let h2 = db.frame_hash(k2).unwrap();
3081        assert_eq!(h1, h2);
3082        assert_eq!(db.find_by_hash(h1).unwrap(), Some(k1));
3083        let uniq = db.unique_frame_keys(&Select::new()).unwrap();
3084        assert!(uniq.contains(&k1));
3085        assert!(!uniq.contains(&k2));
3086    }
3087
3088    #[test]
3089    fn metadata_indexes_forces_velocities_energy() {
3090        let dir = tempfile::tempdir().unwrap();
3091        let db = ConCorpus::open(dir.path()).unwrap();
3092        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3093            .unwrap();
3094        db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
3095            .unwrap();
3096        db.append_trajectory_path(3, fixture("tiny_cuh2.convel"))
3097            .unwrap();
3098        db.append_trajectory_path(4, fixture("tiny_cuh2_vel_forces.con"))
3099            .unwrap();
3100
3101        let forces = db.select(&Select::new().require_forces()).unwrap();
3102        assert!(forces.iter().any(|k| k.traj_id == 2));
3103        assert!(forces.iter().any(|k| k.traj_id == 4));
3104        assert!(!forces.iter().any(|k| k.traj_id == 1));
3105
3106        let vels = db.select(&Select::new().require_velocities()).unwrap();
3107        assert!(vels.iter().any(|k| k.traj_id == 3));
3108        assert!(vels.iter().any(|k| k.traj_id == 4));
3109
3110        let both = db
3111            .select(&Select::new().require_forces().require_velocities())
3112            .unwrap();
3113        assert!(both.iter().any(|k| k.traj_id == 4));
3114        assert!(!both.iter().any(|k| k.traj_id == 2));
3115
3116        let with_e = db.select(&Select::new().require_energy()).unwrap();
3117        assert!(with_e.iter().any(|k| k.traj_id == 2));
3118        assert!(with_e.iter().any(|k| k.traj_id == 4));
3119        assert!(!with_e.iter().any(|k| k.traj_id == 1));
3120
3121        let in_range = db
3122            .select(&Select::new().energy_range(-43.0, -42.0).require_forces())
3123            .unwrap();
3124        assert!(in_range.iter().any(|k| k.traj_id == 2));
3125        let miss = db.select(&Select::new().energy_range(0.0, 1.0)).unwrap();
3126        assert!(miss.is_empty());
3127    }
3128
3129    #[test]
3130    fn composition_and_fmax_indexes() {
3131        let dir = tempfile::tempdir().unwrap();
3132        let db = ConCorpus::open(dir.path()).unwrap();
3133        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3134            .unwrap();
3135        db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
3136            .unwrap();
3137        db.append_trajectory_path(3, fixture("sulfolene.con"))
3138            .unwrap();
3139
3140        let formula = db
3141            .frame_formula(FrameKey {
3142                traj_id: 1,
3143                frame_idx: 0,
3144            })
3145            .unwrap();
3146        assert_eq!(formula, "Cu:2|H:2");
3147
3148        let exact = db
3149            .select(&Select::new().exact_composition("Cu:2|H:2"))
3150            .unwrap();
3151        assert!(exact.iter().any(|k| k.traj_id == 1));
3152        assert!(exact.iter().any(|k| k.traj_id == 2));
3153        assert!(!exact.iter().any(|k| k.traj_id == 3));
3154
3155        let cu2 = db
3156            .select(&Select::new().element_exact("Cu", 2).element_exact("H", 2))
3157            .unwrap();
3158        assert_eq!(cu2.len(), exact.len());
3159
3160        let cu_min2 = db.select(&Select::new().element_min("Cu", 2)).unwrap();
3161        assert!(cu_min2.iter().any(|k| k.traj_id == 1));
3162
3163        let wrong = db.select(&Select::new().exact_composition("Fe:1")).unwrap();
3164        assert!(wrong.is_empty());
3165
3166        let fr_f = db
3167            .get_frame(FrameKey {
3168                traj_id: 2,
3169                frame_idx: 0,
3170            })
3171            .unwrap();
3172        let expected_fmax = frame_fmax(&fr_f).expect("forces fixture has fmax");
3173        let in_fmax = db
3174            .select(
3175                &Select::new()
3176                    .require_forces()
3177                    .fmax_range(0.0, expected_fmax + 1e-6),
3178            )
3179            .unwrap();
3180        assert!(in_fmax.iter().any(|k| k.traj_id == 2));
3181        assert!(!in_fmax.iter().any(|k| k.traj_id == 1));
3182
3183        let too_small = db
3184            .select(&Select::new().fmax_range(0.0, expected_fmax * 0.0 - 1.0))
3185            .unwrap();
3186        // negative max excludes all positive fmax postings that are > hi; lo>hi may be empty
3187        let impossible = db
3188            .select(&Select::new().fmax_range(1e9, 1e9 + 1.0))
3189            .unwrap();
3190        assert!(impossible.is_empty());
3191        let _ = too_small;
3192    }
3193
3194    #[test]
3195    fn reindex_and_append_frames() {
3196        let dir = tempfile::tempdir().unwrap();
3197        let db = ConCorpus::open(dir.path()).unwrap();
3198        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3199            .unwrap();
3200        db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
3201            .unwrap();
3202
3203        let before = db
3204            .select(&Select::new().exact_composition("Cu:2|H:2"))
3205            .unwrap();
3206        assert!(!before.is_empty());
3207
3208        let n = db.reindex().unwrap();
3209        assert!(n >= 2);
3210        let after = db
3211            .select(&Select::new().exact_composition("Cu:2|H:2"))
3212            .unwrap();
3213        assert_eq!(before, after);
3214
3215        let fr = db
3216            .get_frame(FrameKey {
3217                traj_id: 1,
3218                frame_idx: 0,
3219            })
3220            .unwrap();
3221        let n2 = db
3222            .append_trajectory_frames(10, &[fr.clone()], "mem")
3223            .unwrap();
3224        assert_eq!(n2, 1);
3225        let mem = db
3226            .select(&Select::new().trajectory(10).exact_composition("Cu:2|H:2"))
3227            .unwrap();
3228        assert_eq!(mem.len(), 1);
3229
3230        let n3 = db.extend_trajectory_frames(10, &[fr], "mem-ext").unwrap();
3231        assert_eq!(n3, 2);
3232    }
3233
3234    #[test]
3235    fn extend_trajectory_path_creates_then_appends() {
3236        let dir = tempfile::tempdir().unwrap();
3237        let db = ConCorpus::open(dir.path()).unwrap();
3238        let path = fixture("tiny_cuh2.con");
3239        let n1 = db.extend_trajectory_path(3, &path).unwrap();
3240        assert!(n1 >= 1);
3241        let n2 = db.extend_trajectory_path(3, &path).unwrap();
3242        assert_eq!(n2, n1 * 2);
3243    }
3244
3245    #[test]
3246    fn reindex_twice_stable() {
3247        let dir = tempfile::tempdir().unwrap();
3248        let db = ConCorpus::open(dir.path()).unwrap();
3249        db.append_trajectory_path(1, fixture("tiny_cuh2_forces.con"))
3250            .unwrap();
3251        let a = db.reindex().unwrap();
3252        let keys1 = db.select(&Select::new().require_forces()).unwrap();
3253        let b = db.reindex().unwrap();
3254        let keys2 = db.select(&Select::new().require_forces()).unwrap();
3255        assert_eq!(a, b);
3256        assert_eq!(keys1, keys2);
3257    }
3258
3259    #[test]
3260    fn batch_touch_and_texts() {
3261        let dir = tempfile::tempdir().unwrap();
3262        let db = ConCorpus::open(dir.path()).unwrap();
3263        let n = db
3264            .append_trajectory_path(1, fixture("tiny_cuh2.con"))
3265            .unwrap();
3266        assert!(n >= 1);
3267        let (total, checksum) = db.touch_trajectory_blobs(1, n).unwrap();
3268        assert!(total > 0);
3269        assert!(checksum > 0);
3270        let texts = db
3271            .get_frame_texts(&[FrameKey {
3272                traj_id: 1,
3273                frame_idx: 0,
3274            }])
3275            .unwrap();
3276        assert_eq!(texts.len(), 1);
3277        assert_eq!(texts[0].len() as u64, total);
3278        // Checksum must depend on payload bytes (not only length).
3279        let mut expect = 0u64;
3280        for (i, b) in texts[0].as_bytes().iter().enumerate() {
3281            expect = expect.wrapping_add((*b as u64).wrapping_mul((i as u64).wrapping_add(1)));
3282        }
3283        assert_eq!(checksum, expect);
3284    }
3285
3286    /// Concurrent threads append distinct traj IDs; prepare runs in parallel, LMDB commits serialize only.
3287    #[test]
3288    fn concurrent_writers_distinct_traj() {
3289        use std::sync::Arc;
3290        use std::thread;
3291        let dir = tempfile::tempdir().unwrap();
3292        let db = Arc::new(ConCorpus::open(dir.path()).unwrap());
3293        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
3294        let mut handles = Vec::new();
3295        for tid in 1u64..=8 {
3296            let db = Arc::clone(&db);
3297            let text = text.clone();
3298            handles.push(thread::spawn(move || {
3299                db.append_trajectory_str(tid, &text, format!("t{tid}"))
3300                    .unwrap()
3301            }));
3302        }
3303        let mut counts = Vec::new();
3304        for h in handles {
3305            counts.push(h.join().unwrap());
3306        }
3307        assert!(counts.iter().all(|&c| c >= 1));
3308        for tid in 1u64..=8 {
3309            let meta = db.traj_meta(tid).unwrap().expect("meta");
3310            assert_eq!(meta.n_frames, counts[(tid - 1) as usize]);
3311            let _ = db
3312                .get_frame_text(FrameKey {
3313                    traj_id: tid,
3314                    frame_idx: 0,
3315                })
3316                .unwrap();
3317        }
3318        let cu = db.select(&Select::new().require_symbol("Cu")).unwrap();
3319        assert_eq!(cu.len(), 8);
3320    }
3321
3322    /// Multi-process: N OS processes each open the same env via CLI and select.
3323    #[test]
3324    fn multiproc_cli_concurrent_select() {
3325        use std::process::Command;
3326        let dir = tempfile::tempdir().unwrap();
3327        let corpus = dir.path().join("mp_corpus");
3328        {
3329            let db = ConCorpus::open(&corpus).unwrap();
3330            db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3331                .unwrap();
3332        }
3333        let Some(bin) = option_env!("CARGO_BIN_EXE_readcon-db") else {
3334            return;
3335        };
3336        let corpus_s = corpus.to_str().unwrap();
3337        let mut kids = Vec::new();
3338        for _ in 0..4 {
3339            kids.push(
3340                Command::new(bin)
3341                    .args(["select", corpus_s, "--symbol", "Cu"])
3342                    .stdout(std::process::Stdio::null())
3343                    .spawn()
3344                    .expect("spawn"),
3345            );
3346        }
3347        for mut c in kids {
3348            assert!(c.wait().unwrap().success());
3349        }
3350    }
3351
3352    #[test]
3353    fn ase_competitive_mass_volume_pbc_meta_charge() {
3354        use readcon_core::types::ConFrameBuilder;
3355        let dir = tempfile::tempdir().unwrap();
3356        let db = ConCorpus::open(dir.path()).unwrap();
3357
3358        // Frame A: Cu2H2 default cell from fixture path
3359        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3360            .unwrap();
3361        let fr1 = db
3362            .get_frame(FrameKey {
3363                traj_id: 1,
3364                frame_idx: 0,
3365            })
3366            .unwrap();
3367        let mass1 = frame_total_mass(&fr1).expect("mass");
3368        let vol1 = frame_cell_volume(&fr1).expect("vol");
3369
3370        // Frame B: inject metadata via rewrite
3371        let mut fr2 = fr1.clone();
3372        fr2.header.set_frame_index(7);
3373        fr2.header.set_time(1.5);
3374        fr2.header.set_timestep(0.25);
3375        fr2.header.set_neb_bead(2);
3376        fr2.header.set_neb_band(1);
3377        fr2.header.set_pbc([true, true, false]);
3378        fr2.header
3379            .metadata
3380            .insert("charge".into(), serde_json::json!(-1.0));
3381        fr2.header
3382            .metadata
3383            .insert("magmom".into(), serde_json::json!(2.0));
3384        // larger cell → larger volume
3385        fr2.header.boxl = [20.0, 20.0, 20.0];
3386        db.append_trajectory_frames(2, &[fr2], "meta").unwrap();
3387
3388        let by_mass = db
3389            .select(&Select::new().mass_range(mass1 - 1e-6, mass1 + 1e-6))
3390            .unwrap();
3391        assert!(by_mass.iter().any(|k| k.traj_id == 1));
3392
3393        let by_vol_small = db
3394            .select(&Select::new().volume_range(vol1 - 1.0, vol1 + 1.0))
3395            .unwrap();
3396        assert!(by_vol_small.iter().any(|k| k.traj_id == 1));
3397        assert!(!by_vol_small.iter().any(|k| k.traj_id == 2));
3398
3399        let pbc_match = db.select(&Select::new().pbc([true, true, false])).unwrap();
3400        assert_eq!(pbc_match.len(), 1);
3401        assert_eq!(pbc_match[0].traj_id, 2);
3402        let pbc_miss = db.select(&Select::new().pbc([true, true, true])).unwrap();
3403        assert!(pbc_miss.is_empty());
3404
3405        let fi = db
3406            .select(&Select::new().frame_index_range(7.0, 7.0))
3407            .unwrap();
3408        assert_eq!(
3409            fi,
3410            vec![FrameKey {
3411                traj_id: 2,
3412                frame_idx: 0
3413            }]
3414        );
3415        let fi_miss = db
3416            .select(&Select::new().frame_index_range(99.0, 100.0))
3417            .unwrap();
3418        assert!(fi_miss.is_empty());
3419
3420        let bead = db.select(&Select::new().neb_bead_range(2.0, 2.0)).unwrap();
3421        assert_eq!(bead.len(), 1);
3422        let t = db.select(&Select::new().time_range(1.4, 1.6)).unwrap();
3423        assert_eq!(t.len(), 1);
3424
3425        let dt = db
3426            .select(&Select::new().timestep_range(0.25, 0.25))
3427            .unwrap();
3428        assert_eq!(
3429            dt,
3430            vec![FrameKey {
3431                traj_id: 2,
3432                frame_idx: 0
3433            }]
3434        );
3435        let dt_miss = db.select(&Select::new().timestep_range(9.0, 10.0)).unwrap();
3436        assert!(dt_miss.is_empty());
3437        // traj 1 has no timestep metadata → must not satisfy a finite timestep window
3438        assert!(!dt.iter().any(|k| k.traj_id == 1));
3439
3440        let band = db.select(&Select::new().neb_band_range(1.0, 1.0)).unwrap();
3441        assert_eq!(
3442            band,
3443            vec![FrameKey {
3444                traj_id: 2,
3445                frame_idx: 0
3446            }]
3447        );
3448        let band_miss = db
3449            .select(&Select::new().neb_band_range(99.0, 100.0))
3450            .unwrap();
3451        assert!(band_miss.is_empty());
3452
3453        let ch = db.select(&Select::new().charge_range(-1.0, -1.0)).unwrap();
3454        assert_eq!(ch.len(), 1);
3455        let mm = db.select(&Select::new().magmom_range(2.0, 2.0)).unwrap();
3456        assert_eq!(mm.len(), 1);
3457        let ch_miss = db.select(&Select::new().charge_range(0.0, 0.0)).unwrap();
3458        assert!(ch_miss.is_empty());
3459
3460        let n = db.reindex().unwrap();
3461        assert!(n >= 2);
3462        let after = db
3463            .select(&Select::new().frame_index_range(7.0, 7.0))
3464            .unwrap();
3465        assert_eq!(after.len(), 1);
3466        let after_dt = db
3467            .select(&Select::new().timestep_range(0.25, 0.25))
3468            .unwrap();
3469        assert_eq!(after_dt.len(), 1);
3470        let after_band = db.select(&Select::new().neb_band_range(1.0, 1.0)).unwrap();
3471        assert_eq!(after_band.len(), 1);
3472
3473        // still composition
3474        let form = db
3475            .select(&Select::new().exact_composition("Cu:2|H:2"))
3476            .unwrap();
3477        assert!(form.len() >= 2);
3478    }
3479
3480    #[test]
3481    fn cooked_soa_roundtrip_and_authority() {
3482        let dir = tempfile::tempdir().unwrap();
3483        let db = ConCorpus::open(dir.path()).unwrap();
3484        let key = FrameKey {
3485            traj_id: 1,
3486            frame_idx: 0,
3487        };
3488        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3489            .unwrap();
3490        // default: no cooked
3491        assert!(db.get_cooked_soa_bytes(key).unwrap().is_none());
3492        let pos_via_parse = db.get_positions(key).unwrap();
3493        let fr = db.get_frame(key).unwrap();
3494        assert_eq!(pos_via_parse.len(), fr.atom_data.len());
3495        for (i, a) in fr.atom_data.iter().enumerate() {
3496            assert_eq!(pos_via_parse[i], [a.x, a.y, a.z]);
3497        }
3498
3499        let text_before = db.get_frame_text(key).unwrap();
3500        let h_before = db.frame_hash(key).unwrap();
3501        let nbytes = db.cook_frame(key).unwrap();
3502        assert!(nbytes > 0);
3503        let cooked = db.get_cooked_soa(key).unwrap().expect("cooked present");
3504        assert_eq!(cooked.natoms as usize, fr.atom_data.len());
3505        for (i, a) in fr.atom_data.iter().enumerate() {
3506            assert_eq!(cooked.positions[i], [a.x, a.y, a.z]);
3507        }
3508        // CON authority unchanged
3509        assert_eq!(db.get_frame_text(key).unwrap(), text_before);
3510        assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h_before.to_bytes());
3511        assert_eq!(db.find_by_hash(h_before).unwrap(), Some(key));
3512        let pos_cooked = db.get_positions(key).unwrap();
3513        assert_eq!(pos_cooked, cooked.positions);
3514
3515        // select still works
3516        let keys = db.select(&Select::new().require_symbol("Cu")).unwrap();
3517        assert!(keys.contains(&key));
3518
3519        // reindex ignores cooked; text still selects
3520        db.reindex().unwrap();
3521        assert!(db.get_cooked_soa(key).unwrap().is_some());
3522        assert!(db
3523            .select(&Select::new().require_symbol("Cu"))
3524            .unwrap()
3525            .contains(&key));
3526
3527        // delete cooked → fallback parse
3528        db.delete_cooked_soa(key).unwrap();
3529        assert!(db.get_cooked_soa_bytes(key).unwrap().is_none());
3530        let pos_again = db.get_positions(key).unwrap();
3531        assert_eq!(pos_again, pos_via_parse);
3532        assert_eq!(db.get_frame_text(key).unwrap(), text_before);
3533    }
3534
3535    #[test]
3536    fn cooked_soa_forces_and_cook_on_ingest() {
3537        let dir = tempfile::tempdir().unwrap();
3538        let db = ConCorpus::open(dir.path()).unwrap();
3539        db.append_trajectory_path_cook(3, fixture("tiny_cuh2_forces.con"), true)
3540            .unwrap();
3541        let key = FrameKey {
3542            traj_id: 3,
3543            frame_idx: 0,
3544        };
3545        assert!(db.get_cooked_soa_bytes(key).unwrap().is_some());
3546        let fr = db.get_frame(key).unwrap();
3547        let forces = db.get_forces(key).unwrap().expect("forces");
3548        for (i, a) in fr.atom_data.iter().enumerate() {
3549            if let Some(f) = a.force() {
3550                assert_eq!(forces[i], f);
3551            }
3552        }
3553        // corrupt cooked → fallback
3554        {
3555            let _gate = db.write_guard();
3556            let mut wtxn = db.env.write_txn().unwrap();
3557            let fk_b = key.to_bytes();
3558            db.frames_soa
3559                .put(&mut wtxn, &fk_b[..], b"XXXX_not_valid")
3560                .unwrap();
3561            wtxn.commit().unwrap();
3562        }
3563        assert!(db.get_cooked_soa(key).unwrap().is_none());
3564        let pos = db.get_positions(key).unwrap();
3565        assert_eq!(pos.len(), fr.atom_data.len());
3566        assert_eq!(
3567            pos[0],
3568            [fr.atom_data[0].x, fr.atom_data[0].y, fr.atom_data[0].z]
3569        );
3570    }
3571
3572    #[test]
3573    fn cooked_soa_recook_all() {
3574        let dir = tempfile::tempdir().unwrap();
3575        let db = ConCorpus::open(dir.path()).unwrap();
3576        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3577            .unwrap();
3578        db.append_trajectory_path(2, fixture("tiny_cuh2_forces.con"))
3579            .unwrap();
3580        let n = db.recook_all().unwrap();
3581        assert_eq!(n, 2);
3582        assert!(db
3583            .get_cooked_soa(FrameKey {
3584                traj_id: 1,
3585                frame_idx: 0
3586            })
3587            .unwrap()
3588            .is_some());
3589        assert!(db
3590            .get_cooked_soa(FrameKey {
3591                traj_id: 2,
3592                frame_idx: 0
3593            })
3594            .unwrap()
3595            .is_some());
3596    }
3597
3598    /// RCSO cannot stand in for missing CON text (not fully equivalent).
3599    #[test]
3600    fn cooked_soa_cannot_replace_missing_frames_text() {
3601        let dir = tempfile::tempdir().unwrap();
3602        let db = ConCorpus::open(dir.path()).unwrap();
3603        let key = FrameKey {
3604            traj_id: 9,
3605            frame_idx: 0,
3606        };
3607        db.append_trajectory_path_cook(9, fixture("tiny_cuh2.con"), true)
3608            .unwrap();
3609        assert!(db.has_valid_cooked_soa(key).unwrap());
3610        let pos_ok = db.get_positions(key).unwrap();
3611        assert!(!pos_ok.is_empty());
3612        // Delete authoritative CON blob but leave RCSO (unsupported storage shape).
3613        {
3614            let _gate = db.write_guard();
3615            let mut wtxn = db.env.write_txn().unwrap();
3616            let fk_b = key.to_bytes();
3617            db.frames.delete(&mut wtxn, &fk_b[..]).unwrap();
3618            wtxn.commit().unwrap();
3619        }
3620        // Numeric tier may still decode RCSO...
3621        assert!(db.get_cooked_soa(key).unwrap().is_some());
3622        let pos_cooked_only = db.get_positions(key).unwrap();
3623        assert_eq!(pos_cooked_only, pos_ok);
3624        // ...but authority surfaces fail without CON text.
3625        assert!(db.get_frame_text(key).is_err());
3626        assert!(db.get_frame(key).is_err());
3627        // Hash reverse lookup still keyed off CON-era hash_by_frame may exist, but
3628        // reindex cannot rebuild without frames blob.
3629        assert!(db.frame_formula(key).is_err());
3630    }
3631
3632    #[test]
3633    fn numeric_fast_path_matches_con_twice() {
3634        let dir = tempfile::tempdir().unwrap();
3635        let db = ConCorpus::open(dir.path()).unwrap();
3636        let key = FrameKey {
3637            traj_id: 1,
3638            frame_idx: 0,
3639        };
3640        db.append_trajectory_path(1, fixture("tiny_cuh2_forces.con"))
3641            .unwrap();
3642        let fr = db.get_frame(key).unwrap();
3643        let expect_pos: Vec<_> = fr.atom_data.iter().map(|a| [a.x, a.y, a.z]).collect();
3644        let expect_f: Vec<_> = fr
3645            .atom_data
3646            .iter()
3647            .map(|a| a.force().unwrap_or([0.0; 3]))
3648            .collect();
3649        // Without cook: parse path
3650        assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3651        assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
3652        db.cook_frame(key).unwrap();
3653        assert!(db.has_valid_cooked_soa(key).unwrap());
3654        // With cook: fast path must match CON-derived atoms (shipped API, not re-impl)
3655        assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3656        assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
3657        let text = db.get_frame_text(key).unwrap();
3658        let h = db.frame_hash(key).unwrap();
3659        db.delete_cooked_soa(key).unwrap();
3660        assert_eq!(db.get_frame_text(key).unwrap(), text);
3661        assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h.to_bytes());
3662        assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3663    }
3664
3665    /// Positions, forces, and velocities: CON parse and cooked extract must match.
3666    #[test]
3667    fn numeric_extract_velocities_match_con() {
3668        let dir = tempfile::tempdir().unwrap();
3669        let db = ConCorpus::open(dir.path()).unwrap();
3670        let key = FrameKey {
3671            traj_id: 1,
3672            frame_idx: 0,
3673        };
3674        db.append_trajectory_path(1, fixture("tiny_cuh2_vel_forces.con"))
3675            .unwrap();
3676        assert!(!db.has_valid_cooked_soa(key).unwrap());
3677        let fr = db.get_frame(key).unwrap();
3678        let expect_pos: Vec<_> = fr.atom_data.iter().map(|a| [a.x, a.y, a.z]).collect();
3679        let expect_f: Vec<_> = fr
3680            .atom_data
3681            .iter()
3682            .map(|a| a.force().unwrap_or([0.0; 3]))
3683            .collect();
3684        let expect_v: Vec<_> = fr
3685            .atom_data
3686            .iter()
3687            .map(|a| a.velocity().unwrap_or([0.0; 3]))
3688            .collect();
3689        assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3690        assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
3691        assert_eq!(db.get_velocities(key).unwrap().unwrap(), expect_v);
3692        assert!((expect_v[0][0] - 0.001234).abs() < 1e-12);
3693        let text = db.get_frame_text(key).unwrap();
3694        let h = db.frame_hash(key).unwrap();
3695        db.cook_frame(key).unwrap();
3696        assert!(db.has_valid_cooked_soa(key).unwrap());
3697        assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3698        assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
3699        assert_eq!(db.get_velocities(key).unwrap().unwrap(), expect_v);
3700        assert_eq!(db.get_frame_text(key).unwrap(), text);
3701        assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h.to_bytes());
3702        assert_eq!(db.find_by_hash(h).unwrap(), Some(key));
3703        db.delete_cooked_soa(key).unwrap();
3704        assert!(!db.has_valid_cooked_soa(key).unwrap());
3705        assert_eq!(db.get_positions(key).unwrap(), expect_pos);
3706        assert_eq!(db.get_forces(key).unwrap().unwrap(), expect_f);
3707        assert_eq!(db.get_velocities(key).unwrap().unwrap(), expect_v);
3708        assert_eq!(db.get_frame_text(key).unwrap(), text);
3709    }
3710
3711    /// Valid disagreeing RCSO must win on getters; CON text and hash stay.
3712    #[test]
3713    fn numeric_extract_valid_rcso_skips_con_parse() {
3714        let dir = tempfile::tempdir().unwrap();
3715        let db = ConCorpus::open(dir.path()).unwrap();
3716        let key = FrameKey {
3717            traj_id: 1,
3718            frame_idx: 0,
3719        };
3720        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3721            .unwrap();
3722        let text = db.get_frame_text(key).unwrap();
3723        let h = db.frame_hash(key).unwrap();
3724        let mut fr = db.get_frame(key).unwrap();
3725        let con_x = fr.atom_data[0].x;
3726        fr.atom_data[0].x = 99.0;
3727        let poison = crate::cooked_soa::CookedSoa::encode_frame(&fr).unwrap();
3728        {
3729            let _gate = db.write_guard();
3730            let mut wtxn = db.env.write_txn().unwrap();
3731            let fk_b = key.to_bytes();
3732            db.frames_soa
3733                .put(&mut wtxn, &fk_b[..], &poison[..])
3734                .unwrap();
3735            wtxn.commit().unwrap();
3736        }
3737        assert!(db.has_valid_cooked_soa(key).unwrap());
3738        let pos = db.get_positions(key).unwrap();
3739        assert!(
3740            (pos[0][0] - 99.0).abs() < 1e-12,
3741            "valid RCSO must skip CON parse, got {}",
3742            pos[0][0]
3743        );
3744        assert!((con_x - 0.6394).abs() < 1e-4);
3745        assert_eq!(db.get_frame_text(key).unwrap(), text);
3746        assert_eq!(db.frame_hash(key).unwrap().to_bytes(), h.to_bytes());
3747    }
3748
3749    #[test]
3750    fn traj_meta_old_json_loads_without_topo_fields() {
3751        let m: TrajMeta = serde_json::from_str(r#"{"n_frames":3,"source":"a.con"}"#).unwrap();
3752        assert_eq!(m.n_frames, 3);
3753        assert_eq!(m.source, "a.con");
3754        assert!(m.topo_cutoff.is_none());
3755        assert!(m.topo_graph.is_none());
3756        assert!(m.topo_hops.is_none());
3757        assert!(m.topo_method.is_none());
3758    }
3759
3760    fn seed_topo(db: &ConCorpus, fk: FrameKey, hex: &str) {
3761        let (idx, by) = db.require_topo_dbs().unwrap();
3762        let _gate = db.write_guard();
3763        let mut wtxn = db.env.write_txn().unwrap();
3764        db.put_topo_entry(&mut wtxn, idx, by, fk, hex).unwrap();
3765        wtxn.commit().unwrap();
3766    }
3767
3768    fn stamp_traj_topo(db: &ConCorpus, tid: TrajId, method: &str, cutoff: f64) {
3769        let mut meta = db.traj_meta(tid).unwrap().expect("meta");
3770        meta.stamp_topo(cutoff, "cutoff", 2, method);
3771        let _gate = db.write_guard();
3772        let mut wtxn = db.env.write_txn().unwrap();
3773        let tid_key = tid.to_be_bytes();
3774        db.traj_meta
3775            .put(
3776                &mut wtxn,
3777                &tid_key[..],
3778                &serde_json::to_string(&meta).unwrap(),
3779            )
3780            .unwrap();
3781        wtxn.commit().unwrap();
3782    }
3783
3784    #[test]
3785    fn select_topo_key_prefix_without_seams() {
3786        let dir = tempfile::tempdir().unwrap();
3787        let db = ConCorpus::open(dir.path()).unwrap();
3788        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3789            .unwrap();
3790        let k0 = FrameKey {
3791            traj_id: 1,
3792            frame_idx: 0,
3793        };
3794        seed_topo(&db, k0, "abc123");
3795        let hits = db.select(&Select::new().topo_key("ABC123")).unwrap();
3796        assert_eq!(hits, vec![k0]);
3797        let miss = db.select(&Select::new().topo_key("00")).unwrap();
3798        assert!(miss.is_empty());
3799    }
3800
3801    #[test]
3802    fn reindex_without_annotation_leaves_idx_topo_empty() {
3803        let dir = tempfile::tempdir().unwrap();
3804        let db = ConCorpus::open(dir.path()).unwrap();
3805        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3806            .unwrap();
3807        let k0 = FrameKey {
3808            traj_id: 1,
3809            frame_idx: 0,
3810        };
3811        seed_topo(&db, k0, "abc123");
3812        assert_eq!(
3813            db.select(&Select::new().topo_key("abc123")).unwrap(),
3814            vec![k0]
3815        );
3816        db.reindex().unwrap();
3817        assert!(db
3818            .select(&Select::new().topo_key("abc123"))
3819            .unwrap()
3820            .is_empty());
3821        assert_eq!(db.frame_topo_key(k0).unwrap().as_deref(), Some("abc123"));
3822    }
3823
3824    #[test]
3825    fn reindex_rebuilds_idx_topo_when_params_agree() {
3826        let dir = tempfile::tempdir().unwrap();
3827        let db = ConCorpus::open(dir.path()).unwrap();
3828        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3829            .unwrap();
3830        db.append_trajectory_path(2, fixture("tiny_cuh2.con"))
3831            .unwrap();
3832        let k0 = FrameKey {
3833            traj_id: 1,
3834            frame_idx: 0,
3835        };
3836        let k1 = FrameKey {
3837            traj_id: 2,
3838            frame_idx: 0,
3839        };
3840        seed_topo(&db, k0, "deadbeef");
3841        seed_topo(&db, k1, "deadbeef");
3842        stamp_traj_topo(&db, 1, "nauty", 3.0);
3843        stamp_traj_topo(&db, 2, "nauty", 3.0);
3844        let (idx, _) = db.require_topo_dbs().unwrap();
3845        {
3846            let _gate = db.write_guard();
3847            let mut wtxn = db.env.write_txn().unwrap();
3848            idx.clear(&mut wtxn).unwrap();
3849            wtxn.commit().unwrap();
3850        }
3851        assert!(db
3852            .select(&Select::new().topo_key("deadbeef"))
3853            .unwrap()
3854            .is_empty());
3855        db.reindex().unwrap();
3856        let hits = db.select(&Select::new().topo_key("deadbeef")).unwrap();
3857        assert_eq!(hits.len(), 2);
3858    }
3859
3860    #[test]
3861    fn reindex_refuses_mixed_methods() {
3862        let dir = tempfile::tempdir().unwrap();
3863        let db = ConCorpus::open(dir.path()).unwrap();
3864        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3865            .unwrap();
3866        db.append_trajectory_path(2, fixture("tiny_cuh2.con"))
3867            .unwrap();
3868        stamp_traj_topo(&db, 1, "nauty", 3.0);
3869        stamp_traj_topo(&db, 2, "wl", 3.0);
3870        let err = db.reindex().unwrap_err().to_string();
3871        assert!(err.contains("mixed topology"), "{err}");
3872        let cu = db.select(&Select::new().require_symbol("Cu")).unwrap();
3873        assert_eq!(cu.len(), 2);
3874    }
3875
3876    #[test]
3877    fn extend_preserves_topology_fields() {
3878        let dir = tempfile::tempdir().unwrap();
3879        let db = ConCorpus::open(dir.path()).unwrap();
3880        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3881            .unwrap();
3882        stamp_traj_topo(&db, 1, "nauty", 3.0);
3883        let fr = db
3884            .get_frame(FrameKey {
3885                traj_id: 1,
3886                frame_idx: 0,
3887            })
3888            .unwrap();
3889        db.extend_trajectory_frames(1, &[fr], "extra").unwrap();
3890        let meta = db.traj_meta(1).unwrap().unwrap();
3891        assert_eq!(meta.n_frames, 2);
3892        assert_eq!(meta.topo_method.as_deref(), Some("nauty"));
3893        assert_eq!(meta.topo_cutoff, Some(3.0));
3894        assert_eq!(meta.topo_graph.as_deref(), Some("cutoff"));
3895        assert_eq!(meta.topo_hops, Some(2));
3896    }
3897
3898    #[test]
3899    fn annotate_without_seams_names_command() {
3900        let dir = tempfile::tempdir().unwrap();
3901        let db = ConCorpus::open(dir.path()).unwrap();
3902        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3903            .unwrap();
3904        let err = db
3905            .annotate_topology(AnnotateTopologyOpts::new(3.0).seams("/no/such/seams-binary"))
3906            .unwrap_err()
3907            .to_string();
3908        assert!(err.contains("seams fingerprint"), "{err}");
3909    }
3910
3911    #[test]
3912    fn find_by_topology_without_params_errors() {
3913        let dir = tempfile::tempdir().unwrap();
3914        let db = ConCorpus::open(dir.path()).unwrap();
3915        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3916            .unwrap();
3917        let err = db
3918            .find_by_topology_path(fixture("tiny_cuh2.con"))
3919            .unwrap_err()
3920            .to_string();
3921        assert!(err.contains("annotate-topology"), "{err}");
3922    }
3923
3924    #[test]
3925    fn readonly_missing_topo_select_is_empty() {
3926        let dir = tempfile::tempdir().unwrap();
3927        {
3928            let db = ConCorpus::open(dir.path()).unwrap();
3929            db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3930                .unwrap();
3931            db.close();
3932        }
3933        let db = ConCorpus::open_readonly(dir.path()).unwrap();
3934        let hits = db.select(&Select::new().topo_key("abc")).unwrap();
3935        assert!(hits.is_empty());
3936    }
3937
3938    fn seams_or_skip() -> Option<std::path::PathBuf> {
3939        match crate::topology::resolve_seams_binary(None) {
3940            Ok(p) => Some(p),
3941            Err(e) => {
3942                eprintln!("skipping topology seams test: {e}");
3943                None
3944            }
3945        }
3946    }
3947
3948    fn permute_cu_rows(text: &str) -> String {
3949        let mut lines: Vec<&str> = text.lines().collect();
3950        let mut cu = Vec::new();
3951        for (i, line) in lines.iter().enumerate() {
3952            let t = line.trim_start();
3953            if t.starts_with("0.6394") || t.starts_with("3.1969") {
3954                cu.push(i);
3955            }
3956        }
3957        assert_eq!(cu.len(), 2, "expected two Cu coordinate rows");
3958        lines.swap(cu[0], cu[1]);
3959        let mut out = lines.join("\n");
3960        if text.ends_with('\n') {
3961            out.push('\n');
3962        }
3963        out
3964    }
3965
3966    fn break_one_bond(text: &str) -> String {
3967        let mut fr = {
3968            let mut it = ConFrameIterator::new(text);
3969            it.next().unwrap().unwrap()
3970        };
3971        // Default seams typeI is the first atom type (Cu here). H-H is
3972        // inside 3 A but is not in that cutoff graph; displace a Cu.
3973        let cu = fr
3974            .atom_data
3975            .iter_mut()
3976            .find(|a| a.symbol.as_ref() == "Cu")
3977            .expect("Cu");
3978        cu.z += 10.0;
3979        let mut buf = Cursor::new(Vec::new());
3980        {
3981            let mut w = ConFrameWriter::new(&mut buf);
3982            w.write_frame(&fr).unwrap();
3983        }
3984        String::from_utf8(buf.into_inner()).unwrap()
3985    }
3986
3987    #[test]
3988    fn annotate_permutation_shares_hex_broken_does_not() {
3989        let Some(seams) = seams_or_skip() else {
3990            return;
3991        };
3992        let dir = tempfile::tempdir().unwrap();
3993        let db = ConCorpus::open(dir.path()).unwrap();
3994        db.append_trajectory_path(1, fixture("tiny_cuh2.con"))
3995            .unwrap();
3996        let n = db
3997            .annotate_topology(AnnotateTopologyOpts::new(3.0).seams(&seams))
3998            .unwrap();
3999        assert_eq!(n, 1);
4000        let k0 = FrameKey {
4001            traj_id: 1,
4002            frame_idx: 0,
4003        };
4004        let hex = db.frame_topo_key(k0).unwrap().expect("hex");
4005        assert!(!hex.is_empty());
4006        let meta = db.traj_meta(1).unwrap().unwrap();
4007        assert_eq!(meta.topo_cutoff, Some(3.0));
4008        assert_eq!(meta.topo_graph.as_deref(), Some("cutoff"));
4009        assert_eq!(meta.topo_hops, Some(2));
4010        assert!(meta.topo_method.is_some());
4011
4012        let orig = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
4013        let perm = permute_cu_rows(&orig);
4014        let broken = break_one_bond(&orig);
4015        let perm_hits = db.find_by_topology_text(&perm).unwrap();
4016        assert_eq!(perm_hits, vec![k0]);
4017        let broken_hits = db.find_by_topology_text(&broken).unwrap();
4018        assert!(
4019            broken_hits.is_empty(),
4020            "broken Cu-Cu bond must not share topology key, got {broken_hits:?}"
4021        );
4022        let sel = db.select(&Select::new().topo_key(&hex)).unwrap();
4023        assert_eq!(sel, vec![k0]);
4024    }
4025
4026    fn raw_spans(path: &std::path::Path) -> Vec<String> {
4027        let raw = std::fs::read_to_string(path).unwrap();
4028        let mut it = ConFrameIterator::new(&raw);
4029        let mut out = Vec::new();
4030        loop {
4031            match it.next_with_raw_span(&raw) {
4032                None => break,
4033                Some(Ok((_, blob))) => out.push(blob.to_owned()),
4034                Some(Err(e)) => panic!("{e}"),
4035            }
4036        }
4037        out
4038    }
4039
4040    #[test]
4041    fn zstd_round_trip_is_byte_exact_and_does_not_cook() {
4042        let path = fixture("sulfolene.con");
4043        let spans = raw_spans(&path);
4044        assert!(!spans.is_empty());
4045        let dir = tempfile::tempdir().unwrap();
4046        let db = ConCorpus::open(dir.path()).unwrap();
4047        let n = db.append_trajectory_path(1, &path).unwrap();
4048        assert_eq!(n as usize, spans.len());
4049        for (i, span) in spans.iter().enumerate() {
4050            let key = FrameKey {
4051                traj_id: 1,
4052                frame_idx: i as u32,
4053            };
4054            assert_eq!(db.get_frame_text(key).unwrap(), *span);
4055            let stored = db.stored_frame(key).unwrap();
4056            assert!(
4057                crate::frame_codec::is_single_zstd_frame(&stored),
4058                "frame {i} was not one tagged zstd frame"
4059            );
4060            assert_eq!(crate::frame_codec::load_frame(&stored).unwrap(), *span);
4061            assert_eq!(
4062                db.frame_hash(key).unwrap(),
4063                hash_frame_bytes(span.as_bytes())
4064            );
4065            assert!(!db.has_valid_cooked_soa(key).unwrap());
4066        }
4067    }
4068
4069    #[test]
4070    fn grouped_frames_round_trip_and_migrate_does_not_wrap_again() {
4071        let one = std::fs::read_to_string(fixture("sulfolene.con")).unwrap();
4072        let mut text = String::new();
4073        for i in 0..80u32 {
4074            let shifted = one.replacen(
4075                "10.477713",
4076                &format!("{:.6}", 10.477713 + f64::from(i) * 1e-4),
4077                1,
4078            );
4079            text.push_str(&shifted);
4080            if !shifted.ends_with('\n') {
4081                text.push('\n');
4082            }
4083        }
4084        let dir = tempfile::tempdir().unwrap();
4085        let path = dir.path().join("in.con");
4086        std::fs::write(&path, &text).unwrap();
4087        let spans = raw_spans(&path);
4088        assert_eq!(spans.len(), 80);
4089        let db = ConCorpus::open(dir.path().join("db")).unwrap();
4090        let n = db.append_trajectory_str(1, &text, "mem").unwrap();
4091        assert_eq!(n, 80);
4092        let mut stored_sum = 0usize;
4093        let mut single_sum = 0usize;
4094        for (i, span) in spans.iter().enumerate() {
4095            let key = FrameKey {
4096                traj_id: 1,
4097                frame_idx: i as u32,
4098            };
4099            assert_eq!(db.get_frame_text(key).unwrap(), *span);
4100            assert_eq!(
4101                db.frame_hash(key).unwrap(),
4102                hash_frame_bytes(span.as_bytes())
4103            );
4104            let stored = db.stored_frame(key).unwrap();
4105            stored_sum += stored.len();
4106            single_sum += crate::frame_codec::store_frame(span, true).unwrap().len();
4107            if i == 0 {
4108                assert!(crate::frame_codec::is_group(&stored));
4109            } else if i < 64 {
4110                assert!(crate::frame_codec::is_redirect(&stored));
4111            }
4112            assert!(!db.has_valid_cooked_soa(key).unwrap());
4113        }
4114        assert!(stored_sum < single_sum, "{stored_sum} vs {single_sum}");
4115        assert_eq!(db.migrate_frames().unwrap(), 0);
4116        assert_eq!(
4117            db.get_frame_text(FrameKey {
4118                traj_id: 1,
4119                frame_idx: 79
4120            })
4121            .unwrap(),
4122            spans[79]
4123        );
4124    }
4125
4126    #[test]
4127    fn legacy_utf8_reads_and_migrate_never_double_encodes() {
4128        let path = fixture("tiny_cuh2.con");
4129        let spans = raw_spans(&path);
4130        let dir = tempfile::tempdir().unwrap();
4131        let db = ConCorpus::open(dir.path()).unwrap();
4132        db.set_compress_frames(false);
4133        let n = db.append_trajectory_path(1, &path).unwrap();
4134        assert_eq!(n as usize, spans.len());
4135        for (i, span) in spans.iter().enumerate() {
4136            let key = FrameKey {
4137                traj_id: 1,
4138                frame_idx: i as u32,
4139            };
4140            assert_eq!(db.stored_frame(key).unwrap(), span.as_bytes());
4141            assert_eq!(db.get_frame_text(key).unwrap(), *span);
4142            assert_eq!(
4143                db.frame_hash(key).unwrap(),
4144                hash_frame_bytes(span.as_bytes())
4145            );
4146        }
4147        drop(db);
4148        let db = ConCorpus::open(dir.path()).unwrap();
4149        let key0 = FrameKey {
4150            traj_id: 1,
4151            frame_idx: 0,
4152        };
4153        let hash0 = db.frame_hash(key0).unwrap();
4154        assert_eq!(db.migrate_frames().unwrap() as usize, spans.len());
4155        assert_eq!(db.get_frame_text(key0).unwrap(), spans[0]);
4156        assert_eq!(db.frame_hash(key0).unwrap(), hash0);
4157        let tagged = db.stored_frame(key0).unwrap();
4158        assert!(crate::frame_codec::is_single_zstd_frame(&tagged));
4159        assert_eq!(crate::frame_codec::load_frame(&tagged).unwrap(), spans[0]);
4160        assert!(std::str::from_utf8(&tagged).is_err());
4161        db.reindex().unwrap();
4162        assert_eq!(db.frame_hash(key0).unwrap(), hash0);
4163        assert_eq!(db.get_frame_text(key0).unwrap(), spans[0]);
4164        assert_eq!(db.migrate_frames().unwrap(), 0);
4165        assert_eq!(db.stored_frame(key0).unwrap(), tagged);
4166        assert!(!db.has_valid_cooked_soa(key0).unwrap());
4167    }
4168
4169    #[test]
4170    fn repeated_ingest_stores_identical_frame_bytes() {
4171        let path = fixture("tiny_multi_cuh2.con");
4172        let mut dumps = Vec::new();
4173        for _ in 0..2 {
4174            let dir = tempfile::tempdir().unwrap();
4175            let db = ConCorpus::open(dir.path()).unwrap();
4176            db.append_trajectory_path(1, &path).unwrap();
4177            dumps.push(db.canonical_dump().unwrap());
4178        }
4179        assert_eq!(dumps[0], dumps[1]);
4180    }
4181}