Skip to main content

aether_renderer_core/
input.rs

1use glob::glob;
2use std::cmp::Ordering;
3use std::path::{Path, PathBuf};
4
5/// Collect files from `input_folder` matching the optional pattern, in natural
6/// order (`frame_2.png` before `frame_10.png`). Defaults to `*.png`.
7///
8/// Besides glob patterns (`*.png`, `scene1_*.png`) this accepts ffmpeg-style
9/// sequence patterns such as `frame_%04d.png`.
10pub fn collect_input_frames(
11    input_folder: &Path,
12    file_pattern: Option<String>,
13) -> Result<Vec<PathBuf>, Box<dyn std::error::Error>> {
14    let pattern = file_pattern.unwrap_or_else(|| "*.png".to_string());
15    let (glob_pattern, sequence) = match sequence_to_glob(&pattern) {
16        Some(g) => (g, Some(pattern.clone())),
17        None => (pattern.clone(), None),
18    };
19    // Escape the folder so names like "renders [final]" are taken literally.
20    let folder = glob::Pattern::escape(&input_folder.to_string_lossy());
21    let glob_path = format!("{}/{}", folder, glob_pattern);
22    let mut frames: Vec<PathBuf> = glob(&glob_path)?
23        .filter_map(Result::ok)
24        .filter(|p| p.is_file())
25        .filter(|p| {
26            // Skip macOS resource forks like ._frame_0001.png
27            !p.file_name()
28                .and_then(|n| n.to_str())
29                .is_some_and(|n| n.starts_with("._"))
30        })
31        .filter(|p| match &sequence {
32            Some(seq) => matches_sequence(seq, p),
33            None => true,
34        })
35        .collect();
36    frames.sort_by(|a, b| natural_cmp(&a.to_string_lossy(), &b.to_string_lossy()));
37    Ok(frames)
38}
39
40/// `frame_%04d.png` -> `frame_*.png`; `None` if the pattern has no `%d`.
41fn sequence_to_glob(pattern: &str) -> Option<String> {
42    let mut out = String::new();
43    let mut found = false;
44    let mut chars = pattern.chars().peekable();
45    while let Some(c) = chars.next() {
46        if c == '%' {
47            let mut spec = String::new();
48            while let Some(&d) = chars.peek() {
49                if d.is_ascii_digit() {
50                    spec.push(d);
51                    chars.next();
52                } else {
53                    break;
54                }
55            }
56            if chars.peek() == Some(&'d') {
57                chars.next();
58                out.push('*');
59                found = true;
60                continue;
61            }
62            out.push('%');
63            out.push_str(&spec);
64        } else {
65            out.push(c);
66        }
67    }
68    found.then_some(out)
69}
70
71/// Does the file name match `frame_%04d.png` exactly (the `*` from the glob is digits)?
72fn matches_sequence(seq: &str, path: &Path) -> bool {
73    let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
74        return false;
75    };
76    let Some(start) = seq.find('%') else {
77        return true;
78    };
79    let Some(d) = seq[start..].find('d').map(|i| start + i) else {
80        return true;
81    };
82    let (prefix, suffix) = (&seq[..start], &seq[d + 1..]);
83    let width: usize = seq[start + 1..d]
84        .trim_start_matches('0')
85        .parse()
86        .unwrap_or(0);
87    name.len() > prefix.len() + suffix.len()
88        && name.starts_with(prefix)
89        && name.ends_with(suffix)
90        && {
91            let digits = &name[prefix.len()..name.len() - suffix.len()];
92            digits.chars().all(|c| c.is_ascii_digit()) && digits.len() >= width
93        }
94}
95
96/// Compare strings with digit runs compared as numbers: `a2 < a10`.
97pub fn natural_cmp(a: &str, b: &str) -> Ordering {
98    let (mut a, mut b) = (a.as_bytes(), b.as_bytes());
99    loop {
100        match (a.first(), b.first()) {
101            (None, None) => return Ordering::Equal,
102            (None, _) => return Ordering::Less,
103            (_, None) => return Ordering::Greater,
104            (Some(x), Some(y)) if x.is_ascii_digit() && y.is_ascii_digit() => {
105                let na = a.iter().take_while(|c| c.is_ascii_digit()).count();
106                let nb = b.iter().take_while(|c| c.is_ascii_digit()).count();
107                let (da, db) = (&a[..na], &b[..nb]);
108                let ta = trim_zeros(da);
109                let tb = trim_zeros(db);
110                let ord = ta
111                    .len()
112                    .cmp(&tb.len())
113                    .then_with(|| ta.cmp(tb))
114                    .then_with(|| na.cmp(&nb));
115                if ord != Ordering::Equal {
116                    return ord;
117                }
118                a = &a[na..];
119                b = &b[nb..];
120            }
121            (Some(x), Some(y)) => {
122                if x != y {
123                    return x.cmp(y);
124                }
125                a = &a[1..];
126                b = &b[1..];
127            }
128        }
129    }
130}
131
132fn trim_zeros(d: &[u8]) -> &[u8] {
133    let i = d.iter().position(|&c| c != b'0').unwrap_or(d.len());
134    &d[i..]
135}
136
137/// A numbered copy of a frame list that ffmpeg can read on every platform.
138///
139/// ffmpeg's glob input (`-pattern_type glob`) does not exist on Windows, and its
140/// numbered input stops at the first gap. So the frames are linked into a temp
141/// folder as `000000.png`, `000001.png`, … (symlinks on Unix, hard links or copies
142/// on Windows), and ffmpeg reads `<temp>/%06d.png`.
143pub struct FrameSequence {
144    _dir: tempfile::TempDir,
145    pattern: PathBuf,
146}
147
148impl FrameSequence {
149    pub fn new(frames: &[PathBuf]) -> std::io::Result<Self> {
150        let dir = tempfile::Builder::new()
151            .prefix("aether-frames-")
152            .tempdir()?;
153        let ext = frames
154            .first()
155            .and_then(|f| f.extension())
156            .and_then(|e| e.to_str())
157            .unwrap_or("png")
158            .to_lowercase();
159        for (i, frame) in frames.iter().enumerate() {
160            let target = dir.path().join(format!("{i:06}.{ext}"));
161            link_or_copy(frame, &target)?;
162        }
163        let pattern = dir.path().join(format!("%06d.{ext}"));
164        Ok(Self { _dir: dir, pattern })
165    }
166
167    /// The `-i` argument for ffmpeg.
168    pub fn pattern(&self) -> &Path {
169        &self.pattern
170    }
171}
172
173fn link_or_copy(src: &Path, dst: &Path) -> std::io::Result<()> {
174    let src = std::fs::canonicalize(src)?;
175    #[cfg(unix)]
176    if std::os::unix::fs::symlink(&src, dst).is_ok() {
177        return Ok(());
178    }
179    if std::fs::hard_link(&src, dst).is_ok() {
180        return Ok(());
181    }
182    std::fs::copy(&src, dst).map(|_| ())
183}
184
185/// Whether a PNG file can hold transparency: an alpha channel (RGBA or gray +
186/// alpha) or a `tRNS` chunk. Returns `false` for anything that is not a readable PNG.
187pub fn png_has_alpha(path: &Path) -> bool {
188    std::fs::read(path).is_ok_and(|bytes| png_bytes_have_alpha(&bytes))
189}
190
191/// [`png_has_alpha`] for PNG bytes, e.g. the start of a PNG stream. The first
192/// few KB are enough: IHDR and tRNS come before the image data.
193pub fn png_bytes_have_alpha(bytes: &[u8]) -> bool {
194    if bytes.len() < 33 || &bytes[..8] != b"\x89PNG\r\n\x1a\n" {
195        return false;
196    }
197    // IHDR is always first: color type 4 = gray + alpha, 6 = RGBA.
198    if matches!(bytes[25], 4 | 6) {
199        return true;
200    }
201    // Palette / gray / RGB images can still mark colors transparent in tRNS,
202    // which comes before the image data.
203    let mut pos = 8;
204    while pos + 8 <= bytes.len() {
205        let len = u32::from_be_bytes([bytes[pos], bytes[pos + 1], bytes[pos + 2], bytes[pos + 3]])
206            as usize;
207        match &bytes[pos + 4..pos + 8] {
208            b"tRNS" => return true,
209            b"IDAT" | b"IEND" => return false,
210            _ => pos += 12 + len,
211        }
212    }
213    false
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    use std::fs::File;
220    use tempfile::tempdir;
221
222    #[test]
223    fn collects_all_png_by_default() -> Result<(), Box<dyn std::error::Error>> {
224        let dir = tempdir()?;
225        File::create(dir.path().join("a.png"))?;
226        File::create(dir.path().join("b.png"))?;
227        let frames = collect_input_frames(dir.path(), None)?;
228        assert_eq!(frames.len(), 2);
229        Ok(())
230    }
231
232    #[test]
233    fn collects_with_custom_pattern() -> Result<(), Box<dyn std::error::Error>> {
234        let dir = tempdir()?;
235        File::create(dir.path().join("scene1_001.png"))?;
236        File::create(dir.path().join("scene2_001.png"))?;
237        let frames = collect_input_frames(dir.path(), Some("scene1_*.png".into()))?;
238        assert_eq!(frames.len(), 1);
239        assert!(frames[0].ends_with("scene1_001.png"));
240        Ok(())
241    }
242
243    #[test]
244    fn numbered_pattern_and_natural_order() -> Result<(), Box<dyn std::error::Error>> {
245        let dir = tempdir()?;
246        for name in [
247            "f_10.png",
248            "f_9.png",
249            "f_1.png",
250            "f_x.png",
251            "._f_2.png",
252            "g_3.png",
253        ] {
254            File::create(dir.path().join(name))?;
255        }
256        let names = |v: Vec<PathBuf>| -> Vec<String> {
257            v.iter()
258                .map(|p| p.file_name().unwrap().to_string_lossy().into_owned())
259                .collect()
260        };
261        let frames = collect_input_frames(dir.path(), Some("f_%d.png".into()))?;
262        assert_eq!(names(frames), ["f_1.png", "f_9.png", "f_10.png"]);
263        let frames = collect_input_frames(dir.path(), None)?;
264        assert_eq!(
265            names(frames),
266            ["f_1.png", "f_9.png", "f_10.png", "f_x.png", "g_3.png"]
267        );
268        Ok(())
269    }
270
271    #[test]
272    fn folder_with_glob_characters() -> Result<(), Box<dyn std::error::Error>> {
273        let dir = tempdir()?;
274        let sub = dir.path().join("renders [final]");
275        std::fs::create_dir(&sub)?;
276        File::create(sub.join("a.png"))?;
277        assert_eq!(collect_input_frames(&sub, None)?.len(), 1);
278        Ok(())
279    }
280
281    #[test]
282    fn natural_cmp_orders_numbers() {
283        assert_eq!(natural_cmp("a2", "a10"), Ordering::Less);
284        assert_eq!(natural_cmp("a010", "a9"), Ordering::Greater);
285        assert_eq!(natural_cmp("a01", "a1"), Ordering::Greater);
286        assert_eq!(natural_cmp("b", "a9"), Ordering::Greater);
287    }
288
289    #[test]
290    fn sequence_links_frames_in_order() -> Result<(), Box<dyn std::error::Error>> {
291        let dir = tempdir()?;
292        let a = dir.path().join("x 1.png");
293        let b = dir.path().join("x 2.png");
294        std::fs::write(&a, b"one")?;
295        std::fs::write(&b, b"two")?;
296        let seq = FrameSequence::new(&[b.clone(), a.clone()])?;
297        let folder = seq.pattern().parent().unwrap();
298        assert_eq!(std::fs::read(folder.join("000000.png"))?, b"two");
299        assert_eq!(std::fs::read(folder.join("000001.png"))?, b"one");
300        assert!(seq.pattern().ends_with("%06d.png"));
301        Ok(())
302    }
303}