transotf/src/av/mod.rs

use crate::av_err2str;
use ffmpeg_sys_next::AVCodecID::AV_CODEC_ID_HEVC;
use ffmpeg_sys_next::{
    av_codec_is_decoder, av_codec_is_encoder, av_codec_next, av_inv_q, av_register_all,
    avcodec_alloc_context3, avcodec_find_decoder_by_name, avcodec_find_encoder_by_name, AVCodec,
    AVCodecContext, AVCodecID, AVRational, AV_CODEC_CAP_HARDWARE,
};
use num_rational::Ratio;
use std::any::type_name;
use std::borrow::Borrow;
use std::collections::HashMap;
use std::ffi::{CStr, CString};
use std::fmt::Debug;
use std::option::Option::Some;
use std::os::raw::c_int;
use std::ptr::null;

pub mod as_ptr;
pub mod avio;
pub mod codec_parameters;
pub mod decoder;
pub mod decoder_selector;
pub mod dictionary;
pub mod dynbuf;
pub mod encoder;
pub mod encoder_selector;
pub mod format;
pub mod frame;
pub mod mov_avc;
pub mod mov_mux;
pub mod options;
pub mod packet;
pub mod resampler;
pub mod scaler;
pub mod stream;
pub mod xcoder;

static mut CACHE: Option<FfmpegInfoCache> = None;

pub fn get_best_encoder(codec: AVCodecID, hw_accel: Option<bool>) -> Option<String> {
    let encoders = get_cache().encoders.get(&codec)?;
    get_best_codec(encoders, hw_accel)
}

pub fn get_best_codec(
    codecs: &Vec<FfmpegCodec>,
    should_be_hw_accel: Option<bool>,
) -> Option<String> {
    let hw_accel = codecs.iter().find(|x| x.hwaccel);
    if hw_accel.is_some() && should_be_hw_accel != Some(false) {
        hw_accel
    } else if should_be_hw_accel != Some(true) {
        codecs.iter().find(|x| !x.hwaccel)
    } else {
        None
    }
    .map(|x| x.name.clone())
}

pub fn get_best_decoder(codec: AVCodecID, hw_accel: Option<bool>) -> Option<String> {
    let decoders = get_cache().decoders.get(&codec)?;
    get_best_codec(decoders, hw_accel)
}

pub fn decoder_codec_from_name(name: &str) -> Option<*mut AVCodec> {
    let codec = unsafe {
        let name = CString::new(name).unwrap();
        avcodec_find_decoder_by_name(name.as_ptr())
    };

    if codec.is_null() {
        return None;
    }

    Some(codec)
}

pub fn decoder_from_name(name: &str) -> Option<*mut AVCodecContext> {
    let codec = decoder_codec_from_name(name)?;

    let xcoder = unsafe { avcodec_alloc_context3(codec) };
    if xcoder.is_null() {
        return None;
    }

    Some(xcoder)
}

pub fn encoders_for_codec(codec_id: AVCodecID) -> Vec<FfmpegCodec> {
    return if let Some(codecs) = get_cache().encoders.get(&codec_id) {
        let mut codecs: Vec<FfmpegCodec> = codecs.to_vec();
        codecs.sort_by_key(|x| !x.hwaccel);
        codecs
    } else {
        vec![]
    };
}

pub fn decoders_for_codec(codec_id: AVCodecID) -> Vec<FfmpegCodec> {
    return if let Some(codecs) = get_cache().decoders.get(&codec_id) {
        let mut codecs: Vec<FfmpegCodec> = codecs.to_vec();
        codecs.sort_by_key(|x| !x.hwaccel);
        codecs
    } else {
        vec![]
    };
}

pub fn encoder_codec_from_name(name: &str) -> Option<*mut AVCodec> {
    let codec = unsafe {
        let name = CString::new(name).unwrap();
        avcodec_find_encoder_by_name(name.as_ptr())
    };

    if codec.is_null() {
        return None;
    }

    Some(codec)
}

pub fn encoder_from_name(name: &str) -> Option<*mut AVCodecContext> {
    let codec = encoder_codec_from_name(name)?;

    let encoder = unsafe { avcodec_alloc_context3(codec) };
    if encoder.is_null() {
        return None;
    }

    Some(encoder)
}

pub fn get_cache() -> &'static FfmpegInfoCache {
    unsafe {
        return CACHE.get_or_insert_with(|| create_cache());
    }
}

pub fn create_cache() -> FfmpegInfoCache {
    let mut info = FfmpegInfoCache::default();
    unsafe {
        let mut desc = av_codec_next(null());
        while !desc.is_null() {
            let p_desc = *desc;
            let name = CStr::from_ptr(p_desc.name).to_str().unwrap().to_string();
            let long_name = CStr::from_ptr(p_desc.long_name)
                .to_str()
                .unwrap()
                .to_string();

            if av_codec_is_decoder(desc) == 1 {
                info.decoders
                    .entry(p_desc.id)
                    .or_default()
                    .push(FfmpegCodec {
                        name: name.clone(),
                        long_name: long_name.clone(),
                        hwaccel: (p_desc.capabilities & (AV_CODEC_CAP_HARDWARE as i32)) > 0,
                    });
            }

            if av_codec_is_encoder(desc) == 1 {
                info.encoders
                    .entry(p_desc.id)
                    .or_default()
                    .push(FfmpegCodec {
                        name,
                        long_name,
                        hwaccel: (p_desc.capabilities & (AV_CODEC_CAP_HARDWARE as i32)) > 0,
                    });
            }

            desc = av_codec_next(desc);
        }
    }

    // Remove deprecated nvenc codecs if present
    info.encoders.entry(AV_CODEC_ID_HEVC).and_modify(|item| {
        remove_if_exists(item, "h264_nvenc", vec!["nvenc", "nvenc_h264"]);
        remove_if_exists(item, "nvenc_h264", vec!["nvenc"]);
        remove_if_exists(item, "hevc_nvenc", vec!["nvenc_hevc"]);
    });

    info
}

fn remove_if_exists(list: &mut Vec<FfmpegCodec>, needle: &str, to_remove: Vec<&str>) {
    let mut found_deprecated = vec![];
    let mut found_new = false;
    let mut i = 0;
    for codec in Borrow::<Vec<_>>::borrow(list) {
        if codec.name == needle {
            found_new = true;
            break;
        }

        if to_remove.contains(&codec.name.as_str()) {
            found_deprecated.push(i)
        }

        i += 1;
    }

    if found_new {
        found_deprecated.sort();
        found_deprecated.reverse();
        for index in found_deprecated {
            list.remove(index);
        }
    }
}

pub fn init() {
    unsafe {
        av_register_all();
    }
}

#[inline]
pub(crate) fn verify_response(error: &str, resp: c_int) -> Result<c_int, String> {
    if resp < 0 {
        Err(format!("{} (code {}, {})", error, resp, av_err2str(resp)))
    } else {
        Ok(resp)
    }
}

#[derive(Default, Debug)]
pub struct FfmpegInfoCache {
    pub decoders: HashMap<AVCodecID, Vec<FfmpegCodec>>,
    pub encoders: HashMap<AVCodecID, Vec<FfmpegCodec>>,
}

#[derive(Default, Debug, Clone)]
pub struct FfmpegCodec {
    pub name: String,
    pub long_name: String,
    pub hwaccel: bool,
}

pub trait Rational {
    fn new(num: i32, den: i32) -> Self;

    fn to_av(&self) -> AVRational {
        AVRational {
            den: self.den(),
            num: self.num(),
        }
    }

    fn to_num(&self) -> Ratio<i32> {
        Ratio::new(self.num(), self.den())
    }

    fn num(&self) -> i32;
    fn den(&self) -> i32;
    fn invert(&self) -> Self;
    fn trunc(&self) -> Ratio<i32> {
        self.to_num().trunc()
    }

    fn debug(&self) -> String {
        return format!(
            "Rational[{}]({}/{})",
            type_name::<Self>(),
            self.num(),
            self.den()
        );
    }

    fn simplify(&self) -> Self;
}

impl Rational for Ratio<i32> {
    #[inline]
    fn new(num: i32, den: i32) -> Self {
        Ratio::new(num, den)
    }

    #[inline]
    fn to_num(&self) -> Ratio<i32> {
        *self
    }

    #[inline]
    fn num(&self) -> i32 {
        *self.numer()
    }

    #[inline]
    fn den(&self) -> i32 {
        *self.denom()
    }

    #[inline]
    fn invert(&self) -> Self {
        Ratio::new(self.den(), self.num())
    }

    fn simplify(&self) -> Self {
        Self::new(1, self.den() / self.num())
    }
}

impl Rational for AVRational {
    #[inline]
    fn new(num: i32, den: i32) -> Self {
        AVRational { num, den }
    }

    #[inline]
    fn to_av(&self) -> AVRational {
        *self
    }

    #[inline]
    fn num(&self) -> i32 {
        self.num
    }

    #[inline]
    fn den(&self) -> i32 {
        self.den
    }

    #[inline]
    fn invert(&self) -> Self {
        unsafe { av_inv_q(*self) }
    }

    fn simplify(&self) -> Self {
        Self::new(1, (self.den() as f64 / self.num() as f64).round() as i32)
    }
}