Merge pull request #1 from cadubentzen/discord

Add Discord patches

Author: mikepavone

Cargo.toml

@@ -9,14 +9,18 @@ edition = "2021"
 
 [features]
 default = []
-backend = ["anyhow", "byteorder", "thiserror", "crc32fast", "nix", "gbm", "gbm-sys", "drm", "drm-fourcc", "zerocopy"]
+backend = ["anyhow", "byteorder", "thiserror", "crc32fast", "nix", "drm", "drm-fourcc", "zerocopy"]
 vaapi = ["libva", "backend"]
 v4l2 = ["v4l2r", "backend"]
+gbm = ["dep:gbm", "gbm-sys"]
+h265 = []
+av1 = []
+c2-wrapper = ["backend"]
 
 [dependencies]
 anyhow = { version = "1.0.75", optional = true }
 byteorder = { version = "1.4.3", optional = true }
-libva = { version = "0.0.12", package = "cros-libva", optional = true }
+libva = { version = "0.0.13", package = "cros-libva", optional = true }
 v4l2r = { version = "0.0.5", package = "v4l2r", optional = true }
 log = { version = "0", features = ["release_max_level_debug"] }
 thiserror = { version = "1.0.58", optional = true }

src/backend/v4l2/encoder.rs

@@ -523,6 +523,7 @@ where
             "bitrate mode",
             match rate_control {
                 RateControl::ConstantBitrate(_) => VideoBitrateMode::ConstantBitrate,
+                RateControl::VariableBitrate { .. } => todo!(),
                 RateControl::ConstantQuality(_) => VideoBitrateMode::ConstantQuality,
             },
         )?;
@@ -1217,6 +1218,7 @@ pub(crate) mod tests {
                 timestamp: self.counter,
                 layout: self.frame_layout.clone(),
                 force_keyframe: false,
+                force_idr: false,
             };
 
             let handle = TestMmapFrame { meta: meta.clone(), frame_count: self.max_count };
@@ -1326,6 +1328,7 @@ pub(crate) mod tests {
                 timestamp: self.counter,
                 layout: layout.clone(),
                 force_keyframe: false,
+                force_idr: false,
             };
 
             let frame = DmabufFrame { fds, layout };

src/backend/vaapi/decoder.rs

@@ -191,7 +191,7 @@ impl<V: VideoFrame> VaapiDecodedHandle<V> {
 }
 
 pub struct VaapiBackend<V: VideoFrame> {
-    pub display: Rc<Display>,
+    pub display: Arc<Display>,
     pub context: Rc<Context>,
     stream_info: StreamInfo,
     // TODO: We should try to support context reuse
@@ -200,7 +200,7 @@ pub struct VaapiBackend<V: VideoFrame> {
 }
 
 impl<V: VideoFrame> VaapiBackend<V> {
-    pub(crate) fn new(display: Rc<libva::Display>, supports_context_reuse: bool) -> Self {
+    pub(crate) fn new(display: Arc<libva::Display>, supports_context_reuse: bool) -> Self {
         let init_stream_info = StreamInfo {
             format: DecodedFormat::NV12,
             coded_resolution: Resolution::from((16, 16)),

src/backend/vaapi/encoder.rs

@@ -5,6 +5,7 @@
 use std::any::Any;
 use std::marker::PhantomData;
 use std::rc::Rc;
+use std::sync::Arc;
 
 use anyhow::anyhow;
 use libva::Config;
@@ -17,6 +18,7 @@ use libva::PictureEnd;
 use libva::Surface;
 use libva::SurfaceMemoryDescriptor;
 use libva::UsageHint;
+use libva::VAEntrypoint;
 use libva::VAEntrypoint::VAEntrypointEncSlice;
 use libva::VAEntrypoint::VAEntrypointEncSliceLP;
 use libva::VAProfile;
@@ -52,11 +54,14 @@ impl From<libva::VaError> for StatelessBackendError {
 pub(crate) fn tunings_to_libva_rc<const CLAMP_MIN_QP: u32, const CLAMP_MAX_QP: u32>(
     tunings: &Tunings,
 ) -> StatelessBackendResult<libva::EncMiscParameterRateControl> {
-    let bits_per_second = tunings.rate_control.bitrate_target().unwrap_or(0);
+    let bits_per_second = tunings.rate_control.bitrate_maximum().unwrap_or(0);
+    let target_bits_per_second = tunings.rate_control.bitrate_target().unwrap_or(bits_per_second);
+
     let bits_per_second = u32::try_from(bits_per_second).map_err(|e| anyhow::anyhow!(e))?;
+    let target_bits_per_second =
+        u32::try_from(target_bits_per_second).map_err(|e| anyhow::anyhow!(e))?;
 
-    // At the moment we don't support variable bitrate therefore target 100%
-    const TARGET_PERCENTAGE: u32 = 100;
+    let target_percentage: u32 = 100 * target_bits_per_second / bits_per_second;
 
     // Window size in ms that the RC should apply to
     const WINDOW_SIZE: u32 = 1_500;
@@ -70,10 +75,10 @@ pub(crate) fn tunings_to_libva_rc<const CLAMP_MIN_QP: u32, const CLAMP_MAX_QP: u
     const RESET: u32 = 0;
 
     // Don't skip frames
-    const DISABLE_FRAME_SKIP: u32 = 1;
+    const DISABLE_FRAME_SKIP: u32 = 0;
 
     // Allow bit stuffing
-    const DISABLE_BIT_STUFFING: u32 = 0;
+    const DISABLE_BIT_STUFFING: u32 = 1;
 
     // Use default
     const MB_RATE_CONTROL: u32 = 0;
@@ -113,7 +118,7 @@ pub(crate) fn tunings_to_libva_rc<const CLAMP_MIN_QP: u32, const CLAMP_MAX_QP: u
 
     Ok(libva::EncMiscParameterRateControl::new(
         bits_per_second,
-        TARGET_PERCENTAGE,
+        target_percentage,
         WINDOW_SIZE,
         initial_qp,
         min_qp,
@@ -161,7 +166,8 @@ where
     /// VA context used for encoding.
     context: Rc<Context>,
 
-    _va_profile: VAProfile::Type,
+    va_profile: VAProfile::Type,
+    entrypoint: VAEntrypoint::Type,
     scratch_pool: VaSurfacePool<()>,
     _phantom: PhantomData<(M, H)>,
 }
@@ -172,7 +178,7 @@ where
     H: std::borrow::Borrow<Surface<M>>,
 {
     pub fn new(
-        display: Rc<Display>,
+        display: Arc<Display>,
         va_profile: VAProfile::Type,
         fourcc: Fourcc,
         coded_size: Resolution,
@@ -185,6 +191,7 @@ where
             .ok_or_else(|| StatelessBackendError::UnsupportedFormat)?;
 
         let rt_format = format_map.rt_format;
+        let entrypoint = if low_power { VAEntrypointEncSliceLP } else { VAEntrypointEncSlice };
 
         let va_config = display.create_config(
             vec![
@@ -198,7 +205,7 @@ where
                 },
             ],
             va_profile,
-            if low_power { VAEntrypointEncSliceLP } else { VAEntrypointEncSlice },
+            entrypoint,
         )?;
 
         let context = display.create_context::<M>(
@@ -210,7 +217,7 @@ where
         )?;
 
         let mut scratch_pool = VaSurfacePool::new(
-            Rc::clone(&display),
+            Arc::clone(&display),
             rt_format,
             Some(UsageHint::USAGE_HINT_ENCODER),
             coded_size,
@@ -223,7 +230,8 @@ where
             va_config,
             context,
             scratch_pool,
-            _va_profile: va_profile,
+            va_profile,
+            entrypoint,
             _phantom: Default::default(),
         })
     }
@@ -271,6 +279,51 @@ where
 
         Ok(Reconstructed(surface))
     }
+
+    pub(crate) fn supports_max_frame_size(&self) -> StatelessBackendResult<bool> {
+        let display = self.context().display();
+        let mut attrs = [libva::VAConfigAttrib {
+            type_: libva::VAConfigAttribType::VAConfigAttribMaxFrameSize,
+            value: 0,
+        }];
+
+        display.get_config_attributes(self.va_profile, self.entrypoint, &mut attrs)?;
+
+        if attrs[0].value == libva::VA_ATTRIB_NOT_SUPPORTED {
+            return Ok(false);
+        }
+        Ok(true)
+    }
+
+    pub(crate) fn supports_quality_range(&self, quality: u32) -> StatelessBackendResult<bool> {
+        let display = self.context().display();
+        let mut attrs = [libva::VAConfigAttrib {
+            type_: libva::VAConfigAttribType::VAConfigAttribEncQualityRange,
+            value: 0,
+        }];
+
+        display.get_config_attributes(self.va_profile, self.entrypoint, &mut attrs)?;
+
+        if attrs[0].value == libva::VA_ATTRIB_NOT_SUPPORTED || quality > attrs[0].value {
+            return Ok(false);
+        }
+        Ok(true)
+    }
+
+    pub(crate) fn supports_packed_header(&self, header: u32) -> StatelessBackendResult<bool> {
+        let display = self.context().display();
+        let mut attrs = [libva::VAConfigAttrib {
+            type_: libva::VAConfigAttribType::VAConfigAttribEncPackedHeaders,
+            value: 0,
+        }];
+
+        display.get_config_attributes(self.va_profile, self.entrypoint, &mut attrs)?;
+
+        if attrs[0].value == libva::VA_ATTRIB_NOT_SUPPORTED || (header & attrs[0].value) == 0 {
+            return Ok(false);
+        }
+        Ok(true)
+    }
 }
 
 impl<M, Handle> StatelessEncoderBackendImport<Handle, Handle> for VaapiBackend<M, Handle>
@@ -491,6 +544,7 @@ pub(crate) mod tests {
                 layout: self.frame_layout.clone(),
                 force_keyframe: false,
                 timestamp: self.counter,
+                force_idr: false,
             };
 
             self.counter += 1;

src/backend/vaapi/surface_pool.rs

@@ -8,6 +8,7 @@ use std::collections::BTreeMap;
 use std::collections::VecDeque;
 use std::rc::Rc;
 use std::rc::Weak;
+use std::sync::Arc;
 
 use libva::Display;
 use libva::Surface;
@@ -79,7 +80,7 @@ impl<M: SurfaceMemoryDescriptor> Drop for PooledVaSurface<M> {
 }
 
 struct VaSurfacePoolInner<M: SurfaceMemoryDescriptor> {
-    display: Rc<Display>,
+    display: Arc<Display>,
     rt_format: u32,
     usage_hint: Option<libva::UsageHint>,
     coded_resolution: Resolution,
@@ -133,7 +134,7 @@ impl<M: SurfaceMemoryDescriptor> VaSurfacePool<M> {
     /// * `usage_hint` - hint about how the surfaces from this pool will be used.
     /// * `coded_resolution` - resolution of the surfaces.
     pub fn new(
-        display: Rc<Display>,
+        display: Arc<Display>,
         rt_format: u32,
         usage_hint: Option<libva::UsageHint>,
         coded_resolution: Resolution,

src/bitstream_utils.rs

@@ -510,15 +510,18 @@ impl<W: Write> BitWriter<W> {
     }
 
     /// Immediately outputs any cached bits to [`std::io::Write`]
-    pub fn flush(&mut self) -> BitWriterResult<()> {
+    /// and returns the number of trailing bits in the last byte.
+    pub fn flush(&mut self) -> BitWriterResult<u8> {
+        let mut num_trailing_bits = 0;
         if self.nth_bit != 0 {
             self.out.write_all(&[self.curr_byte])?;
+            num_trailing_bits = 8 - self.nth_bit;
             self.nth_bit = 0;
             self.curr_byte = 0;
         }
 
         self.out.flush()?;
-        Ok(())
+        Ok(num_trailing_bits)
     }
 
     /// Returns `true` if ['Self`] hold data that wasn't written to [`std::io::Write`]

src/c2_wrapper/c2_decoder.rs

@@ -32,7 +32,7 @@ use crate::decoder::StreamInfo;
 use crate::image_processing::convert_video_frame;
 use crate::video_frame::frame_pool::FramePool;
 use crate::video_frame::frame_pool::PooledVideoFrame;
-#[cfg(feature = "vaapi")]
+#[cfg(feature = "gbm")]
 use crate::video_frame::gbm_video_frame::{GbmDevice, GbmUsage, GbmVideoFrame};
 #[cfg(feature = "v4l2")]
 use crate::video_frame::v4l2_mmap_video_frame::V4l2MmapVideoFrame;
@@ -62,7 +62,7 @@ pub trait C2DecoderBackend {
     ) -> Result<DynStatelessVideoDecoder<V>, String>;
 }
 
-#[cfg(feature = "vaapi")]
+#[cfg(feature = "gbm")]
 type AuxiliaryVideoFrame = GbmVideoFrame;
 #[cfg(feature = "v4l2")]
 type AuxiliaryVideoFrame = V4l2MmapVideoFrame;
@@ -189,7 +189,7 @@ where
                 ),
             )
         } else {
-            #[cfg(feature = "vaapi")]
+            #[cfg(feature = "gbm")]
             {
                 let gbm_device = Arc::new(
                     GbmDevice::open(PathBuf::from("/dev/dri/renderD128"))

src/c2_wrapper/c2_encoder.rs

@@ -250,6 +250,7 @@ where
                     timestamp: job.timestamp,
                     layout: Default::default(),
                     force_keyframe: false,
+                    force_idr: false,
                 };
                 #[cfg(feature = "vaapi")]
                 let encode_result = self.encoder.encode(meta, frame);

src/codec/h264/nalu_writer.rs

@@ -177,14 +177,20 @@ impl<W: Write> NaluWriter<W> {
     /// Writes a H.264 NALU header.
     pub fn write_header(&mut self, idc: u8, _type: u8) -> NaluWriterResult<()> {
         self.0.flush()?;
-        self.0.inner_mut().write_header(idc, _type)?;
+        let _num_bytes = self.0.inner_mut().write_header(idc, _type)?;
+        // self.0.bits_written += num_bytes * 8;
         Ok(())
     }
 
     /// Returns `true` if next bits will be aligned to 8
     pub fn aligned(&self) -> bool {
         !self.0.has_data_pending()
     }
+
+    /// Returns the number of trailing bits in the last byte.
+    pub fn flush(&mut self) -> NaluWriterResult<u8> {
+        Ok(self.0.flush()?)
+    }
 }
 
 #[cfg(test)]