lumen-engine-ffmpeg 0.2.2

FFmpeg integration for media decode, encode, muxing, and GPU interop in Lumen.
Documentation
use std::ptr;

use crate::{
    FfmpegError, Result,
    ffi::{self, AvFrame, sys},
    format::{InputContext, Packet, Rational},
};
use sys::{
    AVRounding::AV_ROUND_UP,
    AVSampleFormat::{AV_SAMPLE_FMT_FLT, AV_SAMPLE_FMT_NONE, AV_SAMPLE_FMT_S16},
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SampleFormat {
    F32,
    I16,
    Unknown,
}

#[derive(Debug, Clone)]
pub struct AudioFrame {
    pub sample_rate: u32,
    pub channels: u16,
    pub sample_format: SampleFormat,
    pub pts: Option<i64>,
    pub samples: usize,
    pub interleaved_f32: Vec<f32>,
}

#[derive(Debug, Clone, Copy)]
pub struct AudioResamplerConfig {
    pub sample_rate: u32,
    pub channels: u16,
    pub sample_format: SampleFormat,
}

impl Default for AudioResamplerConfig {
    fn default() -> Self {
        Self {
            sample_rate: 48_000,
            channels: 2,
            sample_format: SampleFormat::F32,
        }
    }
}

pub struct AudioDecoder {
    stream_index: usize,
    time_base: Rational,
    context: *mut sys::AVCodecContext,
}

unsafe impl Send for AudioDecoder {}

impl AudioDecoder {
    pub fn open(input: &InputContext, stream_index: usize) -> Result<Self> {
        let parameters = input.stream_parameters(stream_index)?;
        let codec = unsafe { sys::avcodec_find_decoder((*parameters).codec_id) };
        if codec.is_null() {
            return Err(FfmpegError::new(
                "avcodec_find_decoder",
                "no audio decoder found",
            ));
        }
        let context = unsafe { sys::avcodec_alloc_context3(codec) };
        if context.is_null() {
            return Err(FfmpegError::new(
                "avcodec_alloc_context3",
                "failed to allocate audio decoder context",
            ));
        }
        unsafe {
            ffi::check(
                sys::avcodec_parameters_to_context(context, parameters),
                "avcodec_parameters_to_context",
            )?;
            ffi::check(
                sys::avcodec_open2(context, codec, ptr::null_mut()),
                "avcodec_open2",
            )?;
        }
        Ok(Self {
            stream_index,
            time_base: input.stream_time_base(stream_index)?.into(),
            context,
        })
    }

    pub fn stream_index(&self) -> usize {
        self.stream_index
    }

    pub fn time_base(&self) -> Rational {
        self.time_base
    }

    pub fn send_packet(&mut self, packet: &Packet) -> Result<()> {
        if packet.stream_index() != self.stream_index {
            return Ok(());
        }
        unsafe {
            ffi::check(
                sys::avcodec_send_packet(self.context, packet.inner.as_ptr()),
                "avcodec_send_packet",
            )
        }
    }

    pub fn send_eof(&mut self) -> Result<()> {
        unsafe {
            ffi::check(
                sys::avcodec_send_packet(self.context, ptr::null()),
                "avcodec_send_packet",
            )
        }
    }

    pub fn flush(&mut self) {
        unsafe { sys::avcodec_flush_buffers(self.context) };
    }

    pub fn receive_frame(&mut self) -> Result<Option<DecodedAudioFrame>> {
        let mut frame = AvFrame::new()?;
        let result = unsafe { sys::avcodec_receive_frame(self.context, frame.as_mut_ptr()) };
        if result == sys::AVERROR(libc::EAGAIN) || result == sys::AVERROR_EOF {
            return Ok(None);
        }
        if result < 0 {
            return Err(ffi::error_from_code("avcodec_receive_frame", result));
        }
        Ok(Some(DecodedAudioFrame { frame }))
    }
}

impl Drop for AudioDecoder {
    fn drop(&mut self) {
        unsafe { sys::avcodec_free_context(&mut self.context) };
    }
}

pub struct DecodedAudioFrame {
    frame: AvFrame,
}

impl DecodedAudioFrame {
    pub fn samples(&self) -> usize {
        self.frame.nb_samples()
    }

    pub fn pts(&self) -> Option<i64> {
        self.frame.pts()
    }
}

pub struct AudioResampler {
    ptr: *mut sys::SwrContext,
    config: AudioResamplerConfig,
}

unsafe impl Send for AudioResampler {}

impl AudioResampler {
    pub fn new(source: &DecodedAudioFrame, config: AudioResamplerConfig) -> Result<Self> {
        let mut dst_layout = unsafe { std::mem::zeroed() };
        let mut src_layout = unsafe { (*source.frame.as_ptr()).ch_layout };
        unsafe {
            sys::av_channel_layout_default(&mut dst_layout, config.channels as i32);
            if src_layout.nb_channels == 0 {
                sys::av_channel_layout_default(
                    &mut src_layout,
                    (*source.frame.as_ptr()).ch_layout.nb_channels,
                );
            }
        }
        let mut ptr = ptr::null_mut();
        unsafe {
            ffi::check(
                sys::swr_alloc_set_opts2(
                    &mut ptr,
                    ptr::addr_of_mut!(dst_layout),
                    config.sample_format.to_av_sample_format(),
                    config.sample_rate as i32,
                    ptr::addr_of_mut!(src_layout),
                    source.frame.sample_format(),
                    (*source.frame.as_ptr()).sample_rate,
                    0,
                    ptr::null_mut(),
                ),
                "swr_alloc_set_opts2",
            )?;
            ffi::check(sys::swr_init(ptr), "swr_init")?;
        }
        Ok(Self { ptr, config })
    }

    pub fn convert(&mut self, source: &DecodedAudioFrame) -> Result<AudioFrame> {
        if self.config.sample_format != SampleFormat::F32 {
            return Err(FfmpegError::new(
                "AudioResampler::convert",
                "only interleaved f32 output is currently exposed",
            ));
        }
        let input_samples = source.frame.nb_samples() as i64;
        let delay =
            unsafe { sys::swr_get_delay(self.ptr, (*source.frame.as_ptr()).sample_rate as i64) };
        let output_capacity = unsafe {
            sys::av_rescale_rnd(
                delay + input_samples,
                self.config.sample_rate as i64,
                (*source.frame.as_ptr()).sample_rate as i64,
                AV_ROUND_UP,
            )
        } as usize;
        let channels = self.config.channels as usize;
        let mut output = vec![0_f32; output_capacity.saturating_mul(channels)];
        let mut output_ptr = output.as_mut_ptr() as *mut u8;
        let converted = unsafe {
            sys::swr_convert(
                self.ptr,
                ptr::addr_of_mut!(output_ptr),
                output_capacity as i32,
                (*source.frame.as_ptr()).data.as_ptr() as *mut *const u8,
                source.frame.nb_samples() as i32,
            )
        };
        if converted < 0 {
            return Err(ffi::error_from_code("swr_convert", converted));
        }
        let samples = converted as usize;
        output.truncate(samples.saturating_mul(channels));
        Ok(AudioFrame {
            sample_rate: self.config.sample_rate,
            channels: self.config.channels,
            sample_format: self.config.sample_format,
            pts: source.pts(),
            samples,
            interleaved_f32: output,
        })
    }
}

impl Drop for AudioResampler {
    fn drop(&mut self) {
        unsafe { sys::swr_free(&mut self.ptr) };
    }
}

impl SampleFormat {
    fn to_av_sample_format(self) -> sys::AVSampleFormat {
        match self {
            Self::F32 => AV_SAMPLE_FMT_FLT,
            Self::I16 => AV_SAMPLE_FMT_S16,
            Self::Unknown => AV_SAMPLE_FMT_NONE,
        }
    }
}