1use crate::buffer::AudioBuffer;
5use crate::error::{Error, ErrorKind, Result};
6use serde::Serialize;
7use serde_json::{Map, Value};
8use std::path::{Path, PathBuf};
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::sync::Arc;
11
12#[derive(Clone, Debug, Default, Serialize)]
14pub struct SourceInfo {
15 pub path: String,
16 #[serde(skip_serializing_if = "Option::is_none")]
17 pub container: Option<String>,
18 #[serde(skip_serializing_if = "Option::is_none")]
19 pub codec: Option<String>,
20 pub duration_s: f64,
21 pub sample_rate: u32,
22 pub channels: usize,
23 #[serde(skip_serializing_if = "Option::is_none")]
24 pub bit_depth: Option<u16>,
25 #[serde(skip_serializing_if = "Option::is_none")]
26 pub bitrate: Option<u32>,
27 #[serde(skip_serializing_if = "std::collections::BTreeMap::is_empty")]
28 pub tags: std::collections::BTreeMap<String, String>,
29}
30
31#[derive(Clone, Debug, Default, Serialize)]
33pub struct EncodeOptions {
34 pub format: Option<String>,
37 pub bit_depth: Option<u16>,
38 pub float: bool,
40 pub bitrate: Option<u32>,
41 pub vbr: bool,
42 pub dither: Option<String>,
44 pub seed: u64,
45 pub tags: std::collections::BTreeMap<String, String>,
46}
47
48pub trait AudioIo: Send + Sync {
50 fn read(&self, path: &Path) -> Result<(AudioBuffer, SourceInfo)>;
51 fn read_bytes(&self, bytes: &[u8], format_hint: Option<&str>) -> Result<(AudioBuffer, SourceInfo)>;
53 fn probe(&self, path: &Path) -> Result<SourceInfo>;
54 fn write(&self, path: &Path, buf: &AudioBuffer, opts: &EncodeOptions) -> Result<()>;
55 fn write_bytes(&self, buf: &AudioBuffer, opts: &EncodeOptions) -> Result<Vec<u8>>;
56 fn can_encode(&self, ext: &str) -> bool;
58 fn open_reader(&self, _path: &Path) -> Result<Option<Box<dyn AudioReader>>> {
61 Ok(None)
62 }
63 fn open_writer(
67 &self,
68 _path: &Path,
69 _sample_rate: u32,
70 _channels: usize,
71 _frames: u64,
72 _opts: &EncodeOptions,
73 ) -> Result<Option<Box<dyn AudioWriter>>> {
74 Ok(None)
75 }
76}
77
78pub trait AudioReader {
81 fn info(&self) -> &SourceInfo;
83 fn frames_hint(&self) -> Option<u64>;
85 fn next_chunk(&mut self) -> Result<Option<AudioBuffer>>;
86}
87
88pub trait AudioWriter {
90 fn write(&mut self, planes: &[&[f32]]) -> Result<()>;
92 fn finish(self: Box<Self>) -> Result<()>;
94}
95
96pub trait Progress: Send + Sync {
98 fn stage(&self, _stage: &str, _index: usize, _of: usize) {}
99 fn progress(&self, _stage: &str, _fraction: f64) {}
100 fn output(&self, _path: &str, _role: Option<&str>) {}
101 fn warning(&self, _code: &str, _message: &str) {}
102 fn log(&self, _message: &str) {}
103}
104
105pub struct NoProgress;
106impl Progress for NoProgress {}
107
108#[derive(Clone, Debug, Serialize)]
110pub struct Globals {
111 pub output: Option<String>,
112 pub outdir: Option<PathBuf>,
113 pub format: Option<String>,
114 pub input_format: Option<String>,
115 pub overwrite: bool,
116 pub in_place: bool,
117 pub sidecar: bool,
118 pub dry_run: bool,
119 pub explain: bool,
120 pub force: bool,
121 pub sample_rate: Option<u32>,
122 pub bit_depth: Option<u16>,
123 pub channels: Option<usize>,
124 pub bitrate: Option<u32>,
125 pub seed: u64,
126 pub threads: usize,
127 pub device: String,
128 pub deterministic: bool,
129 pub offline: bool,
130 pub yes: bool,
131 pub bpm: Option<f64>,
132 pub sound_paths: Vec<PathBuf>,
133 pub max_memory: u64,
134 pub allow: Vec<PathBuf>,
136 pub config_path: Option<PathBuf>,
137}
138
139impl Default for Globals {
140 fn default() -> Self {
141 Globals {
142 output: None,
143 outdir: None,
144 format: None,
145 input_format: None,
146 overwrite: false,
147 in_place: false,
148 sidecar: true,
149 dry_run: false,
150 explain: false,
151 force: false,
152 sample_rate: None,
153 bit_depth: None,
154 channels: None,
155 bitrate: None,
156 seed: 0,
157 threads: 0,
158 device: "auto".into(),
159 deterministic: false,
160 offline: false,
161 yes: false,
162 bpm: None,
163 sound_paths: Vec::new(),
164 max_memory: 2 * 1024 * 1024 * 1024,
165 allow: Vec::new(),
166 config_path: None,
167 }
168 }
169}
170
171#[derive(Clone, Debug, Serialize)]
172pub struct OutputRecord {
173 pub path: String,
174 #[serde(skip_serializing_if = "Option::is_none")]
175 pub role: Option<String>,
176 #[serde(skip_serializing_if = "Option::is_none")]
177 pub duration_s: Option<f64>,
178 #[serde(skip_serializing_if = "Option::is_none")]
179 pub sample_rate: Option<u32>,
180 #[serde(skip_serializing_if = "Option::is_none")]
181 pub channels: Option<usize>,
182 #[serde(skip_serializing_if = "Map::is_empty")]
183 pub extra: Map<String, Value>,
184}
185
186#[derive(Clone, Debug, Serialize)]
187pub struct WarningRecord {
188 pub code: String,
189 pub message: String,
190}
191
192#[derive(Clone, Debug, Serialize, PartialEq)]
193pub struct ModelRecord {
194 pub name: String,
195 pub version: String,
196 pub sha256: String,
197}
198
199#[derive(Clone, Debug, Serialize)]
200pub struct InputRecord {
201 pub path: String,
202 #[serde(skip_serializing_if = "Option::is_none")]
203 pub duration_s: Option<f64>,
204 #[serde(skip_serializing_if = "Option::is_none")]
205 pub sample_rate: Option<u32>,
206 #[serde(skip_serializing_if = "Option::is_none")]
207 pub channels: Option<usize>,
208}
209
210#[derive(Clone, Default)]
212pub struct CancelToken(Arc<AtomicBool>);
213
214impl CancelToken {
215 pub fn cancel(&self) {
216 self.0.store(true, Ordering::SeqCst)
217 }
218 pub fn is_cancelled(&self) -> bool {
219 self.0.load(Ordering::SeqCst)
220 }
221 pub fn check(&self) -> Result<()> {
222 if self.is_cancelled() {
223 Err(Error::cancelled("cancelled"))
224 } else {
225 Ok(())
226 }
227 }
228}
229
230pub struct OpCtx<'a> {
232 pub command: String,
234 pub inputs: Vec<PathBuf>,
235 pub globals: Globals,
236 pub io: &'a dyn AudioIo,
237 pub progress: &'a dyn Progress,
238 pub cancel: CancelToken,
239 pub registry: Option<&'a crate::registry::Registry>,
241 pub outputs: Vec<OutputRecord>,
242 pub input_records: Vec<InputRecord>,
243 pub warnings: Vec<WarningRecord>,
244 pub metrics: Map<String, Value>,
245 pub notes: Map<String, Value>,
247 pub models: Vec<ModelRecord>,
248 pub plugins: Vec<Value>,
249 pub deterministic: Option<bool>,
251 pub stdin: Option<Vec<u8>>,
253 pub stdout: Option<Vec<u8>>,
255 pub index: usize,
257 pub temp_files: Vec<PathBuf>,
259 pub result: Option<Value>,
261}
262
263impl<'a> OpCtx<'a> {
264 pub fn new(
265 command: &str,
266 inputs: Vec<PathBuf>,
267 globals: Globals,
268 io: &'a dyn AudioIo,
269 progress: &'a dyn Progress,
270 ) -> Self {
271 OpCtx {
272 command: command.to_string(),
273 inputs,
274 globals,
275 io,
276 progress,
277 cancel: CancelToken::default(),
278 registry: None,
279 outputs: Vec::new(),
280 input_records: Vec::new(),
281 warnings: Vec::new(),
282 metrics: Map::new(),
283 notes: Map::new(),
284 models: Vec::new(),
285 plugins: Vec::new(),
286 deterministic: None,
287 stdin: None,
288 stdout: None,
289 index: 0,
290 temp_files: Vec::new(),
291 result: None,
292 }
293 }
294
295 pub fn child(&self, command: &str, inputs: Vec<PathBuf>) -> OpCtx<'a> {
297 let mut c = OpCtx::new(command, inputs, self.globals.clone(), self.io, self.progress);
298 c.cancel = self.cancel.clone();
299 c.registry = self.registry;
300 c
301 }
302
303 pub fn absorb(&mut self, child: OpCtx<'a>) {
305 self.warnings.extend(child.warnings);
306 for m in child.models {
307 if !self.models.contains(&m) {
308 self.models.push(m);
309 }
310 }
311 self.plugins.extend(child.plugins);
312 if child.deterministic == Some(false) {
313 self.deterministic = Some(false);
314 }
315 }
316
317 pub fn bpm(&self) -> Option<f64> {
318 self.globals.bpm
319 }
320
321 pub fn seed(&self) -> u64 {
322 self.globals.seed
323 }
324
325 pub fn input(&self, i: usize) -> Result<&Path> {
326 self.inputs
327 .get(i)
328 .map(|p| p.as_path())
329 .ok_or_else(|| Error::usage(format!("'{}' needs at least {} input(s)", self.command, i + 1)))
330 }
331
332 pub fn check_allowed(&self, path: &Path) -> Result<()> {
334 if self.globals.allow.is_empty() {
335 return Ok(());
336 }
337 let abs = absolutize(path);
338 if self.globals.allow.iter().any(|a| abs.starts_with(absolutize(a))) {
339 Ok(())
340 } else {
341 Err(Error::new(
342 ErrorKind::Usage,
343 "path_not_allowed",
344 format!("{} is outside the allowed directories", path.display()),
345 )
346 .with_hint("serve with --allow <dir> to permit it"))
347 }
348 }
349
350 pub fn read_audio(&mut self, i: usize) -> Result<AudioBuffer> {
352 let path = self.input(i)?.to_path_buf();
353 self.read_audio_path(&path)
354 }
355
356 pub fn read_audio_path(&mut self, path: &Path) -> Result<AudioBuffer> {
357 let (buf, info) = if path.as_os_str() == "-" {
358 let bytes = self.stdin.take().ok_or_else(|| Error::input("stdin was already consumed or is empty"))?;
359 let r = self.io.read_bytes(&bytes, self.globals.input_format.as_deref())?;
360 self.stdin = Some(bytes);
361 r
362 } else {
363 self.check_allowed(path)?;
364 if !path.exists() {
365 return Err(Error::new(
366 ErrorKind::InputMissing,
367 "input_not_found",
368 format!("input not found: {}", path.display()),
369 ));
370 }
371 self.io.read(path)?
372 };
373 self.input_records.push(InputRecord {
374 path: info.path.clone(),
375 duration_s: Some(info.duration_s),
376 sample_rate: Some(info.sample_rate),
377 channels: Some(info.channels),
378 });
379 Ok(buf)
380 }
381
382 pub fn open_audio_stream(&mut self, i: usize) -> Result<Option<Box<dyn AudioReader>>> {
385 let path = self.input(i)?.to_path_buf();
386 if path.as_os_str() == "-" {
387 return Ok(None);
388 }
389 self.check_allowed(&path)?;
390 if !path.exists() {
391 return Err(Error::new(
392 ErrorKind::InputMissing,
393 "input_not_found",
394 format!("input not found: {}", path.display()),
395 ));
396 }
397 self.io.open_reader(&path)
398 }
399
400 pub fn record_stream_input(&mut self, info: &SourceInfo, sample_rate: u32, channels: usize, frames: u64) {
402 self.input_records.push(InputRecord {
403 path: info.path.clone(),
404 duration_s: Some(frames as f64 / sample_rate.max(1) as f64),
405 sample_rate: Some(sample_rate),
406 channels: Some(channels),
407 });
408 }
409
410 pub fn note_input(&mut self, path: &Path) {
412 self.input_records.push(InputRecord {
413 path: path.display().to_string(),
414 duration_s: None,
415 sample_rate: None,
416 channels: None,
417 });
418 }
419
420 pub fn output_path(&self, role: Option<&str>, ext: &str) -> Result<PathBuf> {
428 let input = self.inputs.first().filter(|p| p.as_os_str() != "-");
429 let name = input
430 .and_then(|p| p.file_stem())
431 .map(|s| s.to_string_lossy().to_string())
432 .unwrap_or_else(|| "ferraus".to_string());
433 let in_ext = input.and_then(|p| p.extension()).map(|e| e.to_string_lossy().to_lowercase());
434 let cmd = self.command.replace(' ', "-");
435 let ext = if ext.is_empty() { in_ext.clone().unwrap_or_else(|| "wav".into()) } else { ext.to_string() };
436
437 if self.globals.in_place && role.is_none() {
438 if let Some(p) = input {
439 return Ok(p.clone());
440 }
441 }
442 let p = if let Some(o) = &self.globals.output {
443 if o == "-" {
444 return Ok(PathBuf::from("-"));
445 }
446 let has_stem = o.contains("{stem}");
450 let index_takes_role = role.is_some() && !has_stem && o.contains("{index}");
451 let index = if index_takes_role { role.unwrap_or_default().to_string() } else { self.index.to_string() };
452 let mut s =
453 o.replace("{name}", &name).replace("{index}", &index).replace("{command}", &cmd).replace("{ext}", &ext);
454 s = s.replace("{stem}", role.unwrap_or("out"));
455 let mut p = PathBuf::from(s);
456 if let (Some(role), false, false) = (role, has_stem, index_takes_role) {
457 let stem = p.file_stem().map(|s| s.to_string_lossy().to_string()).unwrap_or_default();
458 let e = p.extension().map(|e| e.to_string_lossy().to_string()).unwrap_or_else(|| ext.clone());
459 p.set_file_name(format!("{stem}.{role}.{e}"));
460 }
461 if let Some(dir) = &self.globals.outdir {
462 if p.is_relative() {
463 p = dir.join(p);
464 }
465 }
466 p
467 } else {
468 let file = match role {
469 Some(r) => format!("{name}.{r}.{ext}"),
470 None => format!("{name}.{cmd}.{ext}"),
471 };
472 match (&self.globals.outdir, input.and_then(|p| p.parent())) {
473 (Some(dir), _) => dir.join(file),
474 (None, Some(parent)) => parent.join(file),
475 (None, None) => PathBuf::from(file),
476 }
477 };
478 Ok(p)
479 }
480
481 pub fn prepare_output(&self, path: &Path) -> Result<()> {
483 if path.as_os_str() == "-" {
484 return Ok(());
485 }
486 self.check_allowed(path)?;
487 let is_input = self.inputs.iter().any(|i| i == path);
488 if path.exists() && !self.globals.overwrite && !self.globals.force && !(self.globals.in_place && is_input) {
489 return Err(Error::new(ErrorKind::Usage, "output_exists", format!("{} already exists", path.display()))
490 .with_hint("pass --overwrite (or --force) to replace it, or choose another -o path"));
491 }
492 if is_input && !self.globals.in_place {
493 return Err(Error::usage(format!("output {} is the same as an input; use --in-place", path.display())));
494 }
495 if let Some(parent) = path.parent() {
496 if !parent.as_os_str().is_empty() && !self.globals.dry_run {
497 std::fs::create_dir_all(parent)?;
498 }
499 }
500 Ok(())
501 }
502
503 pub fn encode_options(&self, path: &Path) -> EncodeOptions {
505 let format = if path.as_os_str() == "-" {
506 Some(self.globals.format.clone().unwrap_or_else(|| "wav".into()))
507 } else {
508 self.globals.format.clone().or_else(|| path.extension().map(|e| e.to_string_lossy().to_lowercase()))
509 };
510 EncodeOptions {
511 format,
512 bit_depth: self.globals.bit_depth,
513 float: false,
514 bitrate: self.globals.bitrate,
515 vbr: true,
516 dither: None,
517 seed: self.globals.seed,
518 tags: Default::default(),
519 }
520 }
521
522 pub fn write_audio(&mut self, path: &Path, buf: &AudioBuffer, role: Option<&str>) -> Result<()> {
525 let opts = self.encode_options(path);
526 self.write_audio_with(path, buf, role, &opts)
527 }
528
529 pub fn write_audio_with(
530 &mut self,
531 path: &Path,
532 buf: &AudioBuffer,
533 role: Option<&str>,
534 opts: &EncodeOptions,
535 ) -> Result<()> {
536 self.prepare_output(path)?;
537 if !self.globals.dry_run {
538 if path.as_os_str() == "-" {
539 let bytes = self.io.write_bytes(buf, opts)?;
540 self.stdout.get_or_insert_with(Vec::new).extend(bytes);
541 } else {
542 if self.globals.in_place && path.exists() {
543 let bak = PathBuf::from(format!("{}.bak", path.display()));
544 std::fs::copy(path, &bak)?;
545 }
546 let tmp = tmp_path(path);
548 self.io.write(&tmp, buf, opts).inspect_err(|_| {
549 let _ = std::fs::remove_file(&tmp);
550 })?;
551 std::fs::rename(&tmp, path)?;
552 }
553 }
554 self.outputs.push(OutputRecord {
555 path: path.display().to_string(),
556 role: role.map(|r| r.to_string()),
557 duration_s: Some(buf.duration_s()),
558 sample_rate: Some(buf.sample_rate),
559 channels: Some(buf.channels()),
560 extra: Map::new(),
561 });
562 self.progress.output(&path.display().to_string(), role);
563 Ok(())
564 }
565
566 #[allow(clippy::too_many_arguments)]
571 pub fn write_audio_blocks(
572 &mut self,
573 path: &Path,
574 role: Option<&str>,
575 opts: &EncodeOptions,
576 sample_rate: u32,
577 channels: usize,
578 frames: usize,
579 produce: &mut dyn FnMut(&mut dyn FnMut(&[&[f32]]) -> Result<()>) -> Result<()>,
580 ) -> Result<()> {
581 if path.as_os_str() != "-" && !self.globals.dry_run && !(self.globals.in_place && path.exists()) {
582 self.prepare_output(path)?;
583 let tmp = tmp_path(path);
584 if let Some(mut w) = self.io.open_writer(&tmp, sample_rate, channels, frames as u64, opts)? {
585 let mut written = 0usize;
586 let mut sink = |p: &[&[f32]]| {
587 if p.len() != channels {
588 return Err(Error::internal("audio block with the wrong channel count"));
589 }
590 written += p.first().map(|c| c.len()).unwrap_or(0);
591 w.write(p)
592 };
593 let mut r = produce(&mut sink);
594 if r.is_ok() && written != frames {
595 r = Err(Error::internal(format!("streamed {written} frames of {frames}")));
596 }
597 let r = r.and_then(|_| w.finish()).and_then(|_| std::fs::rename(&tmp, path).map_err(Error::from));
598 if let Err(e) = r {
599 let _ = std::fs::remove_file(&tmp);
600 return Err(e);
601 }
602 self.outputs.push(OutputRecord {
603 path: path.display().to_string(),
604 role: role.map(|r| r.to_string()),
605 duration_s: Some(frames as f64 / sample_rate.max(1) as f64),
606 sample_rate: Some(sample_rate),
607 channels: Some(channels),
608 extra: Map::new(),
609 });
610 self.progress.output(&path.display().to_string(), role);
611 return Ok(());
612 }
613 }
614 let mut data: Vec<Vec<f32>> = (0..channels).map(|_| Vec::with_capacity(frames)).collect();
615 produce(&mut |p: &[&[f32]]| {
616 for (d, s) in data.iter_mut().zip(p) {
617 d.extend_from_slice(s);
618 }
619 Ok(())
620 })?;
621 self.write_audio_with(path, &AudioBuffer { sample_rate, data }, role, opts)
622 }
623
624 pub fn write_file(&mut self, path: &Path, bytes: &[u8], role: Option<&str>) -> Result<()> {
626 self.prepare_output(path)?;
627 if !self.globals.dry_run {
628 if path.as_os_str() == "-" {
629 self.stdout.get_or_insert_with(Vec::new).extend_from_slice(bytes);
630 } else {
631 let tmp = tmp_path(path);
632 std::fs::write(&tmp, bytes)?;
633 std::fs::rename(&tmp, path)?;
634 }
635 }
636 self.record_output(path, role);
637 Ok(())
638 }
639
640 pub fn record_output(&mut self, path: &Path, role: Option<&str>) {
642 self.outputs.push(OutputRecord {
643 path: path.display().to_string(),
644 role: role.map(|r| r.to_string()),
645 duration_s: None,
646 sample_rate: None,
647 channels: None,
648 extra: Map::new(),
649 });
650 self.progress.output(&path.display().to_string(), role);
651 }
652
653 pub fn set_result(&mut self, value: impl Serialize) {
655 self.result = Some(serde_json::to_value(value).unwrap_or(Value::Null));
656 }
657
658 pub fn warn(&mut self, code: &str, message: impl Into<String>) {
659 let message = message.into();
660 self.progress.warning(code, &message);
661 self.warnings.push(WarningRecord { code: code.to_string(), message });
662 }
663
664 pub fn metric(&mut self, key: &str, value: impl Serialize) {
665 self.metrics.insert(key.to_string(), serde_json::to_value(value).unwrap_or(Value::Null));
666 }
667
668 pub fn note(&mut self, key: &str, value: impl Serialize) {
670 self.notes.insert(key.to_string(), serde_json::to_value(value).unwrap_or(Value::Null));
671 }
672
673 pub fn model_used(&mut self, name: &str, version: &str, sha256: &str) {
674 let r = ModelRecord { name: name.into(), version: version.into(), sha256: sha256.into() };
675 if !self.models.contains(&r) {
676 self.models.push(r);
677 }
678 }
679
680 pub fn set_progress(&self, fraction: f64) {
681 self.progress.progress(&self.command, fraction.clamp(0.0, 1.0));
682 }
683
684 pub fn check_cancel(&self) -> Result<()> {
685 self.cancel.check()
686 }
687
688 pub fn threads(&self) -> usize {
690 if self.globals.threads == 0 {
691 std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1)
692 } else {
693 self.globals.threads
694 }
695 }
696}
697
698pub fn tmp_path(path: &Path) -> PathBuf {
699 let name = path.file_name().map(|n| n.to_string_lossy().to_string()).unwrap_or_default();
700 let ext = path.extension().map(|e| e.to_string_lossy().to_string()).unwrap_or_default();
701 path.with_file_name(format!(".{name}.{}.part.{ext}", std::process::id()))
703}
704
705pub fn absolutize(p: &Path) -> PathBuf {
706 let expanded = if let Ok(s) = p.strip_prefix("~") {
707 std::env::var_os("HOME").map(|h| PathBuf::from(h).join(s)).unwrap_or_else(|| p.to_path_buf())
708 } else {
709 p.to_path_buf()
710 };
711 let abs =
712 if expanded.is_absolute() { expanded } else { std::env::current_dir().unwrap_or_default().join(expanded) };
713 let mut out = PathBuf::new();
715 for c in abs.components() {
716 match c {
717 std::path::Component::ParentDir => {
718 out.pop();
719 }
720 std::path::Component::CurDir => {}
721 other => out.push(other.as_os_str()),
722 }
723 }
724 out
725}