Skip to main content

xil_audio/
fx.rs

1//! ffmpeg-backed dialogue treatments. Port of `audio_fx.py`.
2//!
3//! Audio goes through ffmpeg as raw PCM over pipes, exactly as the Python
4//! does: same filter graphs, same argument lists, same trim-or-pad back to
5//! the input length. For a deterministic graph that makes the output bytes
6//! identical; `film` draws its grain from `anoisesrc` with ffmpeg's default
7//! random seed, so no two renders of it — in either language — match.
8//!
9//! Failures degrade: a treatment that cannot run leaves the stem as it was
10//! and warns once per (treatment, reason). Setting `XIL_STRICT_FX` turns
11//! that into an error instead.
12
13use std::collections::HashSet;
14use std::io::Write;
15use std::process::{Command, Stdio};
16use std::sync::Mutex;
17
18use indexmap::IndexMap;
19use sha2::{Digest, Sha256};
20use xil_core::log;
21
22use crate::segment::Segment;
23
24pub const STRICT_ENV_VAR: &str = "XIL_STRICT_FX";
25
26/// FIFO-bounded result cache, keyed on the input bytes and every parameter.
27const CACHE_MAX_ENTRIES: usize = 512;
28
29/// A named treatment, as written in a cast config `filter` field.
30pub struct Treatment {
31    pub name: &'static str,
32    /// `-filter_complex` graph with `{rate}` / `{layout}` placeholders.
33    pub graph: &'static str,
34    pub codec: Option<&'static str>,
35    pub container: Option<&'static str>,
36    pub codec_rate: Option<u32>,
37}
38
39pub const FILM: Treatment = Treatment {
40    name: "film",
41    graph: concat!(
42        "[0:a]",
43        "highpass=f=90:poles=2,",
44        "lowpass=f=5500:poles=2,",
45        "lowshelf=f=180:g=3,",
46        "equalizer=f=350:w=1.2:t=q:g=2,",
47        "equalizer=f=2800:w=1.6:t=q:g=-6,",
48        "highshelf=f=7000:g=-4,",
49        "acompressor=threshold=-24dB:ratio=3.5:attack=12:release=280:makeup=2:knee=6,",
50        "volume=12dB,",
51        "asoftclip=type=tanh:param=1:oversample=4,",
52        "volume=-12dB,",
53        "volume=8.4dB",
54        "[v];",
55        "anoisesrc=color=pink:amplitude=0.018:sample_rate={rate},",
56        "aformat=channel_layouts={layout}[n];",
57        "[v][n]amix=inputs=2:duration=first:normalize=0[out]",
58    ),
59    codec: None,
60    container: None,
61    codec_rate: None,
62};
63
64pub const SPEAKERPHONE: Treatment = Treatment {
65    name: "speakerphone",
66    graph: concat!(
67        "[0:a]",
68        "highpass=f=350:poles=2,",
69        "highpass=f=350:poles=2,",
70        "lowpass=f=3400:poles=2,",
71        "lowpass=f=3400:poles=2,",
72        "equalizer=f=700:w=1.0:t=q:g=-4,",
73        "equalizer=f=1800:w=1.1:t=q:g=5,",
74        "acompressor=threshold=-24dB:ratio=6:attack=5:release=120:makeup=3:knee=2,",
75        "volume=8dB,",
76        "asoftclip=type=atan:param=1:oversample=4,",
77        "volume=-8dB,",
78        "aecho=0.9:0.9:55:0.22,",
79        "volume=13.6dB",
80        "[out]",
81    ),
82    codec: None,
83    container: None,
84    codec_rate: None,
85};
86
87pub const PHONE: Treatment = Treatment {
88    name: "phone",
89    graph: concat!(
90        "[0:a]",
91        "highpass=f=300:poles=2,",
92        "highpass=f=300:poles=2,",
93        "lowpass=f=3400:poles=2,",
94        "lowpass=f=3400:poles=2,",
95        "equalizer=f=500:w=1.0:t=q:g=-4,",
96        "equalizer=f=1700:w=1.2:t=q:g=6,",
97        "acompressor=threshold=-22dB:ratio=8:attack=3:release=90:makeup=3:knee=2,",
98        "volume=6dB,",
99        "asoftclip=type=atan:param=1:oversample=4,",
100        "volume=-6dB,",
101        "volume=10dB",
102        "[out]",
103    ),
104    codec: Some("libgsm"),
105    container: Some("gsm"),
106    codec_rate: Some(8000),
107};
108
109/// `TREATMENTS`, in the order the Python dict is built.
110pub const TREATMENTS: [&Treatment; 3] = [&FILM, &SPEAKERPHONE, &PHONE];
111
112pub fn treatment(name: &str) -> Option<&'static Treatment> {
113    TREATMENTS.iter().copied().find(|t| t.name == name)
114}
115
116/// `AudioFxError` — only raised under `XIL_STRICT_FX`.
117#[derive(Debug, thiserror::Error)]
118#[error("{0}")]
119pub struct AudioFxError(pub String);
120
121type CacheKey = (
122    [u8; 32],
123    u32,
124    usize,
125    usize,
126    String,
127    bool,
128    Option<String>,
129    Option<String>,
130    Option<u32>,
131);
132
133struct State {
134    cache: IndexMap<CacheKey, Vec<u8>>,
135    warned: HashSet<(String, String)>,
136    encoder_probe: std::collections::HashMap<String, bool>,
137}
138
139static STATE: Mutex<Option<State>> = Mutex::new(None);
140
141fn with_state<R>(f: impl FnOnce(&mut State) -> R) -> R {
142    let mut guard = STATE.lock().unwrap_or_else(|p| p.into_inner());
143    let state = guard.get_or_insert_with(|| State {
144        cache: IndexMap::new(),
145        warned: HashSet::new(),
146        encoder_probe: std::collections::HashMap::new(),
147    });
148    f(state)
149}
150
151/// Python `repr()` of a `str`.
152pub fn py_repr(s: &str) -> String {
153    let quote = if s.contains('\'') && !s.contains('"') {
154        '"'
155    } else {
156        '\''
157    };
158    let mut out = String::new();
159    out.push(quote);
160    for c in s.chars() {
161        match c {
162            '\\' => out.push_str("\\\\"),
163            '\n' => out.push_str("\\n"),
164            '\r' => out.push_str("\\r"),
165            '\t' => out.push_str("\\t"),
166            c if c == quote => {
167                out.push('\\');
168                out.push(c);
169            }
170            c if (c as u32) < 0x20 || c as u32 == 0x7f => {
171                out.push_str(&format!("\\x{:02x}", c as u32))
172            }
173            c => out.push(c),
174        }
175    }
176    out.push(quote);
177    out
178}
179
180/// `_warn_once(label, reason, message)` — the message is already formatted.
181pub fn warn_once(label: &str, reason: &str, message: &str) {
182    let first = with_state(|s| s.warned.insert((label.to_string(), reason.to_string())));
183    if first {
184        log::warning(message);
185    }
186}
187
188/// `_fail`: warn once with the fallback spelled out, or error when strict.
189fn fail(label: &str, reason: &str, message: &str, consequence: &str) -> Result<(), AudioFxError> {
190    if std::env::var_os(STRICT_ENV_VAR).is_some_and(|v| !v.is_empty()) {
191        return Err(AudioFxError(message.to_string()));
192    }
193    warn_once(label, reason, &format!("{message} — {consequence}"));
194    Ok(())
195}
196
197const UNTREATED: &str = "leaving audio untreated";
198const CODEC_FALLBACK: &str = "keeping the filtered audio without codec character";
199
200/// ffmpeg's raw format name for a sample width.
201fn raw_format(width: usize) -> Option<&'static str> {
202    match width {
203        1 => Some("u8"),
204        2 => Some("s16le"),
205        3 => Some("s24le"),
206        4 => Some("s32le"),
207        _ => None,
208    }
209}
210
211/// Run a command with `input` on stdin, collecting stdout and stderr, as
212/// `subprocess.run(..., input=..., capture_output=True)` does.
213fn run_piped(args: &[String], input: &[u8]) -> std::io::Result<std::process::Output> {
214    let mut child = Command::new(&args[0])
215        .args(&args[1..])
216        .stdin(Stdio::piped())
217        .stdout(Stdio::piped())
218        .stderr(Stdio::piped())
219        .spawn()?;
220    let mut stdin = child.stdin.take().expect("piped stdin");
221    let input = input.to_vec();
222    // Write from another thread so a full stdout pipe cannot deadlock us.
223    let writer = std::thread::spawn(move || {
224        let _ = stdin.write_all(&input);
225    });
226    let out = child.wait_with_output();
227    let _ = writer.join();
228    out
229}
230
231fn last_stderr_line(stderr: &[u8]) -> String {
232    String::from_utf8_lossy(stderr)
233        .trim()
234        .lines()
235        .last()
236        .map(str::to_string)
237        .unwrap_or_else(|| "no stderr output".to_string())
238}
239
240fn s(v: &str) -> String {
241    v.to_string()
242}
243
244/// `_codec_round_trip`: encode through a lossy codec and decode straight back.
245fn codec_round_trip(
246    raw: Vec<u8>,
247    segment: &Segment,
248    raw_fmt: &str,
249    label: &str,
250    codec: &str,
251    container: Option<&str>,
252    codec_rate: Option<u32>,
253) -> Result<Vec<u8>, AudioFxError> {
254    let rate = segment.frame_rate.to_string();
255    let channels = segment.channels.to_string();
256    let mux = container.unwrap_or(codec);
257    let encode: Vec<String> = [
258        "ffmpeg",
259        "-hide_banner",
260        "-loglevel",
261        "error",
262        "-nostdin",
263        "-y",
264        "-f",
265        raw_fmt,
266        "-ar",
267        &rate,
268        "-ac",
269        &channels,
270        "-i",
271        "pipe:0",
272        "-ar",
273    ]
274    .iter()
275    .map(|x| s(x))
276    .chain([codec_rate.unwrap_or(segment.frame_rate).to_string()])
277    .chain(
278        ["-ac", "1", "-c:a", codec, "-f", mux, "pipe:1"]
279            .iter()
280            .map(|x| s(x)),
281    )
282    .collect();
283    let decode: Vec<String> = [
284        "ffmpeg",
285        "-hide_banner",
286        "-loglevel",
287        "error",
288        "-nostdin",
289        "-y",
290        "-f",
291        mux,
292        "-i",
293        "pipe:0",
294        "-vn",
295        "-sn",
296        "-dn",
297        "-f",
298        raw_fmt,
299        "-ar",
300        &rate,
301        "-ac",
302        &channels,
303        "pipe:1",
304    ]
305    .iter()
306    .map(|x| s(x))
307    .collect();
308
309    let stage = |name: &str, args: &[String], input: &[u8]| -> Result<Vec<u8>, (String, String)> {
310        match run_piped(args, input) {
311            Ok(o) if o.status.success() && !o.stdout.is_empty() => Ok(o.stdout),
312            Ok(o) => Err((name.to_string(), last_stderr_line(&o.stderr))),
313            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
314                Err(("missing".into(), String::new()))
315            }
316            Err(e) => Err(("oserror".into(), e.to_string())),
317        }
318    };
319
320    let result = stage("encode", &encode, &raw).and_then(|enc| stage("decode", &decode, &enc));
321    match result {
322        Ok(dec) => Ok(dec),
323        Err((kind, detail)) => {
324            let message = match kind.as_str() {
325                "missing" => format!("ffmpeg not found for codec stage of {}", py_repr(label)),
326                "oserror" => format!(
327                    "codec stage failed for treatment {}: {detail}",
328                    py_repr(label)
329                ),
330                stage_name => format!(
331                    "codec {codec} unavailable for treatment {} ({stage_name}): {detail}",
332                    py_repr(label)
333                ),
334            };
335            fail(label, "codec", &message, CODEC_FALLBACK)?;
336            Ok(raw)
337        }
338    }
339}
340
341/// `_fit_length`: trim or pad with silence to exactly `target` bytes.
342fn fit_length(mut raw: Vec<u8>, target: usize, width: usize) -> Vec<u8> {
343    let fill = if width == 1 { 0x80 } else { 0x00 };
344    raw.resize(target, fill);
345    raw
346}
347
348/// `run_ffmpeg_filter(segment, graph, label=..., codec=..., ...)` with
349/// `preserve_length=True`.
350pub fn run_ffmpeg_filter(
351    segment: &Segment,
352    graph: &str,
353    label: &str,
354    codec: Option<&str>,
355    container: Option<&str>,
356    codec_rate: Option<u32>,
357) -> Result<Segment, AudioFxError> {
358    let Some(raw_fmt) = raw_format(segment.sample_width) else {
359        fail(
360            label,
361            "sample_width",
362            &format!(
363                "Unsupported sample width {} bytes for treatment {}",
364                segment.sample_width,
365                py_repr(label)
366            ),
367            UNTREATED,
368        )?;
369        return Ok(segment.clone());
370    };
371    let layout = if segment.channels == 1 {
372        "mono"
373    } else {
374        "stereo"
375    };
376    let resolved = graph
377        .replace("{rate}", &segment.frame_rate.to_string())
378        .replace("{layout}", layout);
379
380    let key: CacheKey = (
381        Sha256::digest(&segment.data).into(),
382        segment.frame_rate,
383        segment.channels,
384        segment.sample_width,
385        resolved.clone(),
386        true,
387        codec.map(str::to_string),
388        container.map(str::to_string),
389        codec_rate,
390    );
391    if let Some(hit) = with_state(|s| s.cache.get(&key).cloned()) {
392        return Ok(Segment::new(
393            hit,
394            segment.sample_width,
395            segment.frame_rate,
396            segment.channels,
397        ));
398    }
399
400    let rate = segment.frame_rate.to_string();
401    let channels = segment.channels.to_string();
402    let cmd: Vec<String> = [
403        "ffmpeg",
404        "-hide_banner",
405        "-loglevel",
406        "error",
407        "-nostdin",
408        "-y",
409        "-f",
410        raw_fmt,
411        "-ar",
412        &rate,
413        "-ac",
414        &channels,
415        "-i",
416        "pipe:0",
417        "-filter_complex",
418        &resolved,
419        "-map",
420        "[out]",
421        "-vn",
422        "-sn",
423        "-dn",
424        "-f",
425        raw_fmt,
426        "-ar",
427        &rate,
428        "-ac",
429        &channels,
430        "pipe:1",
431    ]
432    .iter()
433    .map(|x| s(x))
434    .collect();
435
436    let proc = match run_piped(&cmd, &segment.data) {
437        Ok(p) => p,
438        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
439            fail(
440                label,
441                "missing",
442                &format!("ffmpeg not found for treatment {}", py_repr(label)),
443                UNTREATED,
444            )?;
445            return Ok(segment.clone());
446        }
447        Err(e) => {
448            fail(
449                label,
450                "oserror",
451                &format!("ffmpeg failed for treatment {}: {e}", py_repr(label)),
452                UNTREATED,
453            )?;
454            return Ok(segment.clone());
455        }
456    };
457    if !proc.status.success() || proc.stdout.is_empty() {
458        let rc = proc
459            .status
460            .code()
461            .map(|c| c.to_string())
462            .unwrap_or_else(|| signal_code(&proc.status));
463        fail(
464            label,
465            "returncode",
466            &format!(
467                "ffmpeg treatment {} failed (rc={rc}): {}",
468                py_repr(label),
469                last_stderr_line(&proc.stderr)
470            ),
471            UNTREATED,
472        )?;
473        return Ok(segment.clone());
474    }
475
476    let mut out = proc.stdout;
477    if let Some(codec) = codec {
478        out = codec_round_trip(out, segment, raw_fmt, label, codec, container, codec_rate)?;
479    }
480    let out = fit_length(out, segment.data.len(), segment.sample_width);
481
482    with_state(|s| {
483        if s.cache.len() >= CACHE_MAX_ENTRIES {
484            s.cache.shift_remove_index(0);
485        }
486        s.cache.insert(key, out.clone());
487    });
488    Ok(Segment::new(
489        out,
490        segment.sample_width,
491        segment.frame_rate,
492        segment.channels,
493    ))
494}
495
496/// Python reports a signal death as a negative return code.
497fn signal_code(status: &std::process::ExitStatus) -> String {
498    #[cfg(unix)]
499    {
500        use std::os::unix::process::ExitStatusExt;
501        if let Some(sig) = status.signal() {
502            return format!("-{sig}");
503        }
504    }
505    let _ = status;
506    "None".to_string()
507}
508
509/// `apply_treatment(segment, name)`.
510pub fn apply_treatment(segment: &Segment, name: &str) -> Result<Segment, AudioFxError> {
511    let Some(t) = treatment(name) else {
512        let mut known: Vec<&str> = TREATMENTS.iter().map(|t| t.name).collect();
513        known.sort();
514        warn_once(
515            name,
516            "unknown",
517            &format!(
518                "Unknown ffmpeg treatment {} — known treatments: {}",
519                py_repr(name),
520                known.join(", ")
521            ),
522        );
523        return Ok(segment.clone());
524    };
525    run_ffmpeg_filter(segment, t.graph, t.name, t.codec, t.container, t.codec_rate)
526}
527
528/// `encoder_available(name)`: does `ffmpeg -encoders` list it as a token?
529pub fn encoder_available(name: &str) -> bool {
530    match Command::new("ffmpeg")
531        .args(["-hide_banner", "-encoders"])
532        .output()
533    {
534        Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout)
535            .split_whitespace()
536            .any(|t| t == name),
537        _ => false,
538    }
539}
540
541/// `missing_codecs(names)`: treatment → encoder, for encoders this ffmpeg lacks.
542pub fn missing_codecs<'a>(names: impl IntoIterator<Item = &'a str>) -> IndexMap<String, String> {
543    let mut uniq: Vec<&str> = names.into_iter().collect();
544    uniq.sort();
545    uniq.dedup();
546    let mut missing = IndexMap::new();
547    for name in uniq {
548        let Some(t) = treatment(name) else { continue };
549        let Some(codec) = t.codec else { continue };
550        let available = with_state(|s| s.encoder_probe.get(codec).copied());
551        let available = match available {
552            Some(a) => a,
553            None => {
554                let a = encoder_available(codec);
555                with_state(|s| s.encoder_probe.insert(codec.to_string(), a));
556                a
557            }
558        };
559        if !available {
560            missing.insert(name.to_string(), codec.to_string());
561        }
562    }
563    missing
564}
565
566#[cfg(test)]
567mod tests {
568    use super::*;
569
570    #[test]
571    fn repr_matches_python() {
572        assert_eq!(py_repr("phone"), "'phone'");
573        assert_eq!(py_repr("it's"), "\"it's\"");
574        assert_eq!(py_repr("a\\b"), "'a\\\\b'");
575    }
576
577    #[test]
578    fn fit_length_trims_and_pads() {
579        assert_eq!(fit_length(vec![1, 2, 3, 4], 2, 2), vec![1, 2]);
580        assert_eq!(fit_length(vec![1, 2], 4, 2), vec![1, 2, 0, 0]);
581        assert_eq!(fit_length(vec![1], 3, 1), vec![1, 0x80, 0x80]);
582    }
583
584    #[test]
585    fn treatments_keep_length() {
586        if !crate::ffmpeg::available() {
587            return;
588        }
589        let data: Vec<u8> = (0..4410i32)
590            .map(|i| ((i as f64 * 0.05).sin() * 8000.0) as i16)
591            .flat_map(|v| v.to_le_bytes())
592            .collect();
593        let seg = Segment::new(data, 2, 44100, 1);
594        let out = apply_treatment(&seg, "speakerphone").unwrap();
595        assert_eq!(out.data.len(), seg.data.len());
596        assert_ne!(out.data, seg.data);
597        let again = apply_treatment(&seg, "speakerphone").unwrap();
598        assert_eq!(out, again, "second call is served from the cache");
599    }
600
601    /// `film` mixes in `anoisesrc` grain with ffmpeg's random seed, so it
602    /// cannot be pinned to bytes the way the other treatments are. Check the
603    /// contract instead: stereo in, stereo out, same length, audibly changed,
604    /// and grain present even where the input is silent.
605    #[test]
606    fn film_keeps_length_and_adds_grain() {
607        if !crate::ffmpeg::available() {
608            return;
609        }
610        let mut data: Vec<u8> = (0..8820i32)
611            .map(|i| ((i as f64 * 0.03).sin() * 6000.0) as i16)
612            .flat_map(|v| v.to_le_bytes())
613            .collect();
614        data.extend(vec![0u8; 8820 * 2]);
615        let seg = Segment::new(data, 2, 44100, 2);
616        let out = apply_treatment(&seg, "film").unwrap();
617        assert_eq!(
618            (out.channels, out.frame_rate, out.data.len()),
619            (2, 44100, seg.data.len())
620        );
621        assert_ne!(out.data, seg.data);
622        let tail = &out.data[out.data.len() - 4000..];
623        assert!(
624            tail.iter().any(|&b| b != 0),
625            "the silent tail carries grain"
626        );
627    }
628
629    #[test]
630    fn unknown_treatment_passes_through() {
631        let seg = Segment::new(vec![1, 0, 2, 0], 2, 8000, 1);
632        assert_eq!(apply_treatment(&seg, "robot").unwrap(), seg);
633    }
634}