Skip to main content

readcon_db/
shard.rs

1//! HPC multi-writer: one LMDB env **per shard** so writers do not serialize on a single
2//! write_txn. Route `traj_id % n_shards` to a shard directory.
3//!
4//! One writer owns each `shard_id` across the job (`traj_id % n_shards`). If many
5//! ranks share a shard id, each node keeps a private tree, `drain`s to a unique dest,
6//! then `join_drained_roots`. Global select fans out across shards.
7//!
8//! This is **not** multi-writer inside one LMDB env (impossible). It is **partitioned writers**,
9//! the standard embedded pattern for high write concurrency on one filesystem.
10
11use std::path::{Path, PathBuf};
12
13#[cfg(test)]
14use std::cell::Cell;
15
16use readcon_core::types::ConFrame;
17
18use crate::corpus::ConCorpus;
19use crate::error::{Error, Result};
20use crate::keys::{FrameKey, TrajId};
21use crate::select::Select;
22
23/// Default shard count for HPC campaign roots (power of two aids routing).
24pub const DEFAULT_N_SHARDS: u32 = 64;
25
26/// Manifest file in the corpus root describing shard layout.
27const MANIFEST: &str = "shards.json";
28
29#[cfg(test)]
30thread_local! {
31    pub(crate) static FAIL_DEST_MAN_COPY: Cell<bool> = const { Cell::new(false) };
32}
33
34#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
35pub struct ShardManifest {
36    pub n_shards: u32,
37    pub version: u32,
38}
39
40/// Multi-shard campaign corpus: `root/shard_XXXX/` each holds an independent `ConCorpus`.
41pub struct ShardedConCorpus {
42    root: PathBuf,
43    n_shards: u32,
44    /// Lazily opened shards (only those touched). Avoid opening all 10^6 writers' shards in one process.
45    shards: Vec<Option<ConCorpus>>,
46}
47
48impl ShardedConCorpus {
49    /// Open a sharded root that already has `shards.json`. Does not mkdir.
50    pub fn open_existing(root: impl AsRef<Path>) -> Result<Self> {
51        let root = root.as_ref();
52        if !root.join(MANIFEST).is_file() {
53            return Err(Error::Message(format!(
54                "missing shards.json: {}",
55                root.display()
56            )));
57        }
58        Self::open(root, 1)
59    }
60
61    /// Create or open a sharded root. If manifest missing, writes one with `n_shards`.
62    pub fn open(root: impl AsRef<Path>, n_shards: u32) -> Result<Self> {
63        let root = root.as_ref().to_path_buf();
64        std::fs::create_dir_all(&root)?;
65        let manifest_path = root.join(MANIFEST);
66        let n_shards = if manifest_path.is_file() {
67            let s = std::fs::read_to_string(&manifest_path)?;
68            let m: ShardManifest = serde_json::from_str(&s)?;
69            m.n_shards
70        } else {
71            if n_shards == 0 {
72                return Err(Error::Message("n_shards must be >= 1".into()));
73            }
74            let m = ShardManifest {
75                n_shards,
76                version: 1,
77            };
78            std::fs::write(&manifest_path, serde_json::to_string_pretty(&m)?)?;
79            n_shards
80        };
81        let mut shards = Vec::with_capacity(n_shards as usize);
82        shards.resize_with(n_shards as usize, || None);
83        Ok(Self {
84            root,
85            n_shards,
86            shards,
87        })
88    }
89
90    /// Canonical dumps of every shard with data, in shard order, each
91    /// prefixed by its shard id. See [`ConCorpus::canonical_dump`].
92    pub fn canonical_dump(&mut self) -> Result<Vec<u8>> {
93        let mut out = Vec::new();
94        for sid in 0..self.n_shards {
95            if !self.shard_has_data(sid) {
96                continue;
97            }
98            let d = self.shard_mut(sid)?.canonical_dump()?;
99            out.extend_from_slice(&sid.to_le_bytes());
100            out.extend_from_slice(&(d.len() as u64).to_le_bytes());
101            out.extend_from_slice(&d);
102        }
103        Ok(out)
104    }
105
106    pub fn n_shards(&self) -> u32 {
107        self.n_shards
108    }
109
110    pub fn root(&self) -> &Path {
111        &self.root
112    }
113
114    #[inline]
115    pub fn shard_for_traj(traj_id: TrajId, n_shards: u32) -> u32 {
116        (traj_id % u64::from(n_shards)) as u32
117    }
118
119    fn shard_path(&self, shard_id: u32) -> PathBuf {
120        self.root.join(format!("shard_{shard_id:04}"))
121    }
122
123    fn shard_has_data(&self, shard_id: u32) -> bool {
124        self.shard_path(shard_id).join("data.mdb").is_file()
125    }
126
127    fn release_shards(&mut self) {
128        for slot in &mut self.shards {
129            if let Some(c) = slot.take() {
130                c.close();
131            }
132        }
133    }
134
135    /// Open one shard env (creates dir). Safe for many processes to open **different** shards.
136    pub fn shard_mut(&mut self, shard_id: u32) -> Result<&ConCorpus> {
137        if shard_id >= self.n_shards {
138            return Err(Error::Message(format!(
139                "shard_id {shard_id} >= n_shards {}",
140                self.n_shards
141            )));
142        }
143        let i = shard_id as usize;
144        if self.shards[i].is_none() {
145            let p = self.shard_path(shard_id);
146            self.shards[i] = Some(ConCorpus::open(p)?);
147        }
148        Ok(self.shards[i].as_ref().unwrap())
149    }
150
151    /// Open only the shard for `traj_id` (HPC rank typically owns one shard).
152    pub fn open_shard_for_traj(
153        root: impl AsRef<Path>,
154        traj_id: TrajId,
155    ) -> Result<(u32, ConCorpus)> {
156        let root = root.as_ref();
157        let manifest_path = root.join(MANIFEST);
158        let n_shards = if manifest_path.is_file() {
159            let m: ShardManifest = serde_json::from_str(&std::fs::read_to_string(&manifest_path)?)?;
160            m.n_shards
161        } else {
162            let _ = Self::open(root, DEFAULT_N_SHARDS)?;
163            DEFAULT_N_SHARDS
164        };
165        let sid = Self::shard_for_traj(traj_id, n_shards);
166        let corpus = ConCorpus::open(root.join(format!("shard_{sid:04}")))?;
167        Ok((sid, corpus))
168    }
169
170    /// Open a **single** shard by id. That rank must be the only writer of
171    /// `shard_id` in the job, or drain to a unique dest and `join_drained_roots`.
172    pub fn open_shard(root: impl AsRef<Path>, shard_id: u32) -> Result<ConCorpus> {
173        let root = root.as_ref();
174        let manifest_path = root.join(MANIFEST);
175        if !manifest_path.is_file() {
176            return Err(Error::Message(format!(
177                "missing shards.json: {}",
178                root.display()
179            )));
180        }
181        let m: ShardManifest = serde_json::from_str(&std::fs::read_to_string(&manifest_path)?)?;
182        let n_shards = m.n_shards;
183        if shard_id >= n_shards {
184            return Err(Error::Message(format!(
185                "shard_id {shard_id} >= n_shards {n_shards}"
186            )));
187        }
188        ConCorpus::open(root.join(format!("shard_{shard_id:04}")))
189    }
190
191    pub fn append_trajectory_path(
192        &mut self,
193        traj_id: TrajId,
194        file: impl AsRef<Path>,
195    ) -> Result<u32> {
196        self.append_trajectory_path_units(traj_id, file, None)
197    }
198
199    pub fn append_trajectory_path_units(
200        &mut self,
201        traj_id: TrajId,
202        file: impl AsRef<Path>,
203        units: Option<serde_json::Value>,
204    ) -> Result<u32> {
205        let sid = Self::shard_for_traj(traj_id, self.n_shards);
206        let c = self.shard_mut(sid)?;
207        c.append_trajectory_path_units(traj_id, file, units)
208    }
209
210    pub fn append_trajectory_str(
211        &mut self,
212        traj_id: TrajId,
213        contents: &str,
214        source: impl Into<String>,
215    ) -> Result<u32> {
216        let sid = Self::shard_for_traj(traj_id, self.n_shards);
217        let c = self.shard_mut(sid)?;
218        c.append_trajectory_str(traj_id, contents, source)
219    }
220
221    pub fn append_trajectory_frames(
222        &mut self,
223        traj_id: TrajId,
224        frames: &[ConFrame],
225        source: impl Into<String>,
226    ) -> Result<u32> {
227        let sid = Self::shard_for_traj(traj_id, self.n_shards);
228        let c = self.shard_mut(sid)?;
229        c.append_trajectory_frames(traj_id, frames, source)
230    }
231
232    /// Fan-out select across shards that already exist (read-only; does not mint).
233    pub fn select(&mut self, sel: &Select) -> Result<Vec<FrameKey>> {
234        self.release_shards();
235        let mut out = Vec::new();
236        for sid in 0..self.n_shards {
237            if !self.shard_has_data(sid) {
238                continue;
239            }
240            let c = ConCorpus::open_readonly(self.shard_path(sid))?;
241            out.extend(c.select(sel)?);
242            c.close();
243        }
244        out.sort();
245        if let Some(lim) = sel.limit {
246            out.truncate(lim);
247        }
248        Ok(out)
249    }
250
251    pub fn get_frame_text(&mut self, key: FrameKey) -> Result<String> {
252        self.release_shards();
253        let sid = Self::shard_for_traj(key.traj_id, self.n_shards);
254        if !self.shard_has_data(sid) {
255            return Err(Error::Message(format!(
256                "shard_{sid:04} is not a corpus directory"
257            )));
258        }
259        let c = ConCorpus::open_readonly(self.shard_path(sid))?;
260        let text = c.get_frame_text(key)?;
261        c.close();
262        Ok(text)
263    }
264
265    pub fn reindex_all(&mut self) -> Result<u32> {
266        let mut n = 0u32;
267        for sid in 0..self.n_shards {
268            if self.shard_has_data(sid) {
269                n += self.shard_mut(sid)?.reindex()?;
270            }
271        }
272        Ok(n)
273    }
274
275    /// Compact-copy each present shard onto `dst` (data.mdb only, no lockfile).
276    /// Refuses a dest shard that already exists so two node-local trees that
277    /// share a shard id cannot last-writer-wins.
278    pub fn drain_to(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> Result<u32> {
279        let src = src.as_ref();
280        let dst = dst.as_ref();
281        let man = src.join(MANIFEST);
282        if !man.is_file() {
283            return Err(Error::Message("drain: missing shards.json".into()));
284        }
285        let dest_was_new = !dst.exists();
286        let dest_man = dst.join(MANIFEST);
287        let mut copied_manifest = false;
288        let mut created = Vec::new();
289        let written = (|| -> Result<u32> {
290            std::fs::create_dir_all(dst)?;
291            if dest_man.is_file() {
292                let existing: ShardManifest =
293                    serde_json::from_str(&std::fs::read_to_string(&dest_man)?)?;
294                let incoming: ShardManifest =
295                    serde_json::from_str(&std::fs::read_to_string(&man)?)?;
296                if existing.n_shards != incoming.n_shards {
297                    return Err(Error::Message(
298                        "drain: dest shards.json n_shards does not match src".into(),
299                    ));
300                }
301            }
302            let m: ShardManifest = serde_json::from_str(&std::fs::read_to_string(&man)?)?;
303            for i in 0..m.n_shards {
304                let name = format!("shard_{i:04}");
305                if src.join(&name).join("data.mdb").is_file()
306                    && dst.join(&name).join("data.mdb").is_file()
307                {
308                    return Err(Error::Message(format!(
309                        "drain: dest {name} exists; refuse overwrite. Drain each node to a unique dest, then join-drained."
310                    )));
311                }
312            }
313            copied_manifest = !dest_man.is_file();
314            if copied_manifest {
315                #[cfg(test)]
316                if FAIL_DEST_MAN_COPY.with(|f| f.replace(false)) {
317                    return Err(Error::Message("drain: dest_man copy".into()));
318                }
319                std::fs::copy(&man, &dest_man)?;
320            }
321            let mut n = 0u32;
322            for i in 0..m.n_shards {
323                let name = format!("shard_{i:04}");
324                let from = src.join(&name);
325                if !from.join("data.mdb").is_file() {
326                    continue;
327                }
328                let to = dst.join(&name);
329                created.push(to.clone());
330                let ro = ConCorpus::open_readonly(&from)?;
331                ro.snapshot_to(&to)?;
332                ro.close();
333                n += 1;
334            }
335            Ok(n)
336        })();
337        if written.is_err() {
338            if dest_was_new {
339                let _ = std::fs::remove_dir_all(dst);
340            } else {
341                for p in &created {
342                    let _ = std::fs::remove_dir_all(p);
343                }
344                if copied_manifest {
345                    let _ = std::fs::remove_file(&dest_man);
346                }
347            }
348        }
349        written
350    }
351}
352
353#[cfg(test)]
354mod tests {
355    use super::*;
356    use std::sync::Arc;
357    use std::thread;
358
359    fn fixture(name: &str) -> PathBuf {
360        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
361            .join("resources/test")
362            .join(name)
363    }
364
365    #[test]
366    fn parallel_writers_different_shards() {
367        let dir = tempfile::tempdir().unwrap();
368        let root = dir.path().join("hpc");
369        // 8 shards, 8 threads each write traj_id == shard so zero writer contention across envs.
370        let n_shards = 8u32;
371        ShardedConCorpus::open(&root, n_shards).unwrap();
372        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
373        let root = Arc::new(root);
374        let mut joins = Vec::new();
375        for sid in 0..n_shards {
376            let root = Arc::clone(&root);
377            let text = text.clone();
378            joins.push(thread::spawn(move || {
379                // Each writer opens **only its shard** (HPC rank pattern).
380                let db = ShardedConCorpus::open_shard(root.as_path(), sid).unwrap();
381                let traj = u64::from(sid); // maps to this shard
382                db.append_trajectory_str(traj, &text, format!("shard{sid}"))
383                    .unwrap()
384            }));
385        }
386        let mut ns = Vec::new();
387        for j in joins {
388            ns.push(j.join().unwrap());
389        }
390        assert!(ns.iter().all(|&n| n >= 1));
391        let mut fan = ShardedConCorpus::open(root.as_path(), n_shards).unwrap();
392        let keys = fan.select(&Select::new().require_symbol("Cu")).unwrap();
393        drop(fan);
394        let drained = dir.path().join("pfs");
395        let ncopy = ShardedConCorpus::drain_to(root.as_path(), &drained).unwrap();
396        assert_eq!(ncopy, n_shards);
397        assert!(drained.join("shards.json").is_file());
398        assert!(drained.join("shard_0000").join("data.mdb").is_file());
399        assert!(!drained.join("shard_0000").join("lock.mdb").is_file());
400        let dest_sz = std::fs::metadata(drained.join("shard_0000").join("data.mdb"))
401            .unwrap()
402            .len();
403        assert!(
404            dest_sz < 64 * 1024 * 1024,
405            "compact snapshot must not materialize the 2 GiB map, got {dest_sz}"
406        );
407        assert!(ShardedConCorpus::drain_to(root.as_path(), &drained).is_err());
408        assert!(drained.join("shards.json").is_file());
409        assert_eq!(keys.len(), 8);
410        let joined = dir.path().join("joined");
411        let mut drained_root = ShardedConCorpus::open(&drained, n_shards).unwrap();
412        let njoin = drained_root.join_to_single_env(&joined).unwrap();
413        assert!(njoin >= n_shards);
414        let single = ConCorpus::open(&joined).unwrap();
415        let jk = single.select(&Select::new().require_symbol("Cu")).unwrap();
416        assert_eq!(jk.len(), keys.len());
417    }
418
419    #[test]
420    fn open_shard_for_traj_writes_manifest() {
421        let dir = tempfile::tempdir().unwrap();
422        let root = dir.path().join("fresh");
423        let (sid, db) = ShardedConCorpus::open_shard_for_traj(&root, 0).unwrap();
424        assert_eq!(sid, 0);
425        drop(db);
426        assert!(root.join("shards.json").is_file());
427    }
428
429    #[test]
430    fn join_drained_duplicate_traj_does_not_create_dest() {
431        let dir = tempfile::tempdir().unwrap();
432        let a = dir.path().join("a");
433        let b = dir.path().join("b");
434        ShardedConCorpus::open(&a, 1).unwrap();
435        ShardedConCorpus::open(&b, 1).unwrap();
436        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
437        ShardedConCorpus::open_shard(&a, 0)
438            .unwrap()
439            .append_trajectory_str(7, &text, "a")
440            .unwrap();
441        ShardedConCorpus::open_shard(&b, 0)
442            .unwrap()
443            .append_trajectory_str(7, &text, "b")
444            .unwrap();
445        let dest = dir.path().join("out");
446        let err = join_drained_roots(&[a, b], &dest).unwrap_err();
447        assert!(err.to_string().contains("traj_id"), "{err}");
448        assert!(!dest.exists());
449    }
450
451    #[test]
452    fn join_drained_roots_missing_source_errors() {
453        let dir = tempfile::tempdir().unwrap();
454        let missing = dir.path().join("nope");
455        let dest = dir.path().join("out");
456        let err = join_drained_roots(&[missing], &dest).unwrap_err();
457        assert!(err.to_string().contains("missing shards.json"), "{err}");
458        assert!(!dest.exists());
459    }
460
461    #[test]
462    fn drain_refuse_does_not_write_manifest() {
463        let dir = tempfile::tempdir().unwrap();
464        let root = dir.path().join("hpc");
465        ShardedConCorpus::open(&root, 2).unwrap();
466        let db = ShardedConCorpus::open_shard(&root, 0).unwrap();
467        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
468        db.append_trajectory_str(0, &text, "s0").unwrap();
469        drop(db);
470        let dest = dir.path().join("pfs");
471        std::fs::create_dir_all(dest.join("shard_0000")).unwrap();
472        std::fs::write(dest.join("shard_0000").join("data.mdb"), b"x").unwrap();
473        let err = ShardedConCorpus::drain_to(&root, &dest).unwrap_err();
474        assert!(err.to_string().contains("refuse overwrite"), "{err}");
475        assert!(
476            !dest.join("shards.json").is_file(),
477            "refuse must not leave dest shards.json"
478        );
479    }
480
481    #[test]
482    fn traj_routing_stable() {
483        assert_eq!(ShardedConCorpus::shard_for_traj(0, 64), 0);
484        assert_eq!(ShardedConCorpus::shard_for_traj(65, 64), 1);
485    }
486
487    /// Strong-scaling HPC story: concurrent writers on distinct shards, then
488    /// fan-out select agrees with a **single-env** corpus that ingested the
489    /// same trajectory texts (ground truth membership).
490    #[test]
491    fn multi_shard_writers_select_matches_single_env_baseline() {
492        let dir = tempfile::tempdir().unwrap();
493        let root = dir.path().join("hpc_scale");
494        let baseline = dir.path().join("single_env");
495        let n_shards = 4u32;
496        ShardedConCorpus::open(&root, n_shards).unwrap();
497        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
498        let root_a = Arc::new(root.clone());
499        let mut joins = Vec::new();
500        for sid in 0..n_shards {
501            let root = Arc::clone(&root_a);
502            let text = text.clone();
503            joins.push(thread::spawn(move || {
504                let db = ShardedConCorpus::open_shard(root.as_path(), sid).unwrap();
505                let traj = u64::from(sid);
506                db.append_trajectory_str(traj, &text, format!("s{sid}"))
507                    .unwrap()
508            }));
509        }
510        let mut frames_per_traj = Vec::new();
511        for j in joins {
512            frames_per_traj.push(j.join().unwrap());
513        }
514        assert!(frames_per_traj.iter().all(|&n| n >= 1));
515
516        // Single-env ground truth: same traj_ids and CON text.
517        let single = ConCorpus::open(&baseline).unwrap();
518        for sid in 0..n_shards {
519            let traj = u64::from(sid);
520            let n = single
521                .append_trajectory_str(traj, &text, format!("s{sid}"))
522                .unwrap();
523            assert_eq!(n, frames_per_traj[sid as usize]);
524        }
525
526        let mut fan = ShardedConCorpus::open(&root, n_shards).unwrap();
527        let sharded_keys = fan.select(&Select::new().require_symbol("Cu")).unwrap();
528        let base_keys = single.select(&Select::new().require_symbol("Cu")).unwrap();
529        assert_eq!(sharded_keys.len(), base_keys.len());
530        let mut sk: Vec<_> = sharded_keys
531            .iter()
532            .map(|k| (k.traj_id, k.frame_idx))
533            .collect();
534        let mut bk: Vec<_> = base_keys.iter().map(|k| (k.traj_id, k.frame_idx)).collect();
535        sk.sort_unstable();
536        bk.sort_unstable();
537        assert_eq!(sk, bk, "fan-out select must match single-env membership");
538
539        // Spot-check text blobs agree for each key.
540        for (tid, fidx) in &sk {
541            let key = crate::keys::FrameKey {
542                traj_id: *tid,
543                frame_idx: *fidx,
544            };
545            let a = fan.get_frame_text(key).unwrap();
546            let b = single.get_frame_text(key).unwrap();
547            assert_eq!(a, b);
548        }
549    }
550}
551
552/// Exportable corpus layout kinds for analysis handoff.
553#[derive(Clone, Copy, Debug, PartialEq, Eq)]
554pub enum CorpusExportKind {
555    /// Full sharded root (`shards.json` + `shard_XXXX/`).
556    ShardedLmdb,
557    /// Single-env LMDB directory (one `ConCorpus::open` path).
558    SingleEnvLmdb,
559    /// Filtered extXYZ for external tools (non-LMDB).
560    ExtXyz,
561}
562
563impl CorpusExportKind {
564    pub fn as_str(self) -> &'static str {
565        match self {
566            Self::ShardedLmdb => "sharded-lmdb",
567            Self::SingleEnvLmdb => "single-env-lmdb",
568            Self::ExtXyz => "extxyz",
569        }
570    }
571}
572
573impl ShardedConCorpus {
574    /// **Join:** copy all frames from every shard into a **new single-env** corpus at `dst`
575    /// (traj_id preserved; collision if same traj_id appears on two shards is an error).
576    /// Secondary indexes built via normal append/prepare on each blob. Reversible with
577    /// [`Self::split_single_to_sharded`] using the same `n_shards` and traj routing.
578    pub fn join_to_single_env(&mut self, dst: impl AsRef<Path>) -> Result<u32> {
579        let dst = dst.as_ref();
580        if dst.exists() {
581            return Err(Error::Message(format!(
582                "join dest exists: {}",
583                dst.display()
584            )));
585        }
586        self.release_shards();
587        let mut preview = std::collections::BTreeMap::new();
588        for sid in 0..self.n_shards {
589            if !self.shard_has_data(sid) {
590                continue;
591            }
592            let c = ConCorpus::open_readonly(self.shard_path(sid))?;
593            for fk in c.list_frame_keys()? {
594                preview_traj(&mut preview, fk.traj_id, u64::from(sid))?;
595            }
596            c.close();
597        }
598        let n = (|| -> Result<u32> {
599            let out = ConCorpus::open(dst)?;
600            let mut seen_traj = std::collections::BTreeSet::new();
601            let n = append_sharded_into(self, &out, &mut seen_traj)?;
602            out.close();
603            Ok(n)
604        })();
605        rollback_new_dest(dst, n)
606    }
607
608    /// Join into a temp single-env, export extxyz, then remove the temp dest.
609    pub fn export_extxyz(
610        &mut self,
611        sel: &Select,
612        out: impl AsRef<Path>,
613        energy_key: &str,
614    ) -> Result<u32> {
615        let joined = std::env::temp_dir().join(format!(
616            "readcon_db_join_{}_{}",
617            std::process::id(),
618            std::time::SystemTime::now()
619                .duration_since(std::time::UNIX_EPOCH)
620                .map(|d| d.as_nanos())
621                .unwrap_or(0)
622        ));
623        if joined.exists() {
624            return Err(Error::Message(
625                "export_extxyz: temp join dest exists".into(),
626            ));
627        }
628        struct RemoveOnDrop(PathBuf);
629        impl Drop for RemoveOnDrop {
630            fn drop(&mut self) {
631                let _ = std::fs::remove_dir_all(&self.0);
632            }
633        }
634        let _guard = RemoveOnDrop(joined.clone());
635        self.join_to_single_env(&joined)?;
636        let db = ConCorpus::open_readonly(&joined)?;
637        let keys = db.select(sel)?;
638        let n = db.export_extxyz(&keys, out, energy_key)?;
639        db.close();
640        Ok(n as u32)
641    }
642
643    /// **Split:** read a **single-env** corpus and write a new sharded root at `dst_root`
644    /// with `n_shards` (rewrites manifest). Traj_id preserved; routing is `traj_id % n_shards`.
645    pub fn split_single_to_sharded(
646        single: &ConCorpus,
647        dst_root: impl AsRef<Path>,
648        n_shards: u32,
649    ) -> Result<u32> {
650        if n_shards == 0 {
651            return Err(Error::Message("n_shards must be >= 1".into()));
652        }
653        let dst_root = dst_root.as_ref();
654        if dst_root.exists() {
655            return Err(Error::Message(format!(
656                "split dest exists: {}",
657                dst_root.display()
658            )));
659        }
660        let n = (|| -> Result<u32> {
661            let mut sharded = ShardedConCorpus::open(dst_root, n_shards)?;
662            let keys = single.list_frame_keys()?;
663            let mut by_traj: std::collections::BTreeMap<u64, Vec<FrameKey>> =
664                std::collections::BTreeMap::new();
665            for fk in keys {
666                by_traj.entry(fk.traj_id).or_default().push(fk);
667            }
668            let mut n = 0u32;
669            for (tid, mut fks) in by_traj {
670                fks.sort();
671                let mut concat = String::new();
672                for fk in &fks {
673                    concat.push_str(&single.get_frame_text(*fk)?);
674                }
675                let nf = sharded.append_trajectory_str(tid, &concat, "split-from-single")?;
676                n += nf;
677            }
678            Ok(n)
679        })();
680        rollback_new_dest(dst_root, n)
681    }
682}
683
684fn append_sharded_into(
685    sh: &mut ShardedConCorpus,
686    out: &ConCorpus,
687    seen_traj: &mut std::collections::BTreeSet<u64>,
688) -> Result<u32> {
689    let mut n = 0u32;
690    for sid in 0..sh.n_shards {
691        if !sh.shard_has_data(sid) {
692            continue;
693        }
694        let shard = ConCorpus::open_readonly(sh.shard_path(sid))?;
695        let keys = shard.list_frame_keys()?;
696        let mut by_traj: std::collections::BTreeMap<u64, Vec<FrameKey>> =
697            std::collections::BTreeMap::new();
698        for fk in keys {
699            by_traj.entry(fk.traj_id).or_default().push(fk);
700        }
701        for (tid, mut fks) in by_traj {
702            if !seen_traj.insert(tid) {
703                return Err(Error::Message(format!(
704                    "traj_id {tid} appears in multiple shards or join sources"
705                )));
706            }
707            fks.sort();
708            let mut concat = String::new();
709            for fk in &fks {
710                concat.push_str(&shard.get_frame_text(*fk)?);
711            }
712            n += out.append_trajectory_str(tid, &concat, format!("join-from-shard-{sid}"))?;
713        }
714        shard.close();
715    }
716    Ok(n)
717}
718
719/// Join several drained sharded roots (unique dests after refuse-overwrite)
720/// into one single-env corpus. Traj ids must be unique across sources.
721pub fn join_drained_roots(sources: &[PathBuf], dst: impl AsRef<Path>) -> Result<u32> {
722    if sources.is_empty() {
723        return Err(Error::Message("join-drained: no sources".into()));
724    }
725    let dst = dst.as_ref();
726    if dst.exists() {
727        return Err(Error::Message(format!(
728            "join-drained dest exists: {}",
729            dst.display()
730        )));
731    }
732    for src in sources {
733        if !src.join(MANIFEST).is_file() {
734            return Err(Error::Message(format!(
735                "join-drained: missing shards.json: {}",
736                src.display()
737            )));
738        }
739    }
740    {
741        let mut preview = std::collections::BTreeMap::new();
742        for (si, src) in sources.iter().enumerate() {
743            let sh = ShardedConCorpus::open_existing(src)?;
744            for sid in 0..sh.n_shards {
745                if !sh.shard_has_data(sid) {
746                    continue;
747                }
748                let shard = ConCorpus::open_readonly(sh.shard_path(sid))?;
749                let owner = ((si as u64) << 32) | u64::from(sid);
750                for fk in shard.list_frame_keys()? {
751                    preview_traj(&mut preview, fk.traj_id, owner)?;
752                }
753                shard.close();
754            }
755        }
756    }
757    let n = (|| -> Result<u32> {
758        let out = ConCorpus::open(dst)?;
759        let mut n = 0u32;
760        let mut seen = std::collections::BTreeSet::new();
761        for src in sources {
762            let mut sh = ShardedConCorpus::open_existing(src)?;
763            n += append_sharded_into(&mut sh, &out, &mut seen)?;
764        }
765        out.close();
766        Ok(n)
767    })();
768    rollback_new_dest(dst, n)
769}
770
771/// Join any set of **single-env** corpus directories into one destination (traj_id must be unique).
772pub fn join_corpus_dirs(sources: &[PathBuf], dst: impl AsRef<Path>) -> Result<u32> {
773    let dst = dst.as_ref();
774    if dst.exists() {
775        return Err(Error::Message(format!(
776            "join dest exists: {}",
777            dst.display()
778        )));
779    }
780    for src in sources {
781        if !src.join("data.mdb").is_file() {
782            return Err(Error::Message(format!(
783                "join: missing data.mdb: {}",
784                src.display()
785            )));
786        }
787    }
788    {
789        let mut preview = std::collections::BTreeMap::new();
790        for (si, src) in sources.iter().enumerate() {
791            let c = ConCorpus::open_readonly(src)?;
792            for fk in c.list_frame_keys()? {
793                preview_traj(&mut preview, fk.traj_id, si as u64)?;
794            }
795            c.close();
796        }
797    }
798    let n = (|| -> Result<u32> {
799        let out = ConCorpus::open(dst)?;
800        let mut n = 0u32;
801        let mut seen = std::collections::BTreeSet::new();
802        for src in sources {
803            let c = ConCorpus::open_readonly(src)?;
804            let keys = c.list_frame_keys()?;
805            let mut by_traj: std::collections::BTreeMap<u64, Vec<FrameKey>> =
806                std::collections::BTreeMap::new();
807            for fk in keys {
808                by_traj.entry(fk.traj_id).or_default().push(fk);
809            }
810            for (tid, mut fks) in by_traj {
811                if !seen.insert(tid) {
812                    return Err(Error::Message(format!(
813                        "duplicate traj_id {tid} across join sources"
814                    )));
815                }
816                fks.sort();
817                let mut concat = String::new();
818                for fk in &fks {
819                    concat.push_str(&c.get_frame_text(*fk)?);
820                }
821                n += out.append_trajectory_str(tid, &concat, src.display().to_string())?;
822            }
823            c.close();
824        }
825        out.close();
826        Ok(n)
827    })();
828    rollback_new_dest(dst, n)
829}
830
831fn rollback_new_dest<T>(dst: &Path, r: Result<T>) -> Result<T> {
832    if r.is_err() {
833        let _ = std::fs::remove_dir_all(dst);
834    }
835    r
836}
837
838/// Record `traj_id` under `owner`. Extra frames of the same traj on the
839/// same owner are fine; the same traj on a different owner is a collision.
840fn preview_traj(
841    preview: &mut std::collections::BTreeMap<u64, u64>,
842    traj_id: u64,
843    owner: u64,
844) -> Result<()> {
845    match preview.entry(traj_id) {
846        std::collections::btree_map::Entry::Occupied(e) if *e.get() != owner => {
847            Err(Error::Message(format!(
848                "traj_id {traj_id} appears in multiple shards or join sources"
849            )))
850        }
851        std::collections::btree_map::Entry::Occupied(_) => Ok(()),
852        std::collections::btree_map::Entry::Vacant(e) => {
853            e.insert(owner);
854            Ok(())
855        }
856    }
857}
858
859/// Open a single-env corpus for analysis export. Refuses a sharded root so
860/// `ConCorpus::open` cannot mint `data.mdb` next to `shards.json`.
861pub fn open_single_env_for_export(src: impl AsRef<Path>) -> Result<ConCorpus> {
862    let src = src.as_ref();
863    if src.join(MANIFEST).is_file() {
864        return Err(Error::Message(
865            "compact-export-extxyz: sharded root needs --sharded".into(),
866        ));
867    }
868    ConCorpus::open_readonly(src)
869}
870
871#[cfg(test)]
872mod compaction_tests {
873    use super::*;
874    use crate::keys::FrameKey;
875    use crate::select::Select;
876
877    fn fixture(name: &str) -> PathBuf {
878        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
879            .join("resources/test")
880            .join(name)
881    }
882
883    #[test]
884    fn join_split_reversible_membership() {
885        let dir = tempfile::tempdir().unwrap();
886        let sharded_root = dir.path().join("sharded");
887        let con_text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
888        {
889            let mut s = ShardedConCorpus::open(&sharded_root, 4).unwrap();
890            for tid in [0u64, 1, 2, 3] {
891                s.append_trajectory_str(tid, &con_text, "t").unwrap();
892            }
893        }
894        let mut s = ShardedConCorpus::open(&sharded_root, 4).unwrap();
895        let before = s.select(&Select::new()).unwrap();
896        assert_eq!(before.len(), 4);
897
898        let joined = dir.path().join("joined");
899        let n = s.join_to_single_env(&joined).unwrap();
900        assert_eq!(n, 4);
901        assert!(s.join_to_single_env(&joined).is_err());
902        let joined_c = ConCorpus::open(&joined).unwrap();
903        let mid = joined_c.select(&Select::new()).unwrap();
904        assert_eq!(mid.len(), 4);
905
906        let split_root = dir.path().join("split_again");
907        let n2 = ShardedConCorpus::split_single_to_sharded(&joined_c, &split_root, 4).unwrap();
908        assert_eq!(n2, 4);
909        let mut s2 = ShardedConCorpus::open(&split_root, 4).unwrap();
910        let after = s2.select(&Select::new()).unwrap();
911        assert_eq!(after.len(), before.len());
912        // same traj set
913        let mut bt: Vec<_> = before.iter().map(|k| k.traj_id).collect();
914        let mut at: Vec<_> = after.iter().map(|k| k.traj_id).collect();
915        bt.sort();
916        at.sort();
917        assert_eq!(bt, at);
918    }
919
920    #[test]
921    fn join_drained_roots_keeps_disjoint_trajs_same_shard() {
922        let dir = tempfile::tempdir().unwrap();
923        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
924        let n_shards = 2u32;
925        let node_a = dir.path().join("node_a");
926        let node_b = dir.path().join("node_b");
927        ShardedConCorpus::open(&node_a, n_shards).unwrap();
928        ShardedConCorpus::open(&node_b, n_shards).unwrap();
929        // traj 0 and traj 2 both route to shard_0000
930        ShardedConCorpus::open_shard(&node_a, 0)
931            .unwrap()
932            .append_trajectory_str(0, &text, "a")
933            .unwrap();
934        ShardedConCorpus::open_shard(&node_b, 0)
935            .unwrap()
936            .append_trajectory_str(2, &text, "b")
937            .unwrap();
938        let dest_a = dir.path().join("dest_a");
939        let dest_b = dir.path().join("dest_b");
940        assert_eq!(ShardedConCorpus::drain_to(&node_a, &dest_a).unwrap(), 1);
941        assert_eq!(ShardedConCorpus::drain_to(&node_b, &dest_b).unwrap(), 1);
942        assert!(ShardedConCorpus::drain_to(&node_a, &dest_a).is_err());
943        let joined = dir.path().join("joined");
944        let n = join_drained_roots(&[dest_a.clone(), dest_b.clone()], &joined).unwrap();
945        assert!(n >= 2);
946        let db = ConCorpus::open(&joined).unwrap();
947        let keys = db.select(&Select::new()).unwrap();
948        let mut tids: Vec<u64> = keys.iter().map(|k| k.traj_id).collect();
949        tids.sort_unstable();
950        tids.dedup();
951        assert_eq!(tids, vec![0, 2]);
952        assert!(join_drained_roots(&[dest_a, dest_b], &joined).is_err());
953    }
954
955    #[test]
956    fn join_drained_roots_refuses_existing_dest() {
957        let dir = tempfile::tempdir().unwrap();
958        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
959        let node = dir.path().join("node");
960        ShardedConCorpus::open(&node, 2).unwrap();
961        ShardedConCorpus::open_shard(&node, 0)
962            .unwrap()
963            .append_trajectory_str(0, &text, "a")
964            .unwrap();
965        let dest = dir.path().join("dest");
966        assert_eq!(ShardedConCorpus::drain_to(&node, &dest).unwrap(), 1);
967        let joined = dir.path().join("joined");
968        join_drained_roots(&[dest.clone()], &joined).unwrap();
969        let err = join_drained_roots(&[dest], &joined).unwrap_err();
970        assert!(
971            err.to_string().contains("join-drained dest exists"),
972            "{err}"
973        );
974    }
975
976    #[test]
977    fn export_kinds_documented() {
978        assert_eq!(CorpusExportKind::ShardedLmdb.as_str(), "sharded-lmdb");
979        assert_eq!(CorpusExportKind::SingleEnvLmdb.as_str(), "single-env-lmdb");
980        assert_eq!(CorpusExportKind::ExtXyz.as_str(), "extxyz");
981    }
982
983    #[test]
984    fn select_skips_empty_shard_dir_without_data() {
985        let dir = tempfile::tempdir().unwrap();
986        let root = dir.path().join("sharded");
987        let mut s = ShardedConCorpus::open(&root, 4).unwrap();
988        std::fs::create_dir_all(root.join("shard_0001")).unwrap();
989        let keys = s.select(&Select::new()).unwrap();
990        assert!(keys.is_empty());
991        assert!(!root.join("shard_0001").join("data.mdb").exists());
992        let err = s
993            .get_frame_text(FrameKey {
994                traj_id: 1,
995                frame_idx: 0,
996            })
997            .unwrap_err();
998        assert!(
999            err.to_string().contains("is not a corpus directory"),
1000            "{err}"
1001        );
1002        assert!(!root.join("shard_0001").join("data.mdb").exists());
1003    }
1004
1005    #[test]
1006    fn select_skips_missing_shard_dirs() {
1007        let dir = tempfile::tempdir().unwrap();
1008        let root = dir.path().join("sharded");
1009        let mut s = ShardedConCorpus::open(&root, 8).unwrap();
1010        s.append_trajectory_path(0, fixture("tiny_cuh2.con"))
1011            .unwrap();
1012        drop(s);
1013        let mut s = ShardedConCorpus::open_existing(&root).unwrap();
1014        let keys = s.select(&Select::new()).unwrap();
1015        assert_eq!(keys.len(), 1);
1016        let present: Vec<u32> = (0..8)
1017            .filter(|&i| root.join(format!("shard_{i:04}")).is_dir())
1018            .collect();
1019        assert_eq!(present, vec![0]);
1020    }
1021
1022    #[test]
1023    fn get_frame_text_missing_shard_does_not_mint() {
1024        let dir = tempfile::tempdir().unwrap();
1025        let root = dir.path().join("sharded");
1026        let mut s = ShardedConCorpus::open(&root, 4).unwrap();
1027        let err = s
1028            .get_frame_text(FrameKey {
1029                traj_id: 1,
1030                frame_idx: 0,
1031            })
1032            .unwrap_err();
1033        assert!(
1034            err.to_string().contains("is not a corpus directory"),
1035            "{err}"
1036        );
1037        assert!(!root.join("shard_0001").exists());
1038    }
1039
1040    #[test]
1041    fn join_drained_roots_keeps_multi_frame_traj() {
1042        let dir = tempfile::tempdir().unwrap();
1043        let text = std::fs::read_to_string(fixture("tiny_multi_cuh2.con")).unwrap();
1044        let node = dir.path().join("node");
1045        ShardedConCorpus::open(&node, 2).unwrap();
1046        ShardedConCorpus::open_shard(&node, 0)
1047            .unwrap()
1048            .append_trajectory_str(0, &text, "a")
1049            .unwrap();
1050        let dest = dir.path().join("dest");
1051        assert_eq!(ShardedConCorpus::drain_to(&node, &dest).unwrap(), 1);
1052        let joined = dir.path().join("joined");
1053        let n = join_drained_roots(&[dest], &joined).unwrap();
1054        assert!(n >= 2, "joined frames={n}");
1055        let db = ConCorpus::open(&joined).unwrap();
1056        let p0 = db
1057            .get_positions(crate::keys::FrameKey {
1058                traj_id: 0,
1059                frame_idx: 0,
1060            })
1061            .unwrap();
1062        let p1 = db
1063            .get_positions(crate::keys::FrameKey {
1064                traj_id: 0,
1065                frame_idx: 1,
1066            })
1067            .unwrap();
1068        assert!((p0[0][0] - 0.6394).abs() < 1e-4);
1069        assert!((p1[2][0] - 8.8549).abs() < 1e-4);
1070    }
1071
1072    #[test]
1073    fn join_corpus_dirs_keeps_multi_frame_traj() {
1074        let dir = tempfile::tempdir().unwrap();
1075        let a = dir.path().join("a");
1076        ConCorpus::open(&a)
1077            .unwrap()
1078            .append_trajectory_path(1, fixture("tiny_multi_cuh2.con"))
1079            .unwrap();
1080        let dest = dir.path().join("out");
1081        let n = join_corpus_dirs(&[a], &dest).unwrap();
1082        assert!(n >= 2, "joined frames={n}");
1083        let db = ConCorpus::open(&dest).unwrap();
1084        let p0 = db
1085            .get_positions(crate::keys::FrameKey {
1086                traj_id: 1,
1087                frame_idx: 0,
1088            })
1089            .unwrap();
1090        let p1 = db
1091            .get_positions(crate::keys::FrameKey {
1092                traj_id: 1,
1093                frame_idx: 1,
1094            })
1095            .unwrap();
1096        assert!((p0[0][0] - 0.6394).abs() < 1e-4);
1097        assert!((p1[2][0] - 8.8549).abs() < 1e-4);
1098    }
1099
1100    #[test]
1101    fn drain_to_n_shards_mismatch_refuses() {
1102        let dir = tempfile::tempdir().unwrap();
1103        let src = dir.path().join("src");
1104        let dest = dir.path().join("dest");
1105        ShardedConCorpus::open(&src, 2).unwrap();
1106        ShardedConCorpus::open(&dest, 4).unwrap();
1107        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
1108        ShardedConCorpus::open_shard(&src, 0)
1109            .unwrap()
1110            .append_trajectory_str(0, &text, "a")
1111            .unwrap();
1112        let err = ShardedConCorpus::drain_to(&src, &dest).unwrap_err();
1113        assert!(err.to_string().contains("n_shards"), "{err}");
1114        let man: ShardManifest =
1115            serde_json::from_str(&std::fs::read_to_string(dest.join("shards.json")).unwrap())
1116                .unwrap();
1117        assert_eq!(man.n_shards, 4);
1118        assert!(!dest.join("shard_0000").join("data.mdb").is_file());
1119    }
1120
1121    #[test]
1122    fn drain_to_new_dest_rollback_on_bad_src() {
1123        let dir = tempfile::tempdir().unwrap();
1124        let src = dir.path().join("src");
1125        ShardedConCorpus::open(&src, 2).unwrap();
1126        let shard0 = src.join("shard_0000");
1127        std::fs::create_dir_all(&shard0).unwrap();
1128        std::fs::write(shard0.join("data.mdb"), b"not-lmdb").unwrap();
1129        let dest = dir.path().join("dest");
1130        assert!(!dest.exists());
1131        assert!(ShardedConCorpus::drain_to(&src, &dest).is_err());
1132        assert!(!dest.exists(), "dest_was_new rollback must remove dest");
1133    }
1134
1135    #[test]
1136    fn drain_to_new_dest_rollback_on_empty_src_manifest() {
1137        let dir = tempfile::tempdir().unwrap();
1138        let src = dir.path().join("src");
1139        std::fs::create_dir_all(&src).unwrap();
1140        std::fs::write(src.join("shards.json"), b"").unwrap();
1141        let dest = dir.path().join("dest");
1142        assert!(!dest.exists());
1143        assert!(ShardedConCorpus::drain_to(&src, &dest).is_err());
1144        assert!(
1145            !dest.exists(),
1146            "dest leftover empty dest_man parse must remove dest"
1147        );
1148    }
1149
1150    #[test]
1151    fn drain_to_new_dest_rollback_on_dest_man_copy() {
1152        let dir = tempfile::tempdir().unwrap();
1153        let src = dir.path().join("src");
1154        ShardedConCorpus::open(&src, 2).unwrap();
1155        let dest = dir.path().join("dest");
1156        std::fs::create_dir_all(dest.join("shards.json")).unwrap();
1157        assert!(ShardedConCorpus::drain_to(&src, &dest).is_err());
1158        assert!(
1159            dest.exists(),
1160            "existing dest stays; dest_man copy failed because dest_man is a dir"
1161        );
1162        assert!(dest.join("shards.json").is_dir());
1163        assert!(!dest.join("shards.json").is_file());
1164    }
1165
1166    #[test]
1167    fn drain_to_new_dest_rollback_on_dest_man_copy_fail() {
1168        let dir = tempfile::tempdir().unwrap();
1169        let src = dir.path().join("src");
1170        ShardedConCorpus::open(&src, 2).unwrap();
1171        let dest = dir.path().join("dest");
1172        assert!(!dest.exists());
1173        FAIL_DEST_MAN_COPY.with(|f| f.set(true));
1174        assert!(ShardedConCorpus::drain_to(&src, &dest).is_err());
1175        assert!(
1176            !dest.exists(),
1177            "dest leftover dest_man copy dest_was_new must remove dest"
1178        );
1179    }
1180
1181    #[test]
1182    fn drain_to_failure_removes_shards_created_this_call() {
1183        let dir = tempfile::tempdir().unwrap();
1184        let src = dir.path().join("src");
1185        ShardedConCorpus::open(&src, 2).unwrap();
1186        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
1187        ShardedConCorpus::open_shard(&src, 0)
1188            .unwrap()
1189            .append_trajectory_str(0, &text, "a")
1190            .unwrap();
1191        ShardedConCorpus::open_shard(&src, 1)
1192            .unwrap()
1193            .append_trajectory_str(1, &text, "b")
1194            .unwrap();
1195        let dest = dir.path().join("dest");
1196        std::fs::create_dir_all(&dest).unwrap();
1197        std::fs::write(dest.join("shard_0001"), b"not-a-dir").unwrap();
1198        assert!(ShardedConCorpus::drain_to(&src, &dest).is_err());
1199        assert!(!dest.join("shard_0000").exists());
1200    }
1201
1202    #[test]
1203    fn rollback_new_dest_removes_dir() {
1204        let dir = tempfile::tempdir().unwrap();
1205        let dest = dir.path().join("out");
1206        std::fs::create_dir_all(&dest).unwrap();
1207        std::fs::write(dest.join("x"), b"y").unwrap();
1208        let err = rollback_new_dest::<u32>(&dest, Err(Error::Message("boom".into()))).unwrap_err();
1209        assert!(err.to_string().contains("boom"), "{err}");
1210        assert!(!dest.exists());
1211    }
1212
1213    #[test]
1214    fn join_to_single_env_keeps_multi_frame_traj() {
1215        let dir = tempfile::tempdir().unwrap();
1216        let root = dir.path().join("sharded");
1217        let mut s = ShardedConCorpus::open(&root, 2).unwrap();
1218        s.append_trajectory_path(0, fixture("tiny_multi_cuh2.con"))
1219            .unwrap();
1220        let dest = dir.path().join("out");
1221        let n = s.join_to_single_env(&dest).unwrap();
1222        assert!(n >= 2, "joined frames={n}");
1223        let db = ConCorpus::open(&dest).unwrap();
1224        let keys = db.select(&Select::new()).unwrap();
1225        assert!(keys.len() >= 2);
1226        assert!(keys.iter().all(|k| k.traj_id == 0));
1227        let p0 = db
1228            .get_positions(crate::keys::FrameKey {
1229                traj_id: 0,
1230                frame_idx: 0,
1231            })
1232            .unwrap();
1233        let p1 = db
1234            .get_positions(crate::keys::FrameKey {
1235                traj_id: 0,
1236                frame_idx: 1,
1237            })
1238            .unwrap();
1239        assert!((p0[0][0] - 0.6394).abs() < 1e-4);
1240        assert!((p1[2][0] - 8.8549).abs() < 1e-4);
1241    }
1242
1243    #[test]
1244    fn join_to_single_env_duplicate_traj_does_not_create_dest() {
1245        let dir = tempfile::tempdir().unwrap();
1246        let root = dir.path().join("sharded");
1247        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
1248        let mut s = ShardedConCorpus::open(&root, 2).unwrap();
1249        s.shard_mut(0)
1250            .unwrap()
1251            .append_trajectory_str(0, &text, "a")
1252            .unwrap();
1253        s.shard_mut(1)
1254            .unwrap()
1255            .append_trajectory_str(0, &text, "b")
1256            .unwrap();
1257        let dest = dir.path().join("out");
1258        let err = s.join_to_single_env(&dest).unwrap_err();
1259        assert!(err.to_string().contains("traj_id"), "{err}");
1260        assert!(!dest.exists());
1261    }
1262
1263    #[test]
1264    fn join_corpus_dirs_duplicate_does_not_create_dest() {
1265        let dir = tempfile::tempdir().unwrap();
1266        let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
1267        let a = dir.path().join("a");
1268        let b = dir.path().join("b");
1269        ConCorpus::open(&a)
1270            .unwrap()
1271            .append_trajectory_str(1, &text, "a")
1272            .unwrap();
1273        ConCorpus::open(&b)
1274            .unwrap()
1275            .append_trajectory_str(1, &text, "b")
1276            .unwrap();
1277        let dest = dir.path().join("out");
1278        let err = join_corpus_dirs(&[a, b], &dest).unwrap_err();
1279        assert!(
1280            err.to_string().contains("traj_id") || err.to_string().contains("duplicate"),
1281            "{err}"
1282        );
1283        assert!(!dest.exists());
1284    }
1285
1286    #[test]
1287    fn join_corpus_dirs_missing_mdb_does_not_create_dest() {
1288        let dir = tempfile::tempdir().unwrap();
1289        let missing = dir.path().join("nope");
1290        std::fs::create_dir_all(&missing).unwrap();
1291        let dest = dir.path().join("out");
1292        let err = join_corpus_dirs(&[missing], &dest).unwrap_err();
1293        assert!(err.to_string().contains("missing data.mdb"), "{err}");
1294        assert!(!dest.exists());
1295    }
1296
1297    #[test]
1298    fn open_single_env_for_export_refuses_sharded_root() {
1299        let dir = tempfile::tempdir().unwrap();
1300        let root = dir.path().join("sharded");
1301        ShardedConCorpus::open(&root, 2).unwrap();
1302        assert!(!root.join("data.mdb").exists());
1303        match open_single_env_for_export(&root) {
1304            Ok(_) => panic!("expected sharded-root refuse"),
1305            Err(err) => assert!(err.to_string().contains("--sharded"), "{err}"),
1306        }
1307        assert!(!root.join("data.mdb").exists());
1308    }
1309
1310    #[test]
1311    fn export_extxyz_removes_temp_join() {
1312        let dir = tempfile::tempdir().unwrap();
1313        let root = dir.path().join("sharded");
1314        let mut s = ShardedConCorpus::open(&root, 2).unwrap();
1315        s.append_trajectory_path(0, fixture("tiny_cuh2.con"))
1316            .unwrap();
1317        let out = dir.path().join("out.xyz");
1318        let n = s.export_extxyz(&Select::new(), &out, "energy").unwrap();
1319        assert!(n >= 1);
1320        assert!(out.is_file());
1321        let prefix = format!("readcon_db_join_{}_", std::process::id());
1322        let leftovers: Vec<_> = std::env::temp_dir()
1323            .read_dir()
1324            .unwrap()
1325            .filter_map(|e| e.ok())
1326            .filter(|e| e.file_name().to_string_lossy().starts_with(&prefix))
1327            .collect();
1328        assert!(leftovers.is_empty(), "{leftovers:?}");
1329    }
1330
1331    #[test]
1332    fn open_existing_error_is_missing_shards_json() {
1333        let dir = tempfile::tempdir().unwrap();
1334        match ShardedConCorpus::open_existing(dir.path().join("nope")) {
1335            Ok(_) => panic!("expected missing shards.json"),
1336            Err(err) => {
1337                assert!(err.to_string().starts_with("missing shards.json:"), "{err}");
1338                assert!(!err.to_string().contains("join-drained"));
1339            }
1340        }
1341    }
1342}