Skip to main content

alpha2_native/
audio_input.rs

1//! Audio encoder input: project audio at a rate the codec carries, cut into whole frames plus an
2//! exact last frame, in the encoder's sample format.
3//!
4//! `aac` takes only MPEG-4 rates and `libopus` only 8, 12, 16, 24 and 48 kHz, while projects run at
5//! 8 to 96 kHz. Unsupported rates export at the next one up (48 kHz for Opus above 24), resampled
6//! by `alpha2_audio::Resampler`. The last frame is exactly what is left (`aac` accepts a short
7//! frame, `libopus` pads it), so the track ends where the export's audio does.
8
9use std::collections::VecDeque;
10
11use alpha2_media::{AudioBlock, EncodeError};
12use rsmpeg::avcodec::{AVCodecContext, AVCodecRef};
13use rsmpeg::avutil::{AVChannelLayout, AVRational, ra};
14use rsmpeg::ffi;
15
16use crate::audio_timeline::SourceTimeline;
17use crate::capabilities::NativeAudioCodec;
18use crate::convert::encode_adapter_error;
19
20/// The audio stream an export encodes beside its video.
21#[derive(Copy, Clone, Debug, Eq, PartialEq)]
22pub struct AudioStreamConfig {
23    /// The codec to encode.
24    pub codec: NativeAudioCodec,
25    /// Project audio sample rate.
26    pub sample_rate: u32,
27    /// Channel count.
28    pub channels: u16,
29    /// Bitrate in bits per second; the host resolves the level.
30    pub bit_rate: u32,
31}
32
33/// An opened audio encoder, before it has a stream.
34pub(crate) struct OpenedAudioEncoder {
35    pub(crate) encoder: AVCodecContext,
36    pub(crate) encoder_time_base: AVRational,
37    /// The input the encoder takes its frames from.
38    pub(crate) input: EncoderAudioInput,
39    /// The rate it is encoded at ([`encoder_sample_rate`]).
40    pub(crate) encoder_rate: u32,
41    /// The sample format its frames carry ([`encoder_sample_format`]).
42    pub(crate) sample_format: ffi::AVSampleFormat,
43}
44
45/// Opens the encoder for `config` with global headers, which every export container uses.
46///
47/// # Errors
48/// When the runtime has no encoder for the codec, it takes none of the sample formats alpha2 can
49/// provide, it has no fixed frame size, or it does not open.
50pub(crate) fn open_audio_encoder(
51    config: AudioStreamConfig,
52    global_header: bool,
53) -> Result<OpenedAudioEncoder, EncodeError> {
54    let AudioStreamConfig {
55        codec,
56        sample_rate,
57        channels,
58        bit_rate,
59    } = config;
60    let encoder_codec = crate::capabilities::audio_encoder(codec).ok_or_else(|| {
61        EncodeError::Adapter(format!(
62            "this FFmpeg runtime has no {} encoder",
63            codec.as_str()
64        ))
65    })?;
66    let sample_format = encoder_sample_format(&encoder_codec).ok_or_else(|| {
67        EncodeError::Adapter(format!(
68            "the {} encoder takes neither float nor 16-bit samples",
69            codec.as_str()
70        ))
71    })?;
72    let encoder_rate = encoder_sample_rate(&encoder_codec, sample_rate);
73    let rate = i32::try_from(encoder_rate)
74        .map_err(|_| EncodeError::Adapter("audio sample rate is out of range".into()))?;
75    let encoder_time_base = ra(1, rate);
76    let mut encoder = AVCodecContext::new(&encoder_codec);
77    encoder.set_sample_rate(rate);
78    encoder.set_sample_fmt(sample_format);
79    encoder.set_ch_layout(AVChannelLayout::from_nb_channels(i32::from(channels)).into_inner());
80    encoder.set_time_base(encoder_time_base);
81    encoder.set_bit_rate(i64::from(bit_rate));
82    if global_header {
83        encoder.set_flags(
84            encoder.flags | i32::try_from(ffi::AV_CODEC_FLAG_GLOBAL_HEADER).unwrap_or(0),
85        );
86    }
87    encoder.open(None).map_err(encode_adapter_error)?;
88    let frame_size = usize::try_from(encoder.frame_size)
89        .ok()
90        .filter(|size| *size > 0)
91        .ok_or_else(|| {
92            EncodeError::Adapter(format!(
93                "the {} encoder has no fixed frame size",
94                codec.as_str()
95            ))
96        })?;
97    Ok(OpenedAudioEncoder {
98        encoder,
99        encoder_time_base,
100        input: EncoderAudioInput::new(sample_rate, encoder_rate, usize::from(channels), frame_size),
101        encoder_rate,
102        sample_format,
103    })
104}
105
106/// The sample format alpha2 hands `codec`: the first it takes of planar float, interleaved float,
107/// interleaved 16-bit; planar float if it lists none; `None` if it takes none of the three.
108pub(crate) fn encoder_sample_format(codec: &AVCodecRef<'_>) -> Option<ffi::AVSampleFormat> {
109    let Some(formats) = codec.sample_fmts() else {
110        return Some(ffi::AV_SAMPLE_FMT_FLTP);
111    };
112    [
113        ffi::AV_SAMPLE_FMT_FLTP,
114        ffi::AV_SAMPLE_FMT_FLT,
115        ffi::AV_SAMPLE_FMT_S16,
116    ]
117    .into_iter()
118    .find(|format| formats.contains(format))
119}
120
121/// The rate `codec` encodes a `project_rate` project at: that rate if supported, else the lowest
122/// above it, else its highest.
123pub(crate) fn encoder_sample_rate(codec: &AVCodecRef<'_>, project_rate: u32) -> u32 {
124    let Some(rates) = codec.supported_samplerates() else {
125        return project_rate;
126    };
127    let rates = rates
128        .iter()
129        .filter_map(|rate| u32::try_from(*rate).ok())
130        .filter(|rate| *rate > 0)
131        .collect::<Vec<_>>();
132    if rates.contains(&project_rate) {
133        return project_rate;
134    }
135    rates
136        .iter()
137        .copied()
138        .filter(|rate| *rate > project_rate)
139        .min()
140        .or_else(|| rates.iter().copied().max())
141        .unwrap_or(project_rate)
142}
143
144/// Project-rate audio blocks cut into the encoder's frames at its own rate.
145pub(crate) struct EncoderAudioInput {
146    project_rate: u32,
147    channels: usize,
148    frame_size: usize,
149    /// Brings the project's audio to the encoder's rate, when they differ.
150    resampler: Option<SourceTimeline>,
151    /// Samples at the encoder's rate not yet handed over.
152    pending: Vec<VecDeque<f32>>,
153    /// Project-rate samples accepted.
154    received: u64,
155}
156
157impl EncoderAudioInput {
158    pub(crate) fn new(
159        project_rate: u32,
160        encoder_rate: u32,
161        channels: usize,
162        frame_size: usize,
163    ) -> Self {
164        Self {
165            project_rate,
166            channels,
167            frame_size,
168            resampler: (encoder_rate != project_rate)
169                .then(|| SourceTimeline::new(encoder_rate, channels, Some(0))),
170            pending: (0..channels).map(|_| VecDeque::new()).collect(),
171            received: 0,
172        }
173    }
174
175    /// Takes the next block, which must start where the last ended, and returns the
176    /// whole frames it completes.
177    ///
178    /// # Errors
179    /// A block of another rate or channel count, of unequal channels, or out of order.
180    pub(crate) fn push(&mut self, block: &AudioBlock) -> Result<Vec<Vec<Vec<f32>>>, EncodeError> {
181        if block.sample_rate != self.project_rate
182            || block.channels.len() != self.channels
183            || block.channels.is_empty()
184        {
185            return Err(EncodeError::Adapter("audio block format changed".into()));
186        }
187        let frames = block.channels[0].len();
188        if block.channels.iter().any(|channel| channel.len() != frames) {
189            return Err(EncodeError::Adapter("audio channel lengths differ".into()));
190        }
191        if u64::try_from(block.start_sample).ok() != Some(self.received) {
192            return Err(EncodeError::InvalidTimestamp);
193        }
194        match self.resampler.as_mut() {
195            Some(resampler) => {
196                resampler
197                    .place(
198                        self.project_rate,
199                        block.start_sample,
200                        block.channels.clone(),
201                    )
202                    .map_err(|error| EncodeError::Adapter(error.to_string()))?;
203                resampler.produce(&mut self.pending, usize::MAX);
204            }
205            None => {
206                for (pending, channel) in self.pending.iter_mut().zip(&block.channels) {
207                    pending.extend(channel.iter().copied());
208                }
209            }
210        }
211        self.received = self
212            .received
213            .saturating_add(u64::try_from(frames).unwrap_or(u64::MAX));
214        let mut ready = Vec::new();
215        while self.pending[0].len() >= self.frame_size {
216            ready.push(self.take(self.frame_size));
217        }
218        Ok(ready)
219    }
220
221    /// Ends the input and returns the remaining frames, the last as long as what is left.
222    pub(crate) fn finish(&mut self) -> Vec<Vec<Vec<f32>>> {
223        if let Some(resampler) = self.resampler.as_mut() {
224            resampler.finish();
225            resampler.produce(&mut self.pending, usize::MAX);
226        }
227        let mut frames = Vec::new();
228        while !self.pending[0].is_empty() {
229            frames.push(self.take(self.frame_size.min(self.pending[0].len())));
230        }
231        frames
232    }
233
234    /// The input time accepted so far in ticks, reported back as consumed
235    /// once its frames are encoded.
236    pub(crate) fn received_ticks(&self) -> i64 {
237        i64::try_from(self.received)
238            .unwrap_or(i64::MAX)
239            .saturating_mul(alpha2_time::TICKS_PER_SECOND)
240            .div_euclid(i64::from(self.project_rate))
241    }
242
243    fn take(&mut self, samples: usize) -> Vec<Vec<f32>> {
244        self.pending
245            .iter_mut()
246            .map(|pending| pending.drain(..samples).collect())
247            .collect()
248    }
249}