use std::collections::VecDeque;
use std::fs;
use std::path::{Path, PathBuf};
use alpha2_ffmpeg_bridge::OutputFile;
use alpha2_media::{Container, EncodeError};
use rsmpeg::avcodec::{AVCodecContext, AVCodecParameters, AVPacket};
use rsmpeg::avutil::AVRational;
use rsmpeg::ffi;
use crate::audio_input::{
AudioStreamConfig, EncoderAudioInput, OpenedAudioEncoder, open_audio_encoder,
};
use crate::convert::{encode_adapter_error, encode_adapter_error_at};
use crate::encode::{
AudioCommand, AudioEncodeThread, AudioEncoderConfig, EncodeStageTimings, QueuedPacket,
spawn_audio_encoder, timed,
};
use crate::partial::PartialOutput;
pub(crate) struct VideoStream {
pub(crate) parameters: AVCodecParameters,
pub(crate) time_base: AVRational,
pub(crate) copied: bool,
}
pub(crate) enum AudioStream {
Encode(AudioStreamConfig),
Copy {
parameters: AVCodecParameters,
time_base: AVRational,
},
}
enum MuxedAudioTrack {
Encoded(MuxedAudio),
Copied {
stream_index: i32,
source_time_base: AVRational,
stream_time_base: AVRational,
},
}
pub(crate) struct Muxer {
final_path: PathBuf,
partial_path: PathBuf,
_partial_lock: crate::partial::PartialLock,
format: Option<OutputFile>,
stream_index: i32,
stream_time_base: AVRational,
complete: bool,
audio: Option<MuxedAudioTrack>,
audio_stream_index: i32,
audio_packets: VecDeque<QueuedPacket>,
pending_header: Option<PendingHeader>,
parameter_sets: Option<Vec<u8>>,
audio_rescale: Option<(AVRational, AVRational)>,
}
struct MuxedAudio {
thread: AudioEncodeThread,
input: EncoderAudioInput,
}
impl Muxer {
#[cfg_attr(
not(all(feature = "hwencode", any(target_os = "linux", target_os = "macos"))),
allow(dead_code)
)]
pub(crate) fn open(
path: &Path,
container: Container,
video: &AVCodecContext,
video_time_base: AVRational,
audio: Option<AudioStreamConfig>,
) -> Result<Self, EncodeError> {
Self::open_streams(
path,
container,
VideoStream {
parameters: video.extract_codecpar(),
time_base: video_time_base,
copied: false,
},
audio.map(AudioStream::Encode),
)
}
pub(crate) fn open_streams(
path: &Path,
container: Container,
video: VideoStream,
audio: Option<AudioStream>,
) -> Result<Self, EncodeError> {
let mut unwritten = Unwritten::open(path, container, video, audio)?;
unwritten
.format
.write_header(&mut crate::convert::muxer_header_options(container))
.map_err(encode_adapter_error)?;
let stream_time_base = unwritten.time_base(unwritten.stream_index);
let audio_stream_time_base = unwritten.time_base(unwritten.audio_stream_index);
unwritten.start(stream_time_base, audio_stream_time_base, None)
}
#[cfg_attr(
not(any(test, all(feature = "hwencode", target_os = "windows"))),
allow(dead_code)
)]
pub(crate) fn open_deferred(
path: &Path,
container: Container,
video: &AVCodecContext,
video_time_base: AVRational,
audio: Option<AudioStreamConfig>,
) -> Result<Self, EncodeError> {
let unwritten = Unwritten::open(
path,
container,
VideoStream {
parameters: video.extract_codecpar(),
time_base: video_time_base,
copied: false,
},
audio.map(AudioStream::Encode),
)?;
let audio_time_base = unwritten
.opened_audio
.as_ref()
.map(|opened| opened.encoder_time_base);
unwritten.start(
video_time_base,
audio_time_base.unwrap_or(video_time_base),
Some(PendingHeader {
audio_time_base,
container,
}),
)
}
fn write_deferred_header(&mut self, packet: &mut AVPacket) -> Result<(), EncodeError> {
let Some(pending) = self.pending_header.as_ref() else {
return Ok(());
};
let parameter_sets =
alpha2_ffmpeg_bridge::take_packet_side_data(packet, ffi::AV_PKT_DATA_NEW_EXTRADATA)
.filter(|sets| !sets.is_empty())
.map(|sets| distinct_parameter_sets(&sets))
.ok_or_else(|| {
EncodeError::Adapter(
"the encoder's first packet carried no parameter sets for the header"
.into(),
)
})?;
let (audio_time_base, container) = (pending.audio_time_base, pending.container);
let format = self
.format
.as_mut()
.ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?;
alpha2_ffmpeg_bridge::set_output_stream_extradata(
format,
usize::try_from(self.stream_index).unwrap_or(0),
¶meter_sets,
)
.map_err(|error| encode_adapter_error_at("storing the parameter sets", error))?;
crate::encode::beat_encode_call();
format
.write_header(&mut crate::convert::muxer_header_options(container))
.map_err(encode_adapter_error)?;
self.stream_time_base =
format.streams()[usize::try_from(self.stream_index).unwrap_or(0)].time_base;
if let (Some(from), Some(MuxedAudioTrack::Encoded(_))) = (audio_time_base, &self.audio) {
let to =
format.streams()[usize::try_from(self.audio_stream_index).unwrap_or(0)].time_base;
self.audio_rescale = Some((from, to));
}
self.pending_header = None;
self.parameter_sets = Some(parameter_sets);
Ok(())
}
pub(crate) fn write_copied_audio(&mut self, packet: &mut AVPacket) -> Result<(), EncodeError> {
let Some(MuxedAudioTrack::Copied {
stream_index,
source_time_base,
stream_time_base,
}) = self.audio
else {
return Err(EncodeError::Adapter("this output copies no audio".into()));
};
packet.set_stream_index(stream_index);
packet.rescale_ts(source_time_base, stream_time_base);
self.format
.as_mut()
.ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?
.interleaved_write_frame(packet)
.map_err(|error| encode_adapter_error_at("packet write failed", error))
}
pub(crate) fn write_video(
&mut self,
packet: &mut AVPacket,
encoder_time_base: AVRational,
) -> Result<(), EncodeError> {
if self.pending_header.is_some() {
self.write_deferred_header(packet)?;
} else if let Some(header) = self.parameter_sets.as_deref() {
let repeated =
alpha2_ffmpeg_bridge::take_packet_side_data(packet, ffi::AV_PKT_DATA_NEW_EXTRADATA);
if repeated.is_some_and(|sets| distinct_parameter_sets(&sets) != header) {
return Err(EncodeError::Adapter(
"the encoder changed its parameter sets mid-stream, which one stream \
header cannot carry"
.into(),
));
}
}
packet.set_stream_index(self.stream_index);
fill_frame_duration(packet, 1);
packet.rescale_ts(encoder_time_base, self.stream_time_base);
crate::encode::beat_encode_call();
self.format
.as_mut()
.ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?
.interleaved_write_frame(packet)
.map_err(|error| encode_adapter_error_at("packet write failed", error))
}
pub(crate) fn write_audio(
&mut self,
block: &alpha2_media::AudioBlock,
stages: &mut EncodeStageTimings,
) -> Result<(), EncodeError> {
let Some(MuxedAudioTrack::Encoded(state)) = self.audio.as_mut() else {
return Err(EncodeError::Adapter(
"this output encodes no audio stream".into(),
));
};
let frames = state.input.push(block)?;
let covered = state.input.received_ticks();
for frame in frames {
state.thread.send(AudioCommand::Block(frame, covered))?;
}
self.write_ready_audio(false, stages)
}
pub(crate) fn finish_audio(&mut self) -> Result<(), EncodeError> {
let Some(MuxedAudioTrack::Encoded(state)) = self.audio.as_mut() else {
return Ok(());
};
let covered = state.input.received_ticks();
for frame in state.input.finish() {
state.thread.send(AudioCommand::Block(frame, covered))?;
}
state.thread.send(AudioCommand::Finish(covered))?;
state.thread.shutdown();
if let Some(error) = state.thread.lock().failure.clone() {
return Err(EncodeError::Adapter(error));
}
Ok(())
}
pub(crate) fn write_ready_audio(
&mut self,
flushing: bool,
stages: &mut EncodeStageTimings,
) -> Result<(), EncodeError> {
let Some(MuxedAudioTrack::Encoded(audio)) = self.audio.as_ref() else {
return Ok(());
};
{
let mut state = audio.thread.lock();
if let Some(error) = state.failure.clone() {
return Err(EncodeError::Adapter(error));
}
if flushing && !state.finished {
return Err(EncodeError::Adapter(
"audio encoder stopped before it finished".into(),
));
}
self.audio_packets.extend(state.packets.drain(..));
}
if self.pending_header.is_some() {
return Ok(());
}
let rescale = self.audio_rescale;
while let Some(mut queued) = self.audio_packets.pop_front() {
if let Some((from, to)) = rescale {
queued.packet.rescale_ts(from, to);
}
let format = self
.format
.as_mut()
.ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?;
timed(&mut stages.audio_mux_write, || {
format.interleaved_write_frame(&mut queued.packet)
})
.map_err(|error| encode_adapter_error_at("packet write failed", error))?;
}
Ok(())
}
#[cfg_attr(
not(all(
feature = "hwencode",
any(target_os = "linux", target_os = "macos", target_os = "windows")
)),
allow(dead_code)
)]
pub(crate) fn with_audio_timings(&self, stages: EncodeStageTimings) -> EncodeStageTimings {
let Some(MuxedAudioTrack::Encoded(audio)) = self.audio.as_ref() else {
return stages;
};
let thread = audio.thread.lock().stages;
EncodeStageTimings {
audio_send_frame: thread.audio_send_frame,
audio_receive_packet: thread.audio_receive_packet,
..stages
}
}
pub(crate) fn close(&mut self) -> Result<(), EncodeError> {
if self.pending_header.is_some() {
return Err(EncodeError::Adapter(
"the encoder produced no parameter sets: no video packet reached the header".into(),
));
}
crate::encode::beat_encode_call();
self.format
.as_mut()
.ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?
.write_trailer()
.map_err(encode_adapter_error)?;
if let Some(output) = self.format.take() {
crate::encode::beat_encode_call();
output
.close()
.map_err(|error| encode_adapter_error_at("closing the output", error))?;
}
crate::partial::publish(&self.partial_path, &self.final_path)?;
self.complete = true;
Ok(())
}
pub(crate) fn abort(&mut self) {
if let Some(MuxedAudioTrack::Encoded(state)) = self.audio.as_mut() {
state.thread.shutdown();
}
if !self.complete {
drop(self.format.take());
let _ = fs::remove_file(&self.partial_path);
}
}
}
impl Drop for Muxer {
fn drop(&mut self) {
self.abort();
}
}
struct Unwritten {
final_path: PathBuf,
format: OutputFile,
partial: PartialOutput,
stream_index: i32,
audio_stream_index: i32,
opened_audio: Option<OpenedAudioEncoder>,
copied_audio: Option<AVRational>,
}
impl Unwritten {
fn open(
path: &Path,
container: Container,
video: VideoStream,
audio: Option<AudioStream>,
) -> Result<Self, EncodeError> {
let final_path = path.to_owned();
let partial = PartialOutput::claim(&final_path)?;
let mut format = partial.open_muxer(container)?;
let global_header =
format.oformat().flags & i32::try_from(ffi::AVFMT_GLOBALHEADER).unwrap_or(0) != 0;
if !global_header {
return Err(EncodeError::Adapter(format!(
"the muxer for {container} has no global header"
)));
}
let stream_index;
{
let mut stream = format.new_stream();
stream.set_time_base(video.time_base);
let (codec_id, tag) = (video.parameters.codec_id, video.parameters.codec_tag);
stream.set_codecpar(video.parameters);
alpha2_ffmpeg_bridge::set_stream_codec_tag(
&mut stream,
crate::convert::container_codec_tag(container, codec_id, tag, video.copied),
);
stream_index = stream.index;
}
let mut opened_audio = None;
let mut copied_audio = None;
let mut audio_stream_index = 0_i32;
match audio {
Some(AudioStream::Encode(config)) => {
if !container.offers_audio(config.codec.model()) {
return Err(EncodeError::Adapter(format!(
"{container} does not offer audio codec {}",
config.codec.as_str()
)));
}
let opened = open_audio_encoder(config, global_header)?;
let mut stream = format.new_stream();
stream.set_time_base(opened.encoder_time_base);
stream.set_codecpar(opened.encoder.extract_codecpar());
audio_stream_index = stream.index;
opened_audio = Some(opened);
}
Some(AudioStream::Copy {
parameters,
time_base,
}) => {
let mut stream = format.new_stream();
stream.set_time_base(time_base);
let (codec_id, tag) = (parameters.codec_id, parameters.codec_tag);
stream.set_codecpar(parameters);
alpha2_ffmpeg_bridge::set_stream_codec_tag(
&mut stream,
crate::convert::container_codec_tag(container, codec_id, tag, true),
);
audio_stream_index = stream.index;
copied_audio = Some(time_base);
}
None => {}
}
Ok(Self {
final_path,
format,
partial,
stream_index,
audio_stream_index,
opened_audio,
copied_audio,
})
}
fn time_base(&self, index: i32) -> AVRational {
self.format.streams()[usize::try_from(index).unwrap_or(0)].time_base
}
fn start(
self,
stream_time_base: AVRational,
audio_stream_time_base: AVRational,
pending_header: Option<PendingHeader>,
) -> Result<Muxer, EncodeError> {
let Self {
final_path,
partial,
format,
stream_index,
audio_stream_index,
opened_audio,
copied_audio,
} = self;
let partial_path = partial.path().to_owned();
let audio = match (opened_audio, copied_audio) {
(Some(opened), _) => Some(MuxedAudioTrack::Encoded(MuxedAudio {
input: opened.input,
thread: spawn_audio_encoder(AudioEncoderConfig {
encoder: opened.encoder,
stream_index: audio_stream_index,
encoder_time_base: opened.encoder_time_base,
stream_time_base: audio_stream_time_base,
sample_rate: opened.encoder_rate,
sample_format: opened.sample_format,
})?,
})),
(None, Some(source_time_base)) => Some(MuxedAudioTrack::Copied {
stream_index: audio_stream_index,
source_time_base,
stream_time_base: audio_stream_time_base,
}),
(None, None) => None,
};
let partial_lock = partial.keep();
Ok(Muxer {
final_path,
partial_path,
_partial_lock: partial_lock,
format: Some(format),
stream_index,
stream_time_base,
complete: false,
audio,
audio_stream_index,
audio_packets: VecDeque::new(),
pending_header,
parameter_sets: None,
audio_rescale: None,
})
}
}
struct PendingHeader {
audio_time_base: Option<AVRational>,
container: Container,
}
fn distinct_parameter_sets(sets: &[u8]) -> Vec<u8> {
if !(sets.starts_with(&[0, 0, 1]) || sets.starts_with(&[0, 0, 0, 1])) {
return sets.to_vec();
}
let mut starts = Vec::new();
let mut at = 0;
while at + 3 <= sets.len() {
if sets[at..at + 3] == [0, 0, 1] {
starts.push(at + 3);
at += 3;
} else {
at += 1;
}
}
let mut units: Vec<&[u8]> = Vec::new();
for (index, &start) in starts.iter().enumerate() {
let end = starts.get(index + 1).map_or(sets.len(), |next| next - 3);
let mut unit = &sets[start..end];
while let [rest @ .., 0] = unit {
unit = rest;
}
if !unit.is_empty() && !units.contains(&unit) {
units.push(unit);
}
}
units
.into_iter()
.flat_map(|unit| [&[0, 0, 0, 1][..], unit])
.flatten()
.copied()
.collect()
}
pub(crate) fn fill_frame_duration(packet: &mut AVPacket, frame_units: i64) {
if packet.duration == 0 {
packet.set_duration(frame_units);
}
}
pub(crate) fn frame_units(rate: alpha2_time::FrameRate, time_base: AVRational) -> i64 {
rsmpeg::avutil::av_rescale_q(
1,
rsmpeg::avutil::ra(
i32::try_from(rate.denominator()).unwrap_or(1),
i32::try_from(rate.numerator()).unwrap_or(i32::MAX),
),
time_base,
)
.max(1)
}