1use std::{
2 fs,
3 path::{Path, PathBuf},
4};
5
6use crate::LadduDataResult;
7
8use super::{WritePlan, sink_error};
9
10#[derive(Clone, Debug)]
12pub struct OutputPath {
13 base: PathBuf,
14 mode: OutputMode,
15}
16
17#[derive(Clone, Copy, Debug, Default)]
19pub enum OutputMode {
20 #[default]
22 Auto,
23 SingleFile,
25 PerRankFiles,
27}
28
29impl OutputPath {
30 pub fn new(path: impl Into<PathBuf>) -> Self {
32 Self {
33 base: path.into(),
34 mode: OutputMode::Auto,
35 }
36 }
37
38 pub fn with_mode(mut self, mode: OutputMode) -> Self {
40 self.mode = mode;
41 self
42 }
43
44 pub fn base(&self) -> &Path {
46 &self.base
47 }
48
49 pub fn mode(&self) -> OutputMode {
51 self.mode
52 }
53
54 pub fn resolve(&self, plan: WritePlan, default_extension: &str) -> LadduDataResult<PathBuf> {
61 let mode = match self.mode {
62 OutputMode::Auto if plan.is_distributed() => OutputMode::PerRankFiles,
63 OutputMode::Auto => OutputMode::SingleFile,
64 mode => mode,
65 };
66
67 match mode {
68 OutputMode::SingleFile => {
69 if plan.is_distributed() {
70 return Err(sink_error(
71 "resolve output path",
72 self.base.display(),
73 "single-file output is unsafe with multiple MPI ranks; use per-rank output",
74 ));
75 }
76
77 Ok(self.base.clone())
78 }
79
80 OutputMode::PerRankFiles => Ok(per_rank_path(
81 &self.base,
82 plan.rank(),
83 plan.nranks(),
84 default_extension,
85 )),
86
87 OutputMode::Auto => unreachable!(),
88 }
89 }
90
91 pub fn create_parent_dirs(path: &Path) -> LadduDataResult<()> {
97 if let Some(parent) = path.parent()
98 && !parent.as_os_str().is_empty()
99 {
100 fs::create_dir_all(parent)
101 .map_err(|e| sink_error("create output directory", parent.display(), e))?;
102 }
103
104 Ok(())
105 }
106}
107
108fn per_rank_path(base: &Path, rank: usize, nranks: usize, default_extension: &str) -> PathBuf {
109 if base.extension().is_none() {
110 let ext = default_extension.trim_start_matches('.');
111 return base.join(format!("part-rank{rank:05}-of{nranks:05}.{ext}"));
112 }
113
114 let parent = base.parent().unwrap_or_else(|| Path::new(""));
115 let stem = base.file_stem().unwrap_or_default().to_string_lossy();
116 let ext = base.extension().unwrap_or_default().to_string_lossy();
117
118 parent.join(format!("{stem}.rank{rank:05}-of{nranks:05}.{ext}"))
119}