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