1use 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
26const CACHE_MAX_ENTRIES: usize = 512;
28
29pub struct Treatment {
31 pub name: &'static str,
32 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
109pub 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#[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
151pub 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
180pub 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
188fn 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
200fn 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
211fn 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 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
244fn 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
341fn 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
348pub 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
496fn 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
509pub 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
528pub 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
541pub 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 #[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}