Skip to main content

pito_codecs/media/front/
stream.rs

1use std::io::Read;
2use std::path::Path;
3
4use crate::aac::Resampler;
5
6use super::feed::Feed;
7use super::ffmpeg::Child;
8use super::native::{self, PAD, shape};
9use super::probe::{container_ffmpeg, probe};
10use super::route::{Caps, Prefs, pick_decode};
11use super::sound::Sound;
12use crate::media::{AudioBuffer, Codec, Error, Result, Route};
13
14pub struct AudioStream {
15    pub route: Route,
16    pub rate: u32,
17    pub channels: u16,
18    inner: Inner,
19}
20
21impl std::fmt::Debug for AudioStream {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        write!(
24            f,
25            "AudioStream({} {} Hz, {} channels)",
26            self.route, self.rate, self.channels
27        )
28    }
29}
30
31enum Inner {
32    Native(Box<Native>),
33    Ffmpeg { child: Child, got: bool },
34    Mix(Vec<Part>),
35    Done,
36}
37
38struct Part {
39    stream: AudioStream,
40    gain: f32,
41    pending: Vec<f32>,
42    at: usize,
43    ended: bool,
44}
45
46impl Part {
47    fn frames(&self) -> usize {
48        self.pending.len() / 2 - self.at
49    }
50}
51
52struct Native {
53    feed: Feed,
54    resampler: Option<Resampler>,
55    up_first: bool,
56    rate: u32,
57    channels: u16,
58    raw: Vec<f32>,
59    fed: usize,
60    out: Vec<f32>,
61    drop: usize,
62    sent: usize,
63    start: usize,
64    eof: bool,
65}
66
67pub fn audio_stream(
68    path: &Path,
69    track: usize,
70    from: f64,
71    rate: u32,
72    channels: u16,
73    prefs: &Prefs,
74) -> Result<AudioStream> {
75    if rate == 0 || channels == 0 {
76        return Err(Error::Invalid(format!(
77            "a sample rate of {rate} with {channels} channels"
78        )));
79    }
80    if !from.is_finite() {
81        return Err(Error::Invalid(format!("a stream from {from}")));
82    }
83    let picked = pick_decode(Codec::Aac, prefs, &Caps::detect(prefs))?;
84    if picked.route != Route::Ffmpeg {
85        match Native::open(path, track, from, rate, channels) {
86            Ok(n) => {
87                return Ok(AudioStream {
88                    route: Route::Cpu,
89                    rate,
90                    channels,
91                    inner: Inner::Native(Box::new(n)),
92                });
93            }
94            Err(e) if prefs.user_ffmpeg().is_none() => return Err(e),
95            Err(_) => {}
96        }
97    }
98    let f = container_ffmpeg(prefs)?;
99    let from = from.max(0.0);
100    let mut args: Vec<String> = Vec::new();
101    if from > 0.0 {
102        args.extend(["-ss".into(), format!("{from:.6}")]);
103    }
104    args.extend([
105        "-i".into(),
106        path.display().to_string(),
107        "-map".into(),
108        format!("0:a:{track}"),
109        "-vn".into(),
110        "-ac".into(),
111        channels.to_string(),
112        "-ar".into(),
113        rate.to_string(),
114        "-f".into(),
115        "f32le".into(),
116        "-".into(),
117    ]);
118    let child = f.spawn(&args, false, true)?;
119    Ok(AudioStream {
120        route: Route::Ffmpeg,
121        rate,
122        channels,
123        inner: Inner::Ffmpeg { child, got: false },
124    })
125}
126
127pub fn audio_mix_stream(
128    path: &Path,
129    gains: &[f32],
130    from: f64,
131    rate: u32,
132    prefs: &Prefs,
133) -> Result<AudioStream> {
134    if rate == 0 {
135        return Err(Error::Invalid(format!("a sample rate of {rate}")));
136    }
137    if !from.is_finite() {
138        return Err(Error::Invalid(format!("a stream from {from}")));
139    }
140    let mut parts = Vec::new();
141    for t in probe(path, prefs)?.audio {
142        let gain = gains.get(t.index).copied().unwrap_or(1.0);
143        if gain > 0.0 {
144            parts.push(Part {
145                stream: audio_stream(path, t.index, from, rate, 2, prefs)?,
146                gain,
147                pending: Vec::new(),
148                at: 0,
149                ended: false,
150            });
151        }
152    }
153    let route = if parts.iter().any(|p| p.stream.route == Route::Ffmpeg) {
154        Route::Ffmpeg
155    } else {
156        Route::Cpu
157    };
158    Ok(AudioStream {
159        route,
160        rate,
161        channels: 2,
162        inner: Inner::Mix(parts),
163    })
164}
165
166fn mix(parts: &mut [Part], rate: u32) -> Result<Option<AudioBuffer>> {
167    let want = rate as usize;
168    for p in parts.iter_mut() {
169        while !p.ended && p.frames() < want {
170            match p.stream.next() {
171                Some(b) => {
172                    if p.at > 0 {
173                        p.pending.drain(..p.at * 2);
174                        p.at = 0;
175                    }
176                    p.pending.extend_from_slice(&b?.samples);
177                }
178                None => p.ended = true,
179            }
180        }
181    }
182    let n = parts
183        .iter()
184        .map(|p| p.frames().min(want))
185        .max()
186        .unwrap_or(0);
187    if n == 0 {
188        return Ok(None);
189    }
190    let mut out = AudioBuffer::new(rate, 2);
191    out.samples = vec![0.0; n * 2];
192    for p in parts.iter_mut() {
193        let m = p.frames().min(n);
194        let from = p.at * 2;
195        let got = p.pending.get(from..from + m * 2).unwrap_or_default();
196        for (o, s) in out.samples.iter_mut().zip(got) {
197            *o += s * p.gain;
198        }
199        p.at += m;
200    }
201    Ok(Some(out))
202}
203
204impl Native {
205    fn open(path: &Path, track: usize, from: f64, rate: u32, channels: u16) -> Result<Native> {
206        let d = native::audio_demux(path)?;
207        let i = native::audio_track(&d, track)?;
208        let pad = Sound::open(&d, i)?.pre_roll().max(PAD);
209        let start = (from.max(0.0) * rate as f64).round() as usize;
210        let pad = start.min((pad * rate as f64).round() as usize);
211        let t0 = (start - pad) as f64 / rate as f64;
212        let feed = Feed::open(&d, i, t0)?;
213        if channels != feed.channels && channels > 2 {
214            return Err(Error::Invalid(format!(
215                "{channels} channels from a {}-channel track: ask for 1, 2 or {}",
216                feed.channels, feed.channels
217            )));
218        }
219        let up_first = channels > feed.channels;
220        let resampler = (feed.rate != rate).then(|| {
221            Resampler::new(
222                feed.rate,
223                rate,
224                if up_first { feed.channels } else { channels },
225            )
226        });
227        Ok(Native {
228            feed,
229            resampler,
230            up_first,
231            rate,
232            channels,
233            raw: Vec::new(),
234            fed: 0,
235            out: Vec::new(),
236            drop: pad,
237            sent: 0,
238            start,
239            eof: false,
240        })
241    }
242
243    fn convert(&mut self, buf: AudioBuffer) -> Result<AudioBuffer> {
244        if self.up_first {
245            let buf = match self.resampler.as_mut() {
246                Some(r) => r.process(&buf),
247                None => buf,
248            };
249            shape(buf, self.channels)
250        } else {
251            let buf = shape(buf, self.channels)?;
252            Ok(match self.resampler.as_mut() {
253                Some(r) => r.process(&buf),
254                None => buf,
255            })
256        }
257    }
258
259    fn keep(&mut self, buf: AudioBuffer) {
260        let ch = self.channels as usize;
261        let frames = buf.samples.len() / ch;
262        let drop = self.drop.min(frames);
263        self.drop -= drop;
264        self.out
265            .extend_from_slice(buf.samples.get(drop * ch..frames * ch).unwrap_or_default());
266    }
267
268    fn pump(&mut self) -> Result<()> {
269        let src_ch = self.feed.channels as usize;
270        let batch = self.feed.rate as usize;
271        self.raw.clear();
272        let mut ended = false;
273        loop {
274            let have = self.fed + self.raw.len() / src_ch;
275            if let Some(w) = self.feed.want(f64::INFINITY)
276                && have >= w
277            {
278                ended = true;
279                break;
280            }
281            if self.raw.len() / src_ch >= batch {
282                break;
283            }
284            if !self.feed.next(&mut self.raw)? {
285                ended = true;
286                break;
287            }
288        }
289        if ended && let Some(w) = self.feed.want(f64::INFINITY) {
290            let keep = w.saturating_sub(self.fed) * src_ch;
291            self.raw.resize(keep, 0.0);
292        }
293        self.fed += self.raw.len() / src_ch;
294        let mut buf = AudioBuffer::new(self.feed.rate, self.feed.channels);
295        buf.samples = std::mem::take(&mut self.raw);
296        let done = self.convert(buf)?;
297        self.raw = Vec::with_capacity(batch * src_ch);
298        self.keep(done);
299        if ended {
300            self.eof = true;
301            if let Some(tail) = self.resampler.as_mut().map(Resampler::flush) {
302                let tail = if self.up_first {
303                    shape(tail, self.channels)?
304                } else {
305                    tail
306                };
307                self.keep(tail);
308            }
309            let length = self.feed.end();
310            if length > 0.0 {
311                let end = ((length * self.rate as f64).round() as usize).max(self.start);
312                let total = (end - self.start).saturating_sub(self.sent);
313                self.out.resize(total * self.channels as usize, 0.0);
314            }
315        }
316        Ok(())
317    }
318
319    fn read(&mut self) -> Result<Option<AudioBuffer>> {
320        let block = self.rate as usize * self.channels as usize;
321        while self.out.len() < block && !self.eof {
322            self.pump()?;
323        }
324        if self.out.is_empty() {
325            return Ok(None);
326        }
327        let n = block.min(self.out.len());
328        let mut buf = AudioBuffer::new(self.rate, self.channels);
329        buf.samples = self.out.drain(..n).collect();
330        self.sent += buf.frames();
331        Ok(Some(buf))
332    }
333}
334
335pub(super) fn native_stream(
336    path: &Path,
337    track: usize,
338    from: f64,
339    rate: u32,
340    channels: u16,
341) -> Result<AudioStream> {
342    if rate == 0 || channels == 0 || !from.is_finite() {
343        return Err(Error::Invalid(format!(
344            "a stream from {from} at {rate} Hz with {channels} channels"
345        )));
346    }
347    let n = Native::open(path, track, from, rate, channels)?;
348    Ok(AudioStream {
349        route: Route::Cpu,
350        rate,
351        channels,
352        inner: Inner::Native(Box::new(n)),
353    })
354}
355
356impl AudioStream {
357    fn read(&mut self) -> Result<Option<AudioBuffer>> {
358        let block = self.rate as usize * self.channels as usize * 4;
359        match &mut self.inner {
360            Inner::Native(n) => n.read(),
361            Inner::Mix(parts) => mix(parts, self.rate),
362            Inner::Done => Ok(None),
363            Inner::Ffmpeg { child, got } => {
364                let mut bytes = vec![0u8; block];
365                let mut filled = 0;
366                if let Some(out) = child.stdout.as_mut() {
367                    while filled < block {
368                        match out.read(&mut bytes[filled..]) {
369                            Ok(0) => break,
370                            Ok(n) => filled += n,
371                            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
372                            Err(e) => return Err(Error::Io(e)),
373                        }
374                    }
375                }
376                let whole = filled / 4 / self.channels as usize * self.channels as usize * 4;
377                if whole == 0 {
378                    let first = !*got;
379                    if let Inner::Ffmpeg { child, .. } =
380                        std::mem::replace(&mut self.inner, Inner::Done)
381                    {
382                        let done = child.finish();
383                        if first {
384                            done?;
385                        }
386                    }
387                    return Ok(None);
388                }
389                *got = true;
390                let mut buf = AudioBuffer::new(self.rate, self.channels);
391                buf.samples = bytes[..whole]
392                    .chunks_exact(4)
393                    .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
394                    .collect();
395                Ok(Some(buf))
396            }
397        }
398    }
399}
400
401impl Iterator for AudioStream {
402    type Item = Result<AudioBuffer>;
403
404    fn next(&mut self) -> Option<Result<AudioBuffer>> {
405        let got = self.read();
406        if got.is_err() {
407            self.inner = Inner::Done;
408        }
409        got.transpose()
410    }
411}