Skip to main content

alpha2_native/
encode.rs

1//! Muxed A/V and WAV encode sinks.
2
3use std::collections::VecDeque;
4use std::fs;
5use std::io::{BufWriter, Seek, SeekFrom, Write};
6use std::path::{Path, PathBuf};
7use std::sync::mpsc::{self, Receiver, SyncSender};
8use std::sync::{Arc, Mutex};
9use std::thread::{self, JoinHandle};
10use std::time::{Duration, Instant};
11
12use crate::aac::{AacInput, encoder_sample_rate};
13use crate::capabilities::{
14    AvEncodeConfig, NativeVideoEncoderSelection, VideoEncoderOpenRequest, open_video_encoder,
15};
16use crate::convert::{
17    codec_thread_budget, configure_rgba_to_yuv, encode_adapter_error, encode_adapter_error_at,
18    tick_base,
19};
20use crate::partial::PartialOutput;
21use alpha2_ffmpeg_bridge::OutputFile;
22use alpha2_ffmpeg_bridge::Yuv420pFrameBuffer;
23use alpha2_media::{
24    AudioBlock, AudioBlockSink, CancelToken, DecodedVideoFrame, EncodeError, EncodeEvent,
25    EncodeSink, MAX_AUDIO_CHANNELS, PixelFormat, SpeakerLayout,
26};
27use alpha2_time::MediaTime;
28use rsmpeg::avcodec::{AVCodec, AVCodecContext, AVPacket};
29use rsmpeg::avutil::{AVChannelLayout, AVFrameWithImage, AVImage, AVRational, av_rescale_q, ra};
30use rsmpeg::error::RsmpegError;
31use rsmpeg::ffi;
32use rsmpeg::swscale::SwsContext;
33
34/// Where one encode sink's wall time went.
35///
36/// The codec, the packet drain and the mux write move for different reasons:
37/// a codec change moves the first, a channel-depth change moves the second,
38/// and a container change moves the third. A single total such as
39/// `videoEncodeMs` cannot say which one did, so each is timed apart.
40#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
41pub struct EncodeStageTimings {
42    /// Waiting, before `avcodec_send_frame`, for the GPU to finish writing a
43    /// surface the encoder cannot wait for on its own. Zero on every path
44    /// whose encoder waits on the GPU itself or takes CPU frames.
45    pub video_input_ready_wait: Duration,
46    /// `avcodec_send_frame` for video: the codec itself.
47    pub video_send_frame: Duration,
48    /// `avcodec_receive_packet` for video.
49    pub video_receive_packet: Duration,
50    /// `av_interleaved_write_frame` for video packets.
51    pub video_mux_write: Duration,
52    /// `avcodec_send_frame` for audio.
53    pub audio_send_frame: Duration,
54    /// `avcodec_receive_packet` for audio.
55    pub audio_receive_packet: Duration,
56    /// `av_interleaved_write_frame` for audio packets.
57    pub audio_mux_write: Duration,
58}
59
60impl EncodeStageTimings {
61    /// What elapsed between `self` and a later snapshot.
62    #[must_use]
63    pub fn since(self, earlier: Self) -> Self {
64        Self {
65            video_input_ready_wait: self
66                .video_input_ready_wait
67                .saturating_sub(earlier.video_input_ready_wait),
68            video_send_frame: self
69                .video_send_frame
70                .saturating_sub(earlier.video_send_frame),
71            video_receive_packet: self
72                .video_receive_packet
73                .saturating_sub(earlier.video_receive_packet),
74            video_mux_write: self.video_mux_write.saturating_sub(earlier.video_mux_write),
75            audio_send_frame: self
76                .audio_send_frame
77                .saturating_sub(earlier.audio_send_frame),
78            audio_receive_packet: self
79                .audio_receive_packet
80                .saturating_sub(earlier.audio_receive_packet),
81            audio_mux_write: self.audio_mux_write.saturating_sub(earlier.audio_mux_write),
82        }
83    }
84}
85
86thread_local! {
87    /// What this thread calls before each encoder and muxer call; see
88    /// [`watch_encode_calls`].
89    static ENCODE_CALL_WATCH: std::cell::RefCell<Option<Box<dyn Fn()>>> =
90        const { std::cell::RefCell::new(None) };
91}
92
93/// Calls `beat` on this thread before every encoder and muxer call an encode
94/// sink makes on it, until the returned guard drops.
95///
96/// Some of those calls wait with no timeout: `h264_vulkan` waits on its
97/// execution fences with `UINT64_MAX`, and `write(2)` to a hard-mounted share
98/// that stopped answering is retried forever. Nothing can interrupt them, so
99/// the thread that drives a sink reports where it is instead, and its owner
100/// decides that a call which has not returned in time is stuck. A sink's calls
101/// all run on the thread that calls the sink, except the AAC encoder's own
102/// thread, which the sink waits on in a call this beat already covers.
103#[must_use = "the watch ends when the guard drops"]
104pub fn watch_encode_calls(beat: impl Fn() + 'static) -> EncodeCallWatch {
105    let previous = ENCODE_CALL_WATCH.with(|watch| watch.replace(Some(Box::new(beat))));
106    EncodeCallWatch {
107        previous,
108        thread: std::marker::PhantomData,
109    }
110}
111
112/// Ends a [`watch_encode_calls`] watch, restoring the one before it.
113pub struct EncodeCallWatch {
114    previous: Option<Box<dyn Fn()>>,
115    /// The watch belongs to the thread that set it.
116    thread: std::marker::PhantomData<*const ()>,
117}
118
119impl Drop for EncodeCallWatch {
120    fn drop(&mut self) {
121        let previous = self.previous.take();
122        ENCODE_CALL_WATCH.with(|watch| {
123            watch.replace(previous);
124        });
125    }
126}
127
128/// An encoder or muxer call starts now on this thread.
129pub(crate) fn beat_encode_call() {
130    ENCODE_CALL_WATCH.with(|watch| {
131        if let Some(beat) = watch.borrow().as_ref() {
132            beat();
133        }
134    });
135}
136
137/// Runs `body`, adding its elapsed time to `total`. `body` is an encoder or
138/// muxer call, so it is watched ([`watch_encode_calls`]).
139pub(crate) fn timed<T>(total: &mut Duration, body: impl FnOnce() -> T) -> T {
140    beat_encode_call();
141    let started = Instant::now();
142    let value = body();
143    *total = total.saturating_add(started.elapsed());
144    value
145}
146
147/// Transactional IEEE-f32 WAV encoder used by native offline audio export.
148pub struct NativeWavEncodeSink {
149    pub(crate) final_path: PathBuf,
150    pub(crate) partial_path: PathBuf,
151    /// Held until the partial file is renamed into place or removed.
152    pub(crate) _partial_lock: crate::partial::PartialLock,
153    pub(crate) file: Option<BufWriter<fs::File>>,
154    pub(crate) sample_rate: u32,
155    pub(crate) channels: u16,
156    /// The speakers every block feeds.
157    pub(crate) speakers: SpeakerLayout,
158    /// The `WAVE_FORMAT_EXTENSIBLE` channel mask a layout other than its
159    /// count's default is written with; `None` writes a plain IEEE-float
160    /// header, which readers take as the default.
161    pub(crate) channel_mask: Option<u32>,
162    /// The block plane each interleaved channel comes from: WAVE requires
163    /// the mask's bit order.
164    pub(crate) plane_order: Vec<usize>,
165    pub(crate) frames_written: u64,
166    pub(crate) complete: bool,
167    /// One block's samples, interleaved, reused across blocks.
168    pub(crate) interleaved: Vec<u8>,
169}
170
171/// Bytes the WAV sink gathers before each write: about five seconds of
172/// 48 kHz stereo, so a minute of audio takes a few dozen writes.
173const WAV_WRITE_BUFFER: usize = 1 << 20;
174
175/// Audio/video sink that encodes and interleaves through the linked `FFmpeg`
176/// libraries without requiring command-line programs on the host. Video frames
177/// and audio blocks must be submitted in nondecreasing start-timestamp order;
178/// ready packets are muxed incrementally instead of retained for the full file.
179pub struct NativeAvEncodeSink {
180    pub(crate) final_path: PathBuf,
181    pub(crate) partial_path: PathBuf,
182    /// Held until the partial file is renamed into place or removed.
183    pub(crate) _partial_lock: crate::partial::PartialLock,
184    pub(crate) output: Option<OutputFile>,
185    pub(crate) video_encoder: AVCodecContext,
186    pub(crate) video_stream_index: i32,
187    pub(crate) video_encoder_time_base: AVRational,
188    /// One frame in `video_encoder_time_base`.
189    pub(crate) video_frame_units: i64,
190    pub(crate) video_stream_time_base: AVRational,
191    pub(crate) video_scaler: SwsContext,
192    /// Turns decoded YUV of any layout into the encoder's 4:2:0.
193    pub(crate) yuv_input: crate::convert::Yuv420pConverter,
194    pub(crate) width: i32,
195    pub(crate) height: i32,
196    pub(crate) last_video_pts: Option<i64>,
197    /// AAC encoding, backlog and packet draining run on a separate thread
198    /// so video encoding cannot serialize audio work.
199    pub(crate) audio: AudioEncodeThread,
200    /// The project's audio, as the AAC encoder takes it.
201    pub(crate) audio_input: AacInput,
202    pub(crate) video_packets: VecDeque<QueuedPacket>,
203    pub(crate) last_input_ticks: Option<i64>,
204    pub(crate) complete: bool,
205    pub(crate) video_selection: NativeVideoEncoderSelection,
206    pub(crate) stages: EncodeStageTimings,
207}
208
209/// One AAC block handed to the audio encode thread.
210pub(crate) enum AudioCommand {
211    /// One frame of samples to encode, and the input tick the input reached.
212    /// Every frame is the encoder's frame size but the last, which may be
213    /// shorter.
214    Block(Vec<Vec<f32>>, i64),
215    /// No more input: flush the encoder and stop, the input having reached
216    /// the tick.
217    Finish(i64),
218}
219
220/// What the audio encode thread has produced, and how far it has got.
221///
222/// The muxer stays on the video worker beside `pop_next_packet`, so the only
223/// thing crossing threads is finished packets and a watermark.
224#[derive(Default)]
225pub(crate) struct AudioQueueState {
226    /// Packets in production order, which for one stream is DTS order.
227    pub(crate) packets: VecDeque<QueuedPacket>,
228    /// Input ticks the thread has consumed and produced every packet for.
229    ///
230    /// The muxer will not write a video packet past this. That is what keeps
231    /// the container's interleave independent of when the audio thread runs:
232    /// it writes exactly what the single-threaded sink would have written,
233    /// because the single-threaded sink had also consumed every audio block up
234    /// to the video frame it was about to encode.
235    pub(crate) consumed_ticks: i64,
236    /// Set once the thread has flushed its encoder and will produce no more.
237    pub(crate) finished: bool,
238    /// The thread's first failure, surfaced by the muxer.
239    pub(crate) failure: Option<String>,
240    /// The thread's share of the sink's stage totals.
241    pub(crate) stages: EncodeStageTimings,
242}
243
244/// Handle to the audio encode thread.
245pub(crate) struct AudioEncodeThread {
246    commands: Option<SyncSender<AudioCommand>>,
247    join: Option<JoinHandle<()>>,
248    state: Arc<Mutex<AudioQueueState>>,
249}
250
251impl AudioEncodeThread {
252    pub(crate) fn lock(&self) -> std::sync::MutexGuard<'_, AudioQueueState> {
253        self.state
254            .lock()
255            .unwrap_or_else(std::sync::PoisonError::into_inner)
256    }
257
258    /// Hands one command over, reporting the thread's failure if it has
259    /// stopped.
260    pub(crate) fn send(&self, command: AudioCommand) -> Result<(), EncodeError> {
261        let sent = match self.commands.as_ref() {
262            Some(sender) => sender.send(command).is_ok(),
263            None => return Err(EncodeError::Adapter("audio encoder is closed".into())),
264        };
265        if sent {
266            return Ok(());
267        }
268        let failure = self.lock().failure.clone();
269        Err(EncodeError::Adapter(failure.unwrap_or_else(|| {
270            "audio encoder stopped before accepting a block".into()
271        })))
272    }
273
274    /// Stops the thread, leaving whatever it produced in the queue.
275    pub(crate) fn shutdown(&mut self) {
276        drop(self.commands.take());
277        if let Some(join) = self.join.take() {
278            let _ = join.join();
279        }
280    }
281}
282
283impl Drop for AudioEncodeThread {
284    fn drop(&mut self) {
285        self.shutdown();
286    }
287}
288
289/// A finished packet waiting for the muxer, ordered by its DTS in ticks.
290pub(crate) struct QueuedPacket {
291    pub(crate) sort_ticks: i64,
292    pub(crate) packet: AVPacket,
293}
294
295impl NativeAvEncodeSink {
296    /// Creates linked-library video, audio, and container encoders without
297    /// touching the final destination.
298    ///
299    /// # Errors
300    /// Returns an encoder error for invalid configuration, paths, or stale
301    /// component outputs.
302    pub fn create(path: impl AsRef<Path>, config: AvEncodeConfig) -> Result<Self, EncodeError> {
303        Self::create_selected(path, config, None, &codec_thread_budget)
304    }
305
306    /// Creates the requested concrete video encoder by its `FFmpeg` short name.
307    ///
308    /// # Errors
309    /// Returns an encoder error when the name is unavailable, implements a
310    /// different codec, or the remaining configuration is invalid.
311    pub fn create_with_video_encoder(
312        path: impl AsRef<Path>,
313        config: AvEncodeConfig,
314        encoder_name: &str,
315    ) -> Result<Self, EncodeError> {
316        Self::create_selected(path, config, Some(encoder_name), &codec_thread_budget)
317    }
318
319    /// Creates an encoder whose software thread count the caller decides.
320    ///
321    /// `software_threads` is asked, at most once, only when a candidate that
322    /// runs on the CPU is opened: the caller sizes the rest of its work from
323    /// [`NativeVideoEncoderSelection`] afterwards, so a hardware candidate that
324    /// fails to open and falls through to software still gets a software
325    /// encoder's share of the CPU.
326    ///
327    /// `encoder_name` is the concrete encoder the caller pinned. `None` lets
328    /// the candidates for the codec be tried in preference order, which is
329    /// what makes the open-time fallback reachable.
330    ///
331    /// # Errors
332    /// Returns an encoder error when no candidate opens, or for invalid
333    /// configuration, paths, codecs, or muxers.
334    pub fn create_with_video_encoder_threads(
335        path: impl AsRef<Path>,
336        config: AvEncodeConfig,
337        encoder_name: Option<&str>,
338        software_threads: &dyn Fn() -> i64,
339    ) -> Result<Self, EncodeError> {
340        Self::create_selected(path, config, encoder_name, software_threads)
341    }
342
343    #[allow(clippy::too_many_lines)]
344    pub(crate) fn create_selected(
345        path: impl AsRef<Path>,
346        config: AvEncodeConfig,
347        encoder_name: Option<&str>,
348        software_threads: &dyn Fn() -> i64,
349    ) -> Result<Self, EncodeError> {
350        crate::video_encode::check_video_config(&config.video)?;
351        if config.audio_sample_rate == 0
352            || config.audio_channels == 0
353            || config.audio_channels > 8
354            || config.audio_bit_rate == 0
355        {
356            return Err(EncodeError::Adapter(format!(
357                "the audio output ({} channels at {} Hz, {} bit/s) is refused: 1 to 8 channels, a positive rate and bitrate",
358                config.audio_channels, config.audio_sample_rate, config.audio_bit_rate
359            )));
360        }
361        let final_path = path.as_ref().to_owned();
362        let partial = PartialOutput::claim(&final_path)?;
363        let partial_path = partial.path().to_owned();
364        let mut output = partial.open_muxer()?;
365        let global_header =
366            output.oformat().flags & i32::try_from(ffi::AVFMT_GLOBALHEADER).unwrap_or(0) != 0;
367
368        let width = i32::try_from(config.video.width)
369            .map_err(|_| EncodeError::Adapter("width is out of range".into()))?;
370        let height = i32::try_from(config.video.height)
371            .map_err(|_| EncodeError::Adapter("height is out of range".into()))?;
372        let (video_encoder, video_selection) = open_video_encoder(&VideoEncoderOpenRequest {
373            config: &config.video,
374            requested_name: encoder_name,
375            threads: crate::capabilities::SoftwareEncoderThreads::new(software_threads),
376            global_header,
377            width,
378            height,
379        })?;
380        // The frame length, or the finest unit an encoder that cannot hold it takes.
381        let video_encoder_time_base = video_encoder.time_base;
382
383        let audio_channels = i32::from(config.audio_channels);
384        let audio_codec =
385            AVCodec::find_encoder(config.audio_codec.codec_id()).ok_or_else(|| {
386                EncodeError::Adapter(format!(
387                    "{} encoder is unavailable",
388                    config.audio_codec.as_str()
389                ))
390            })?;
391        if audio_codec
392            .sample_fmts()
393            .is_some_and(|formats| !formats.contains(&ffi::AV_SAMPLE_FMT_FLTP))
394        {
395            return Err(EncodeError::Adapter(format!(
396                "the {} encoder does not take planar float samples",
397                config.audio_codec.as_str()
398            )));
399        }
400        let encoder_rate = encoder_sample_rate(&audio_codec, config.audio_sample_rate);
401        let audio_sample_rate = i32::try_from(encoder_rate)
402            .map_err(|_| EncodeError::Adapter("audio sample rate is out of range".into()))?;
403        let audio_encoder_time_base = ra(1, audio_sample_rate);
404        let mut audio_encoder = AVCodecContext::new(&audio_codec);
405        audio_encoder.set_sample_rate(audio_sample_rate);
406        audio_encoder.set_sample_fmt(ffi::AV_SAMPLE_FMT_FLTP);
407        audio_encoder.set_ch_layout(AVChannelLayout::from_nb_channels(audio_channels).into_inner());
408        audio_encoder.set_time_base(audio_encoder_time_base);
409        audio_encoder.set_bit_rate(i64::from(config.audio_bit_rate));
410        if global_header {
411            audio_encoder.set_flags(
412                audio_encoder.flags | i32::try_from(ffi::AV_CODEC_FLAG_GLOBAL_HEADER).unwrap_or(0),
413            );
414        }
415        audio_encoder.open(None).map_err(encode_adapter_error)?;
416        let audio_frame_size = usize::try_from(audio_encoder.frame_size)
417            .ok()
418            .filter(|size| *size > 0)
419            .ok_or_else(|| EncodeError::Adapter("AAC encoder has no fixed frame size".into()))?;
420
421        let video_stream_index;
422        {
423            let mut stream = output.new_stream();
424            stream.set_time_base(video_encoder_time_base);
425            stream.set_codecpar(video_encoder.extract_codecpar());
426            crate::video_encode::describe_output_stream(
427                &mut stream,
428                &video_encoder,
429                config.video.display,
430            )?;
431            video_stream_index = stream.index;
432        }
433        let audio_stream_index;
434        {
435            let mut stream = output.new_stream();
436            stream.set_time_base(audio_encoder_time_base);
437            stream.set_codecpar(audio_encoder.extract_codecpar());
438            audio_stream_index = stream.index;
439        }
440        output
441            .write_header(&mut None)
442            .map_err(encode_adapter_error)?;
443        let video_stream_time_base =
444            output.streams()[usize::try_from(video_stream_index).unwrap_or(0)].time_base;
445        let audio_stream_time_base =
446            output.streams()[usize::try_from(audio_stream_index).unwrap_or(0)].time_base;
447        let mut video_scaler = SwsContext::get_context(
448            width,
449            height,
450            ffi::AV_PIX_FMT_RGBA,
451            width,
452            height,
453            ffi::AV_PIX_FMT_YUV420P,
454            ffi::SWS_BILINEAR,
455            None,
456            None,
457            None,
458        )
459        .ok_or_else(|| EncodeError::Adapter("cannot create RGBA to YUV scaler".into()))?;
460        configure_rgba_to_yuv(&mut video_scaler, config.video.color)
461            .map_err(encode_adapter_error)?;
462        let audio = spawn_audio_encoder(AudioEncoderConfig {
463            encoder: audio_encoder,
464            stream_index: audio_stream_index,
465            encoder_time_base: audio_encoder_time_base,
466            stream_time_base: audio_stream_time_base,
467            sample_rate: encoder_rate,
468        })?;
469        let partial_lock = partial.keep();
470        Ok(Self {
471            final_path,
472            partial_path,
473            _partial_lock: partial_lock,
474            output: Some(output),
475            video_encoder,
476            video_stream_index,
477            video_encoder_time_base,
478            video_frame_units: crate::mux::frame_units(
479                config.video.frame_rate,
480                video_encoder_time_base,
481            ),
482            video_stream_time_base,
483            video_scaler,
484            yuv_input: crate::convert::Yuv420pConverter::new(config.video.color),
485            width,
486            height,
487            last_video_pts: None,
488            audio,
489            audio_input: AacInput::new(
490                config.audio_sample_rate,
491                encoder_rate,
492                usize::from(config.audio_channels),
493                audio_frame_size,
494            ),
495            video_packets: VecDeque::new(),
496            last_input_ticks: None,
497            complete: false,
498            video_selection,
499            stages: EncodeStageTimings::default(),
500        })
501    }
502
503    /// Audio packets the thread has produced and the muxer has not written.
504    #[cfg(test)]
505    pub(crate) fn queued_audio_packets(&self) -> usize {
506        self.audio.lock().packets.len()
507    }
508
509    /// Where this sink's time has gone so far, both threads together.
510    #[must_use]
511    pub fn stage_timings(&self) -> EncodeStageTimings {
512        let audio = self.audio.lock().stages;
513        EncodeStageTimings {
514            audio_send_frame: audio.audio_send_frame,
515            audio_receive_packet: audio.audio_receive_packet,
516            ..self.stages
517        }
518    }
519
520    /// The video encoder this sink opened, and the candidates it passed over.
521    #[must_use]
522    pub fn video_selection(&self) -> &NativeVideoEncoderSelection {
523        &self.video_selection
524    }
525
526    pub(crate) fn cleanup(&mut self) -> Result<(), EncodeError> {
527        // Stop the audio thread before the muxer goes, so nothing is queued against a closed
528        // output.
529        self.audio.shutdown();
530        drop(self.output.take());
531        match fs::remove_file(&self.partial_path) {
532            Ok(()) => Ok(()),
533            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
534            Err(error) => Err(encode_adapter_error(error)),
535        }
536    }
537
538    /// Queues one finished video packet for the muxer.
539    fn queue_video_packet(&mut self, mut packet: AVPacket) {
540        packet.set_stream_index(self.video_stream_index);
541        crate::mux::fill_frame_duration(&mut packet, self.video_frame_units);
542        packet.rescale_ts(self.video_encoder_time_base, self.video_stream_time_base);
543        let timestamp = if packet.dts == ffi::AV_NOPTS_VALUE {
544            packet.pts
545        } else {
546            packet.dts
547        };
548        self.video_packets.push_back(QueuedPacket {
549            sort_ticks: av_rescale_q(timestamp, self.video_stream_time_base, tick_base()),
550            packet,
551        });
552    }
553
554    /// Takes the next packet the muxer may write, or `None` to wait.
555    ///
556    /// Deterministic interleaving follows two rules:
557    /// - Earlier DTS wins; video wins ties so thread scheduling cannot affect output.
558    /// - Hold video until audio has consumed all input ticks through its timestamp.
559    ///   The coordinator's channel bounds how far video can get ahead.
560    fn pop_next_packet(&mut self, flushing: bool) -> Result<Option<QueuedPacket>, EncodeError> {
561        let mut state = self.audio.lock();
562        if let Some(error) = state.failure.clone() {
563            return Err(EncodeError::Adapter(error));
564        }
565        let take_audio = match (self.video_packets.front(), state.packets.front()) {
566            (Some(video), Some(audio)) => audio.sort_ticks < video.sort_ticks,
567            (None, Some(_)) => true,
568            (_, None) => false,
569        };
570        if take_audio {
571            return Ok(state.packets.pop_front());
572        }
573        let Some(video) = self.video_packets.front() else {
574            return Ok(None);
575        };
576        if !flushing && !state.finished && state.consumed_ticks < video.sort_ticks {
577            return Ok(None);
578        }
579        drop(state);
580        Ok(self.video_packets.pop_front())
581    }
582
583    pub(crate) fn write_ready_packets(&mut self, flushing: bool) -> Result<(), EncodeError> {
584        while let Some(mut queued) = self.pop_next_packet(flushing)? {
585            // The muxer is shared, so the write is charged to the stream whose packet it is.
586            let video = queued.packet.stream_index == self.video_stream_index;
587            let total = if video {
588                &mut self.stages.video_mux_write
589            } else {
590                &mut self.stages.audio_mux_write
591            };
592            let output = self
593                .output
594                .as_mut()
595                .ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?;
596            timed(total, || output.interleaved_write_frame(&mut queued.packet))
597                .map_err(|error| encode_adapter_error_at("MP4 packet write failed", error))?;
598        }
599        Ok(())
600    }
601
602    pub(crate) fn drain_video_packets(&mut self, flushing: bool) -> Result<(), EncodeError> {
603        loop {
604            let received = timed(&mut self.stages.video_receive_packet, || {
605                self.video_encoder.receive_packet()
606            });
607            match received {
608                Ok(packet) => self.queue_video_packet(packet),
609                Err(RsmpegError::EncoderDrainError) if !flushing => return Ok(()),
610                Err(RsmpegError::EncoderFlushedError | RsmpegError::EncoderDrainError)
611                    if flushing =>
612                {
613                    return Ok(());
614                }
615                Err(error) => {
616                    return Err(encode_adapter_error_at(
617                        "video encoder packet drain failed",
618                        error,
619                    ));
620                }
621            }
622        }
623    }
624
625    /// Encodes a frame the caller already packed, without copying it.
626    ///
627    /// See [`crate::NativeVideoEncodeSink::write_prepared_video`]; this is the
628    /// same handover for the muxed audio/video sink.
629    ///
630    /// # Errors
631    /// Returns an error for a frame sized differently from this encoder, a
632    /// non-monotonic timestamp, or a codec failure.
633    pub fn write_prepared_video(
634        &mut self,
635        frame: &mut Yuv420pFrameBuffer,
636        pts: MediaTime,
637    ) -> Result<(), EncodeError> {
638        let (width, height) = frame.size();
639        if width != usize::try_from(self.width).unwrap_or(0)
640            || height != usize::try_from(self.height).unwrap_or(0)
641        {
642            return Err(EncodeError::Adapter(
643                "encoder input dimensions do not match its configuration".into(),
644            ));
645        }
646        let pts = av_rescale_q(pts.ticks(), tick_base(), self.video_encoder_time_base);
647        if self.last_video_pts.is_some_and(|previous| pts <= previous) {
648            return Err(EncodeError::InvalidTimestamp);
649        }
650        frame.set_presentation(pts, self.video_encoder_time_base);
651        timed(&mut self.stages.video_send_frame, || {
652            self.video_encoder.send_frame(Some(frame.frame()))
653        })
654        .map_err(|error| {
655            encode_adapter_error_at("video encoder rejected a YUV420P frame", error)
656        })?;
657        self.last_video_pts = Some(pts);
658        self.drain_video_packets(false)?;
659        self.write_ready_packets(false)
660    }
661
662    pub(crate) fn write_video(&mut self, frame: &DecodedVideoFrame) -> Result<(), EncodeError> {
663        if frame.visible_rect.width != u32::try_from(self.width).unwrap_or(0)
664            || frame.visible_rect.height != u32::try_from(self.height).unwrap_or(0)
665        {
666            return Err(EncodeError::Adapter(
667                "encoder input dimensions do not match its configuration".into(),
668            ));
669        }
670        let pts = av_rescale_q(frame.pts.ticks(), tick_base(), self.video_encoder_time_base);
671        if self.last_video_pts.is_some_and(|previous| pts <= previous) {
672            return Err(EncodeError::InvalidTimestamp);
673        }
674        if let Some(mut yuv) = self.yuv_input.convert(frame, self.width, self.height)? {
675            yuv.set_pts(pts);
676            yuv.set_time_base(self.video_encoder_time_base);
677            timed(&mut self.stages.video_send_frame, || {
678                self.video_encoder.send_frame(Some(&yuv))
679            })
680            .map_err(|error| {
681                encode_adapter_error_at("video encoder rejected a YUV420P frame", error)
682            })?;
683            self.last_video_pts = Some(pts);
684            self.drain_video_packets(false)?;
685            return self.write_ready_packets(false);
686        }
687        if frame.pixel_format != PixelFormat::Rgba8 || frame.planes.len() != 1 {
688            return Err(EncodeError::Adapter(
689                "encoder expects YUV planes or one RGBA8 plane".into(),
690            ));
691        }
692        let plane = &frame.planes[0];
693        let expected = usize::try_from(self.width)
694            .ok()
695            .and_then(|width| {
696                usize::try_from(self.height)
697                    .ok()
698                    .and_then(|height| width.checked_mul(height))
699            })
700            .and_then(|pixels| pixels.checked_mul(4))
701            .ok_or_else(|| EncodeError::Adapter("frame dimensions overflow".into()))?;
702        if plane.bytes.len() != expected
703            || plane.stride != u32::try_from(self.width).unwrap_or(0) * 4
704        {
705            return Err(EncodeError::Adapter(
706                "RGBA plane is not tightly packed".into(),
707            ));
708        }
709        // No alpha in the output, so premultiplied RGB is kept: frames arrive opaque over the
710        // project background, and dividing alpha out would show a translucent layer at full
711        // strength.
712        let rgba = alpha2_ffmpeg_bridge::rgba_frame(&plane.bytes, self.width, self.height)
713            .map_err(encode_adapter_error)?;
714        let yuv_image = AVImage::new(ffi::AV_PIX_FMT_YUV420P, self.width, self.height, 1)
715            .ok_or_else(|| EncodeError::Adapter("cannot allocate YUV input".into()))?;
716        let mut yuv = AVFrameWithImage::new(yuv_image);
717        self.video_scaler
718            .scale_frame(&rgba, 0, self.height, &mut yuv)
719            .map_err(encode_adapter_error)?;
720        yuv.set_pts(pts);
721        yuv.set_time_base(self.video_encoder_time_base);
722        timed(&mut self.stages.video_send_frame, || {
723            self.video_encoder.send_frame(Some(&yuv))
724        })
725        .map_err(|error| {
726            encode_adapter_error_at("video encoder rejected a converted YUV420P frame", error)
727        })?;
728        self.last_video_pts = Some(pts);
729        self.drain_video_packets(false)?;
730        self.write_ready_packets(false)
731    }
732
733    /// Buffers one audio block and hands the thread every whole AAC frame it
734    /// completes.
735    pub(crate) fn write_audio(&mut self, block: &AudioBlock) -> Result<(), EncodeError> {
736        let frames = self.audio_input.push(block)?;
737        let covered = self.audio_input.received_ticks();
738        for frame in frames {
739            self.audio.send(AudioCommand::Block(frame, covered))?;
740        }
741        // An audio-only stretch has no video packet to trigger a drain, so the queue would grow.
742        self.write_ready_packets(false)
743    }
744}
745
746/// What the audio encode thread needs from the sink that created it.
747pub(crate) struct AudioEncoderConfig {
748    pub(crate) encoder: AVCodecContext,
749    pub(crate) stream_index: i32,
750    pub(crate) encoder_time_base: AVRational,
751    pub(crate) stream_time_base: AVRational,
752    /// The rate the encoder runs at.
753    pub(crate) sample_rate: u32,
754}
755
756/// Starts the AAC encoder on its own thread.
757///
758/// The channel is two deep: the coordinator hands blocks in faster than the
759/// video encoder consumes frames, and a deeper queue would only buy backlog
760/// the muxer has to hold anyway.
761///
762/// # Errors
763/// When the OS refuses the thread (a process or cgroup thread limit, or no
764/// memory for its stack).
765pub(crate) fn spawn_audio_encoder(
766    config: AudioEncoderConfig,
767) -> Result<AudioEncodeThread, EncodeError> {
768    let (commands, receiver) = mpsc::sync_channel(2);
769    let state = Arc::new(Mutex::new(AudioQueueState::default()));
770    let worker_state = Arc::clone(&state);
771    let run = move || run_audio_encoder(config, &receiver, &worker_state);
772    let builder = thread::Builder::new().name("alpha2-audio-encode".to_owned());
773    #[cfg(test)]
774    let spawned = if REFUSE_AUDIO_ENCODE_THREAD.with(std::cell::Cell::take) {
775        drop(run);
776        Err(std::io::Error::from(std::io::ErrorKind::WouldBlock))
777    } else {
778        builder.spawn(run)
779    };
780    #[cfg(not(test))]
781    let spawned = builder.spawn(run);
782    let join = spawned.map_err(|error| {
783        encode_adapter_error_at("failed to start the audio encode thread", error)
784    })?;
785    Ok(AudioEncodeThread {
786        commands: Some(commands),
787        join: Some(join),
788        state,
789    })
790}
791
792#[cfg(test)]
793thread_local! {
794    /// Makes the next audio encode thread on this thread fail to start, as a
795    /// thread limit does.
796    pub(crate) static REFUSE_AUDIO_ENCODE_THREAD: std::cell::Cell<bool> =
797        const { std::cell::Cell::new(false) };
798}
799
800/// Encodes blocks into packets until the channel closes or a `Finish` lands.
801fn run_audio_encoder(
802    mut config: AudioEncoderConfig,
803    commands: &Receiver<AudioCommand>,
804    state: &Arc<Mutex<AudioQueueState>>,
805) {
806    let mut encoded_samples = 0_u64;
807    loop {
808        let Ok(command) = commands.recv() else {
809            // The sink dropped the sender without finishing (abort or cancel); nothing will read
810            // the queue.
811            return;
812        };
813        let (channels, covered, last) = match command {
814            AudioCommand::Block(channels, covered) => (Some(channels), covered, false),
815            AudioCommand::Finish(covered) => (None, covered, true),
816        };
817        let mut stages = EncodeStageTimings::default();
818        let outcome = encode_audio_block(
819            &mut config,
820            channels.as_deref(),
821            &mut encoded_samples,
822            &mut stages,
823        );
824        let mut guard = state
825            .lock()
826            .unwrap_or_else(std::sync::PoisonError::into_inner);
827        guard.stages.audio_send_frame = guard
828            .stages
829            .audio_send_frame
830            .saturating_add(stages.audio_send_frame);
831        guard.stages.audio_receive_packet = guard
832            .stages
833            .audio_receive_packet
834            .saturating_add(stages.audio_receive_packet);
835        match outcome {
836            Ok(packets) => guard.packets.extend(packets),
837            Err(error) => {
838                if guard.failure.is_none() {
839                    guard.failure = Some(error.to_string());
840                }
841                guard.finished = true;
842                return;
843            }
844        }
845        guard.consumed_ticks = guard.consumed_ticks.max(covered);
846        if last {
847            guard.finished = true;
848            return;
849        }
850    }
851}
852
853/// Encodes one AAC frame, or with none flushes the encoder (which is where
854/// its delayed tail packets come from), and drains whatever packets that
855/// produced.
856fn encode_audio_block(
857    config: &mut AudioEncoderConfig,
858    channels: Option<&[Vec<f32>]>,
859    encoded_samples: &mut u64,
860    stages: &mut EncodeStageTimings,
861) -> Result<Vec<QueuedPacket>, EncodeError> {
862    let last = channels.is_none();
863    if let Some(channels) = channels {
864        let pts = i64::try_from(*encoded_samples).map_err(|_| EncodeError::InvalidTimestamp)?;
865        let frame = alpha2_ffmpeg_bridge::audio_frame_f32p(
866            channels,
867            i32::try_from(config.sample_rate).unwrap_or(i32::MAX),
868            pts,
869        )
870        .map_err(encode_adapter_error)?;
871        timed(&mut stages.audio_send_frame, || {
872            config.encoder.send_frame(Some(&frame))
873        })
874        .map_err(|error| encode_adapter_error_at("audio encoder rejected a frame", error))?;
875        let samples = channels.first().map_or(0, Vec::len);
876        *encoded_samples =
877            encoded_samples.saturating_add(u64::try_from(samples).unwrap_or(u64::MAX));
878    } else {
879        timed(&mut stages.audio_send_frame, || {
880            config.encoder.send_frame(None)
881        })
882        .map_err(encode_adapter_error)?;
883    }
884    let mut packets = Vec::new();
885    loop {
886        let received = timed(&mut stages.audio_receive_packet, || {
887            config.encoder.receive_packet()
888        });
889        match received {
890            Ok(mut packet) => {
891                packet.set_stream_index(config.stream_index);
892                packet.rescale_ts(config.encoder_time_base, config.stream_time_base);
893                let timestamp = if packet.dts == ffi::AV_NOPTS_VALUE {
894                    packet.pts
895                } else {
896                    packet.dts
897                };
898                packets.push(QueuedPacket {
899                    sort_ticks: av_rescale_q(timestamp, config.stream_time_base, tick_base()),
900                    packet,
901                });
902            }
903            Err(RsmpegError::EncoderDrainError) if !last => return Ok(packets),
904            Err(RsmpegError::EncoderFlushedError | RsmpegError::EncoderDrainError) if last => {
905                return Ok(packets);
906            }
907            Err(error) => {
908                return Err(encode_adapter_error_at(
909                    "audio encoder packet drain failed",
910                    error,
911                ));
912            }
913        }
914    }
915}
916
917impl EncodeSink for NativeAvEncodeSink {
918    fn write(&mut self, event: EncodeEvent, cancel: &CancelToken) -> Result<(), EncodeError> {
919        if cancel.is_cancelled() {
920            self.cleanup()?;
921            return Err(EncodeError::Cancelled);
922        }
923        let ticks = match &event {
924            EncodeEvent::Video(frame) => frame.pts.ticks(),
925            EncodeEvent::Audio(block) => block.start_time().ticks(),
926        };
927        if self
928            .last_input_ticks
929            .is_some_and(|previous| ticks < previous)
930        {
931            return Err(EncodeError::InvalidTimestamp);
932        }
933        match event {
934            EncodeEvent::Video(frame) => self.write_video(&frame)?,
935            EncodeEvent::Audio(block) => self.write_audio(&block)?,
936        }
937        self.last_input_ticks = Some(ticks);
938        Ok(())
939    }
940    fn finish(&mut self, cancel: &CancelToken) -> Result<(), EncodeError> {
941        if cancel.is_cancelled() {
942            self.cleanup()?;
943            return Err(EncodeError::Cancelled);
944        }
945        // Let the thread flush its encoder and wait, so every audio packet exists before the muxer
946        // decides the final interleave.
947        let covered = self.audio_input.received_ticks();
948        for frame in self.audio_input.finish() {
949            self.audio.send(AudioCommand::Block(frame, covered))?;
950        }
951        self.audio.send(AudioCommand::Finish(covered))?;
952        self.audio.shutdown();
953        if let Some(error) = self.audio.lock().failure.clone() {
954            return Err(EncodeError::Adapter(error));
955        }
956        timed(&mut self.stages.video_send_frame, || {
957            self.video_encoder.send_frame(None)
958        })
959        .map_err(encode_adapter_error)?;
960        self.drain_video_packets(true)?;
961        self.write_ready_packets(true)?;
962        beat_encode_call();
963        self.output
964            .as_mut()
965            .ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?
966            .write_trailer()
967            .map_err(encode_adapter_error)?;
968        // A write error reported only at close fails the export here, before the rename.
969        if let Some(output) = self.output.take() {
970            beat_encode_call();
971            output
972                .close()
973                .map_err(|error| encode_adapter_error_at("closing the output", error))?;
974        }
975        crate::partial::publish(&self.partial_path, &self.final_path)?;
976        self.complete = true;
977        Ok(())
978    }
979    fn abort(&mut self) -> Result<(), EncodeError> {
980        self.cleanup()
981    }
982}
983
984impl Drop for NativeAvEncodeSink {
985    fn drop(&mut self) {
986        if !self.complete {
987            let _ = self.cleanup();
988        }
989    }
990}
991
992impl NativeWavEncodeSink {
993    /// Creates a non-overwriting sibling-partial WAV output of `channels` in
994    /// the count's default layout.
995    ///
996    /// # Errors
997    /// Returns an error for invalid PCM format, existing paths, or filesystem
998    /// failures.
999    pub fn create(
1000        path: impl AsRef<Path>,
1001        sample_rate: u32,
1002        channels: u16,
1003    ) -> Result<Self, EncodeError> {
1004        let speakers = SpeakerLayout::default_for(usize::from(channels))
1005            .ok_or_else(|| EncodeError::Adapter("invalid WAV format".into()))?;
1006        Self::create_with_speakers(path, sample_rate, speakers)
1007    }
1008
1009    /// Creates a WAV output whose blocks feed `speakers`.
1010    ///
1011    /// The count's default layout gets a plain header, which every reader
1012    /// takes as that layout. Any other is written as `WAVE_FORMAT_EXTENSIBLE`
1013    /// with its channel mask, the planes in the mask's bit order, which WAVE
1014    /// requires.
1015    ///
1016    /// # Errors
1017    /// Returns an error for invalid PCM format, a layout a WAVE channel mask
1018    /// cannot say (an unnamed or repeated speaker, or one past the 18 WAVE
1019    /// positions), existing paths, or filesystem failures.
1020    pub fn create_with_speakers(
1021        path: impl AsRef<Path>,
1022        sample_rate: u32,
1023        speakers: SpeakerLayout,
1024    ) -> Result<Self, EncodeError> {
1025        let channels = u16::try_from(speakers.channels())
1026            .ok()
1027            .filter(|_| sample_rate != 0 && speakers.channels() <= MAX_AUDIO_CHANNELS)
1028            .ok_or_else(|| EncodeError::Adapter("invalid WAV format".into()))?;
1029        let (channel_mask, plane_order) = wav_channel_layout(&speakers)?;
1030        let final_path = path.as_ref().to_owned();
1031        let partial = PartialOutput::claim(&final_path)?;
1032        let partial_path = partial.path().to_owned();
1033        let mut file = BufWriter::with_capacity(WAV_WRITE_BUFFER, partial.open_file()?);
1034        file.write_all(&wav_header_with_mask(
1035            sample_rate,
1036            channels,
1037            channel_mask,
1038            0,
1039        )?)
1040        .map_err(encode_adapter_error)?;
1041        let partial_lock = partial.keep();
1042        Ok(Self {
1043            final_path,
1044            partial_path,
1045            _partial_lock: partial_lock,
1046            file: Some(file),
1047            sample_rate,
1048            channels,
1049            speakers,
1050            channel_mask,
1051            plane_order,
1052            frames_written: 0,
1053            complete: false,
1054            interleaved: Vec::new(),
1055        })
1056    }
1057
1058    pub(crate) fn write_block(&mut self, block: &AudioBlock) -> Result<(), EncodeError> {
1059        if block.sample_rate != self.sample_rate
1060            || block.channels.len() != usize::from(self.channels)
1061            || block.channels.is_empty()
1062        {
1063            return Err(EncodeError::Adapter("audio block format changed".into()));
1064        }
1065        if block.speakers != self.speakers {
1066            return Err(EncodeError::Adapter("audio block speakers changed".into()));
1067        }
1068        let frames = block.channels[0].len();
1069        if block.channels.iter().any(|channel| channel.len() != frames) {
1070            return Err(EncodeError::Adapter("audio channel lengths differ".into()));
1071        }
1072        if u64::try_from(block.start_sample).ok() != Some(self.frames_written) {
1073            return Err(EncodeError::InvalidTimestamp);
1074        }
1075        // Refused before writing, rather than after writing 4 GiB that is then removed.
1076        let frames_after = self
1077            .frames_written
1078            .saturating_add(u64::try_from(frames).unwrap_or(u64::MAX));
1079        self.data_bytes(frames_after)?;
1080        let file = self
1081            .file
1082            .as_mut()
1083            .ok_or_else(|| EncodeError::Adapter("WAV sink is closed".into()))?;
1084        self.interleaved.clear();
1085        for frame in 0..frames {
1086            for plane in &self.plane_order {
1087                self.interleaved
1088                    .extend_from_slice(&block.channels[*plane][frame].to_le_bytes());
1089            }
1090        }
1091        file.write_all(&self.interleaved)
1092            .map_err(encode_adapter_error)?;
1093        self.frames_written = self
1094            .frames_written
1095            .saturating_add(u64::try_from(frames).unwrap_or(u64::MAX));
1096        Ok(())
1097    }
1098
1099    /// The `data` chunk's size once `frames` frames are written, when the
1100    /// RIFF chunk holding it stays within its 32-bit size.
1101    fn data_bytes(&self, frames: u64) -> Result<u32, EncodeError> {
1102        let overhead = if self.channel_mask.is_some() {
1103            WAV_EXTENSIBLE_OVERHEAD
1104        } else {
1105            WAV_PLAIN_OVERHEAD
1106        };
1107        #[cfg(test)]
1108        let limit = RIFF_LIMIT.with(std::cell::Cell::get);
1109        #[cfg(not(test))]
1110        let limit = u64::from(u32::MAX);
1111        frames
1112            .checked_mul(u64::from(self.channels))
1113            .and_then(|samples| samples.checked_mul(4))
1114            .filter(|bytes| bytes.saturating_add(u64::from(overhead)) <= limit)
1115            .and_then(|bytes| u32::try_from(bytes).ok())
1116            .ok_or_else(|| {
1117                EncodeError::Adapter(format!(
1118                    "{frames} frames of {} channels are more than the 4 GiB a WAV file holds",
1119                    self.channels
1120                ))
1121            })
1122    }
1123
1124    pub(crate) fn commit(&mut self) -> Result<(), EncodeError> {
1125        let data_bytes = self.data_bytes(self.frames_written)?;
1126        let header = wav_header_with_mask(
1127            self.sample_rate,
1128            self.channels,
1129            self.channel_mask,
1130            data_bytes,
1131        )?;
1132        let mut file = self
1133            .file
1134            .take()
1135            .ok_or_else(|| EncodeError::Adapter("WAV sink is closed".into()))?;
1136        file.seek(SeekFrom::Start(0))
1137            .and_then(|_| file.write_all(&header))
1138            .and_then(|()| file.flush())
1139            .map_err(encode_adapter_error)?;
1140        drop(file);
1141        crate::partial::publish(&self.partial_path, &self.final_path)?;
1142        self.complete = true;
1143        Ok(())
1144    }
1145}
1146
1147impl AudioBlockSink for NativeWavEncodeSink {
1148    fn write_audio(&mut self, block: &AudioBlock, cancel: &CancelToken) -> Result<(), EncodeError> {
1149        if cancel.is_cancelled() {
1150            self.abort_audio()?;
1151            return Err(EncodeError::Cancelled);
1152        }
1153        self.write_block(block)
1154    }
1155    fn finish_audio(&mut self, cancel: &CancelToken) -> Result<(), EncodeError> {
1156        if cancel.is_cancelled() {
1157            self.abort_audio()?;
1158            return Err(EncodeError::Cancelled);
1159        }
1160        self.commit()
1161    }
1162    fn abort_audio(&mut self) -> Result<(), EncodeError> {
1163        // Close without writing out samples the file is about to lose.
1164        drop(self.file.take().map(BufWriter::into_parts));
1165        match fs::remove_file(&self.partial_path) {
1166            Ok(()) => Ok(()),
1167            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
1168            Err(error) => Err(encode_adapter_error(error)),
1169        }
1170    }
1171}
1172
1173impl EncodeSink for NativeWavEncodeSink {
1174    fn write(&mut self, event: EncodeEvent, cancel: &CancelToken) -> Result<(), EncodeError> {
1175        match event {
1176            EncodeEvent::Audio(block) => self.write_audio(&block, cancel),
1177            EncodeEvent::Video(_) => {
1178                Err(EncodeError::Adapter("WAV sink accepts audio only".into()))
1179            }
1180        }
1181    }
1182    fn finish(&mut self, cancel: &CancelToken) -> Result<(), EncodeError> {
1183        self.finish_audio(cancel)
1184    }
1185    fn abort(&mut self) -> Result<(), EncodeError> {
1186        self.abort_audio()
1187    }
1188}
1189
1190impl Drop for NativeWavEncodeSink {
1191    fn drop(&mut self) {
1192        if !self.complete {
1193            let _ = self.abort_audio();
1194        }
1195    }
1196}
1197
1198/// The channel mask a WAV of `speakers` is written with, `None` for the
1199/// count's default layout, and the plane each interleaved channel comes from.
1200fn wav_channel_layout(speakers: &SpeakerLayout) -> Result<(Option<u32>, Vec<usize>), EncodeError> {
1201    let identity = (0..speakers.channels()).collect::<Vec<_>>();
1202    if SpeakerLayout::default_for(speakers.channels()) == Some(*speakers) {
1203        return Ok((None, identity));
1204    }
1205    let mut mask = 0_u32;
1206    for speaker in speakers.speakers() {
1207        let bit = speaker.ffmpeg_channel();
1208        if bit >= 18 || mask & (1 << bit) != 0 {
1209            return Err(EncodeError::Adapter(format!(
1210                "a WAVE channel mask cannot say the speakers {}",
1211                speakers.describe()
1212            )));
1213        }
1214        mask |= 1 << bit;
1215    }
1216    let mut order = identity;
1217    order.sort_by_key(|plane| speakers.speakers()[*plane].ffmpeg_channel());
1218    Ok((Some(mask), order))
1219}
1220
1221/// The RIFF chunk's size beyond its `data` chunk's samples, plain header.
1222const WAV_PLAIN_OVERHEAD: u32 = 36;
1223/// The same with a `WAVE_FORMAT_EXTENSIBLE` header.
1224const WAV_EXTENSIBLE_OVERHEAD: u32 = 60;
1225
1226#[cfg(test)]
1227thread_local! {
1228    /// The most bytes a RIFF chunk holds, for WAV sinks on this thread: a
1229    /// test lowers it from `u32::MAX`.
1230    pub(crate) static RIFF_LIMIT: std::cell::Cell<u64> =
1231        const { std::cell::Cell::new(u32::MAX as u64) };
1232}
1233
1234fn riff_size(data_bytes: u32, overhead: u32) -> Result<u32, EncodeError> {
1235    data_bytes
1236        .checked_add(overhead)
1237        .ok_or_else(|| EncodeError::Adapter("WAV exceeds the RIFF size limit".into()))
1238}
1239
1240/// A WAV header: a plain IEEE-float one, or `WAVE_FORMAT_EXTENSIBLE` with
1241/// `channel_mask` and the IEEE-float subformat.
1242///
1243/// # Errors
1244/// When the RIFF chunk's size does not fit its 32 bits.
1245pub(crate) fn wav_header_with_mask(
1246    sample_rate: u32,
1247    channels: u16,
1248    channel_mask: Option<u32>,
1249    data_bytes: u32,
1250) -> Result<Vec<u8>, EncodeError> {
1251    let Some(mask) = channel_mask else {
1252        return Ok(wav_header(sample_rate, channels, data_bytes)?.to_vec());
1253    };
1254    let plain = wav_header(sample_rate, channels, 0)?;
1255    let mut header = Vec::with_capacity(68);
1256    header.extend_from_slice(b"RIFF");
1257    header.extend_from_slice(&riff_size(data_bytes, WAV_EXTENSIBLE_OVERHEAD)?.to_le_bytes());
1258    header.extend_from_slice(b"WAVEfmt ");
1259    header.extend_from_slice(&40_u32.to_le_bytes());
1260    header.extend_from_slice(&0xfffe_u16.to_le_bytes());
1261    // Channels, rate, byte rate, block align and bits per sample, as the plain header has them.
1262    header.extend_from_slice(&plain[22..36]);
1263    header.extend_from_slice(&22_u16.to_le_bytes());
1264    header.extend_from_slice(&32_u16.to_le_bytes());
1265    header.extend_from_slice(&mask.to_le_bytes());
1266    // KSDATAFORMAT_SUBTYPE_IEEE_FLOAT.
1267    header.extend_from_slice(&[
1268        0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x80, 0x00, 0x00, 0xaa, 0x00, 0x38, 0x9b,
1269        0x71,
1270    ]);
1271    header.extend_from_slice(b"data");
1272    header.extend_from_slice(&data_bytes.to_le_bytes());
1273    Ok(header)
1274}
1275
1276/// A plain IEEE-float WAV header.
1277///
1278/// # Errors
1279/// When the RIFF chunk's size does not fit its 32 bits.
1280pub(crate) fn wav_header(
1281    sample_rate: u32,
1282    channels: u16,
1283    data_bytes: u32,
1284) -> Result<[u8; 44], EncodeError> {
1285    let mut header = [0_u8; 44];
1286    header[0..4].copy_from_slice(b"RIFF");
1287    header[4..8].copy_from_slice(&riff_size(data_bytes, WAV_PLAIN_OVERHEAD)?.to_le_bytes());
1288    header[8..12].copy_from_slice(b"WAVE");
1289    header[12..16].copy_from_slice(b"fmt ");
1290    header[16..20].copy_from_slice(&16_u32.to_le_bytes());
1291    header[20..22].copy_from_slice(&3_u16.to_le_bytes());
1292    header[22..24].copy_from_slice(&channels.to_le_bytes());
1293    header[24..28].copy_from_slice(&sample_rate.to_le_bytes());
1294    let byte_rate = sample_rate
1295        .saturating_mul(u32::from(channels))
1296        .saturating_mul(4);
1297    header[28..32].copy_from_slice(&byte_rate.to_le_bytes());
1298    header[32..34].copy_from_slice(&channels.saturating_mul(4).to_le_bytes());
1299    header[34..36].copy_from_slice(&32_u16.to_le_bytes());
1300    header[36..40].copy_from_slice(b"data");
1301    header[40..44].copy_from_slice(&data_bytes.to_le_bytes());
1302    Ok(header)
1303}