alpha2_native/
audio_input.rs1use 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#[derive(Copy, Clone, Debug, Eq, PartialEq)]
22pub struct AudioStreamConfig {
23 pub codec: NativeAudioCodec,
25 pub sample_rate: u32,
27 pub channels: u16,
29 pub bit_rate: u32,
31}
32
33pub(crate) struct OpenedAudioEncoder {
35 pub(crate) encoder: AVCodecContext,
36 pub(crate) encoder_time_base: AVRational,
37 pub(crate) input: EncoderAudioInput,
39 pub(crate) encoder_rate: u32,
41 pub(crate) sample_format: ffi::AVSampleFormat,
43}
44
45pub(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
106pub(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
121pub(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
144pub(crate) struct EncoderAudioInput {
146 project_rate: u32,
147 channels: usize,
148 frame_size: usize,
149 resampler: Option<SourceTimeline>,
151 pending: Vec<VecDeque<f32>>,
153 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 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 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 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}