Skip to main content

laddu_data/io/
output.rs

1use std::{
2    fs,
3    path::{Path, PathBuf},
4};
5
6use crate::LadduDataResult;
7
8use super::{WritePlan, sink_error};
9
10/// Resolves a base output path for serial or distributed writes.
11#[derive(Clone, Debug)]
12pub struct OutputPath {
13    base: PathBuf,
14    mode: OutputMode,
15}
16
17/// Policy for resolving a concrete output path.
18#[derive(Clone, Copy, Debug, Default)]
19pub enum OutputMode {
20    /// Select single-file or per-rank output from the write plan.
21    #[default]
22    Auto,
23    /// Write exactly one file; invalid for distributed writes.
24    SingleFile,
25    /// Write a rank-specific file.
26    PerRankFiles,
27}
28
29impl OutputPath {
30    /// Creates an automatically resolved output path.
31    pub fn new(path: impl Into<PathBuf>) -> Self {
32        Self {
33            base: path.into(),
34            mode: OutputMode::Auto,
35        }
36    }
37
38    /// Returns this path with an explicit output mode.
39    pub fn with_mode(mut self, mode: OutputMode) -> Self {
40        self.mode = mode;
41        self
42    }
43
44    /// Returns the unresolved base path.
45    pub fn base(&self) -> &Path {
46        &self.base
47    }
48
49    /// Returns the output mode.
50    pub fn mode(&self) -> OutputMode {
51        self.mode
52    }
53
54    /// Resolves the concrete path for a write plan.
55    ///
56    /// # Errors
57    ///
58    /// Returns [`crate::LadduDataError`] when single-file output is requested for a
59    /// distributed plan.
60    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    /// Creates a file's parent directories when absent.
92    ///
93    /// # Errors
94    ///
95    /// Returns [`crate::LadduDataError`] when a required directory cannot be created.
96    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}