alpha2_native/
audio_input.rs1use 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#[derive(Copy, Clone, Debug, Eq, PartialEq)]
28pub struct AudioStreamConfig {
29 pub codec: NativeAudioCodec,
31 pub sample_rate: u32,
33 pub channels: u16,
35 pub bit_rate: u32,
37}
38
39pub(crate) struct OpenedAudioEncoder {
41 pub(crate) encoder: AVCodecContext,
42 pub(crate) encoder_time_base: AVRational,
43 pub(crate) input: EncoderAudioInput,
45 pub(crate) encoder_rate: u32,
47 pub(crate) sample_format: ffi::AVSampleFormat,
49}
50
51pub(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
114pub(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
131pub(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
154pub(crate) struct EncoderAudioInput {
157 project_rate: u32,
158 channels: usize,
159 frame_size: usize,
160 resampler: Option<SourceTimeline>,
162 pending: Vec<VecDeque<f32>>,
164 received: u64,
166}
167
168impl EncoderAudioInput {
169 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 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 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 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}