shrivel 0.1.3

//! HEVC encoder backends and the FFmpeg arguments each one needs.

use std::process::{Command, Stdio};

use crate::cli::EncoderChoice;

/// A concrete HEVC encoder backend.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Encoder {
    Nvenc,
    VideoToolbox,
    Qsv,
    Amf,
    Software,
}

impl Encoder {
    /// Hardware backends tried by auto-detection, in order of preference.
    /// The software encoder is the final fallback and is not listed here.
    const HARDWARE_PRIORITY: [Encoder; 4] = [
        Encoder::Nvenc,
        Encoder::VideoToolbox,
        Encoder::Qsv,
        Encoder::Amf,
    ];

    /// Maps a CLI choice to a concrete encoder. `Auto` yields `None`.
    pub fn from_choice(choice: EncoderChoice) -> Option<Self> {
        match choice {
            EncoderChoice::Auto => None,
            EncoderChoice::Nvenc => Some(Self::Nvenc),
            EncoderChoice::Qsv => Some(Self::Qsv),
            EncoderChoice::Amf => Some(Self::Amf),
            EncoderChoice::Videotoolbox => Some(Self::VideoToolbox),
            EncoderChoice::Software => Some(Self::Software),
        }
    }

    /// The FFmpeg codec name of this encoder.
    pub fn codec_name(self) -> &'static str {
        match self {
            Self::Nvenc => "hevc_nvenc",
            Self::VideoToolbox => "hevc_videotoolbox",
            Self::Qsv => "hevc_qsv",
            Self::Amf => "hevc_amf",
            Self::Software => "libx265",
        }
    }

    /// Input-side arguments (placed before `-i`), e.g. hardware decoding.
    pub fn input_args(self) -> &'static [&'static str] {
        match self {
            Self::Nvenc => &["-hwaccel", "cuda"],
            _ => &[],
        }
    }

    /// Output-side video encoder arguments for the given quality value.
    ///
    /// `cq` uses the NVENC/x265 scale (0-51, lower is better). Backends with
    /// a different quality scale receive an approximate mapping.
    pub fn video_args(self, cq: u8) -> Vec<String> {
        let cq_str = cq.to_string();
        let args: Vec<&str> = match self {
            Self::Nvenc => vec![
                "-c:v",
                "hevc_nvenc",
                "-preset",
                "p7",
                "-tune",
                "hq",
                "-rc",
                "vbr",
                "-cq",
                &cq_str,
                "-b:v",
                "0",
                "-spatial-aq",
                "1",
                "-temporal-aq",
                "1",
                "-aq-strength",
                "4",
                "-bf",
                "3",
                "-b_ref_mode",
                "middle",
                "-rc-lookahead",
                "32",
            ],
            Self::Qsv => vec![
                "-c:v",
                "hevc_qsv",
                "-preset",
                "veryslow",
                "-global_quality",
                &cq_str,
                "-look_ahead",
                "1",
            ],
            Self::Amf => vec![
                "-c:v", "hevc_amf", "-quality", "quality", "-rc", "cqp", "-qp_i", &cq_str, "-qp_p",
                &cq_str,
            ],
            Self::VideoToolbox => {
                let q = videotoolbox_quality(cq).to_string();
                return ["-c:v", "hevc_videotoolbox", "-q:v", &q]
                    .iter()
                    .map(|s| s.to_string())
                    .collect();
            }
            Self::Software => vec!["-c:v", "libx265", "-preset", "slow", "-crf", &cq_str],
        };
        args.into_iter().map(String::from).collect()
    }

    /// Checks that this encoder actually works on this machine by encoding a
    /// single tiny synthetic frame. Presence in `ffmpeg -encoders` is not
    /// enough: hardware encoders are often compiled in but unusable.
    pub fn is_usable(self) -> bool {
        let mut cmd = Command::new("ffmpeg");
        cmd.args([
            "-hide_banner",
            "-loglevel",
            "error",
            "-nostdin",
            "-f",
            "lavfi",
            "-i",
            "color=c=black:s=640x360:r=1:d=1",
            "-frames:v",
            "1",
        ]);
        cmd.args(self.video_args(30));
        cmd.args(["-f", "null", "-"]);
        cmd.stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::null());
        cmd.status().map(|s| s.success()).unwrap_or(false)
    }

    /// Picks the first working hardware encoder, falling back to x265.
    pub fn detect() -> Self {
        Self::HARDWARE_PRIORITY
            .into_iter()
            .find(|e| e.is_usable())
            .unwrap_or(Self::Software)
    }
}

/// Approximate mapping from the CQ scale (lower is better) to VideoToolbox's
/// `-q:v` scale (0-100, higher is better).
fn videotoolbox_quality(cq: u8) -> u8 {
    (100.0 - f64::from(cq) * 1.5).round().clamp(1.0, 100.0) as u8
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn nvenc_args_carry_the_cq_value() {
        let args = Encoder::Nvenc.video_args(34);
        let pos = args.iter().position(|a| a == "-cq").unwrap();
        assert_eq!(args[pos + 1], "34");
        assert_eq!(args[0], "-c:v");
        assert_eq!(args[1], "hevc_nvenc");
    }

    #[test]
    fn software_args_use_crf() {
        let args = Encoder::Software.video_args(28);
        assert!(args.windows(2).any(|w| w == ["-crf", "28"]));
    }

    #[test]
    fn videotoolbox_quality_is_clamped() {
        assert_eq!(videotoolbox_quality(0), 100);
        assert_eq!(videotoolbox_quality(30), 55);
        assert!(videotoolbox_quality(51) >= 1);
    }

    #[test]
    fn only_nvenc_requests_hardware_decoding() {
        assert!(!Encoder::Nvenc.input_args().is_empty());
        assert!(Encoder::Software.input_args().is_empty());
        assert!(Encoder::Qsv.input_args().is_empty());
    }

    #[test]
    fn auto_has_no_concrete_encoder() {
        assert_eq!(Encoder::from_choice(EncoderChoice::Auto), None);
        assert_eq!(
            Encoder::from_choice(EncoderChoice::Software),
            Some(Encoder::Software)
        );
    }
}