Skip to main content

alpha2_native/
audio_input.rs

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