Implement a decoder

Author: lookback-hugotunius

build.rs

@@ -11,6 +11,9 @@ fn main() {
 
     headers.push("libavcodec/avcodec.h");
     headers.push("libavutil/opt.h");
+    headers.push("libavutil/mem.h");
+    headers.push("libavutil/imgutils.h");
+    headers.push("libavutil/pixdesc.h");
 
     let lib1 = pkg_config::probe_library("libavcodec").expect("find libavcodec");
     let lib2 = pkg_config::probe_library("libavutil").expect("find libavutil");
@@ -45,8 +48,12 @@ fn main() {
         .allowlist_item("av_frame_.*")
         .allowlist_item("av_init_packet")
         .allowlist_item("av_packet_.*")
+        .allowlist_item("av_buffer_.*")
         .allowlist_item("av_strerror")
         .allowlist_item("av_log_set_level")
+        .allowlist_item("av_malloc")
+        .allowlist_item("av_image_.*")
+        .allowlist_item("av_pix_.*")
         .allowlist_item("log_to_string.*")
         .default_enum_style(EnumVariation::Rust {
             non_exhaustive: false,

src/decoder.rs

@@ -0,0 +1,336 @@
+use std::ffi::c_void;
+use std::ptr;
+
+use super::{
+    av_log_set_callback, err_code_to_string, log_callback, set_log_level, sys, Codec, CodecKind,
+    Error, FrameRef, PixelFormat,
+};
+
+use tracing::Level;
+
+pub struct Decoder {
+    ctx: *mut sys::AVCodecContext,
+    /// Maps rotation values to the PTS of the incoming packet.
+    pts_map: PtsMap,
+}
+
+unsafe impl Send for Decoder {}
+
+struct PtsMap {
+    map: [(i64, usize); 16],
+    cur: usize,
+}
+
+pub trait DecoderPacket {
+    /// Returns
+    fn data(&mut self) -> PacketData;
+    fn pts(&self) -> i64;
+    fn rotation(&self) -> usize;
+}
+
+pub struct PacketData {
+    inner: Box<[u8]>,
+}
+
+/// A single frame of video or audio.
+pub struct DecodedFrame(*mut sys::AVFrame);
+
+impl Decoder {
+    /// Create a new decoder
+    pub fn new(codec: &Codec) -> Result<Self, Error> {
+        set_log_level(Level::DEBUG);
+        unsafe {
+            av_log_set_callback(Some(log_callback));
+        }
+
+        if codec.kind() != CodecKind::Decoder {
+            return Err(Error::CodecIsNotDecoder(codec.name));
+        }
+
+        let codec = codec.ptr;
+        let ctx: *mut sys::AVCodecContext = unsafe { sys::avcodec_alloc_context3(codec) };
+        if ctx.is_null() {
+            return Err(Error::CreateContextFailed);
+        }
+
+        let dec = Decoder {
+            ctx,
+            pts_map: PtsMap::default(),
+        };
+
+        // TODO: options
+
+        let err = unsafe { sys::avcodec_open2(ctx, codec, ptr::null_mut()) };
+        if err < 0 {
+            return Err(Error::CodecOpenError(err, err_code_to_string(err)));
+        }
+
+        Ok(dec)
+    }
+
+    /// Decode some compressed data.
+    ///
+    /// Returns an iterator over the resulting frames.
+    pub fn decode(
+        &mut self,
+        packet: &mut dyn DecoderPacket,
+    ) -> Result<impl Iterator<Item = Result<DecodedFrame, Error>> + '_, Error> {
+        let mut pkt = unsafe {
+            let pkt = sys::av_packet_alloc();
+            if pkt.is_null() {
+                return Err(Error::AlllocateFailed("av_malloc for Decoder::decode"));
+            }
+
+            let data = packet.data();
+            // The buffer used for the packet is required to have
+            // `sys::AV_INPUT_BUFFER_PADDING_SIZE` padding bytes, this is guaranteed for us by
+            // `PacketData`.
+            let len = data.inner.len();
+            // This is a fat pointer i.e. 2 words
+            let data_ptr = Box::into_raw(data.inner);
+            let buf = sys::av_buffer_create(
+                data_ptr.cast(),
+                // This might look useless, but depending on the version of libavcodec used it's
+                // required.
+                #[allow(clippy::useless_conversion)]
+                len.try_into().unwrap(),
+                Some(free_boxed_slice),
+                // We store the length of the slice as the opaque data so we can re-create the fat
+                // pointer for freeing in `free_boxed_slice`.
+                len as *mut c_void,
+                0,
+            );
+            assert!(!buf.is_null());
+            (*pkt).buf = buf;
+            (*pkt).data = data_ptr.cast();
+            (*pkt).pts = packet.pts();
+            // This should be the size of the data without the padding
+            (*pkt).size = (len as i32) - sys::AV_INPUT_BUFFER_PADDING_SIZE as i32;
+
+            pkt
+        };
+        self.pts_map.set(packet.pts(), packet.rotation());
+        let ret = unsafe { sys::avcodec_send_packet(self.ctx, pkt) };
+        // Regardless of errors we are done with this packet, parts of the packet might have been
+        // retained in the decoder.
+        unsafe {
+            sys::av_packet_free(&mut pkt);
+        }
+        if ret < 0 {
+            return Err(Error::DecodePacketFailed(ret, err_code_to_string(ret)));
+        }
+
+        Ok(DecoderIterator {
+            dec: self,
+            ended: false,
+        })
+    }
+
+    pub fn timebase(&self) -> u32 {
+        unsafe {
+            let num = (*self.ctx).time_base.num;
+            let den = (*self.ctx).time_base.den;
+            // Assumption here that numerator is 1
+            assert_eq!(num, 1);
+            assert!(den.is_positive());
+            den as u32
+        }
+    }
+}
+
+struct DecoderIterator<'a> {
+    dec: &'a mut Decoder,
+    ended: bool,
+}
+
+impl<'a> Iterator for DecoderIterator<'a> {
+    type Item = Result<DecodedFrame, Error>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.ended {
+            return None;
+        }
+
+        let frame = DecodedFrame::default();
+
+        let ret = unsafe { sys::avcodec_receive_frame(self.dec.ctx, frame.0) };
+        if ret == sys::AVErrorEAgain || ret == sys::AVErrorEof {
+            self.ended = true;
+            return None;
+        } else if ret < 0 {
+            self.ended = true;
+            return Some(Err(Error::ReceiveFrameFailed(ret, err_code_to_string(ret))));
+        }
+        unsafe {
+            // This is a pointer but it's entirely opaque to libavcodec so we can use it to store
+            // some arbitrary pointer sized data.
+            (*frame.0).opaque = self.dec.pts_map.get(frame.pts()).unwrap_or(0) as *mut c_void;
+        };
+
+        Some(Ok(frame))
+    }
+}
+
+impl PacketData {
+    pub fn new(mut data: Vec<u8>) -> Self {
+        data.extend_from_slice(&[0; sys::AV_INPUT_BUFFER_PADDING_SIZE as usize]);
+
+        Self {
+            inner: data.into_boxed_slice(),
+        }
+    }
+}
+
+impl From<&[u8]> for PacketData {
+    fn from(value: &[u8]) -> Self {
+        let new_size = value.len() + sys::AV_INPUT_BUFFER_PADDING_SIZE as usize;
+        let mut vec = Vec::with_capacity(new_size);
+        vec.extend_from_slice(value);
+
+        Self::new(vec)
+    }
+}
+
+impl DecodedFrame {
+    /// Return the inner ptr for the frame.
+    ///
+    /// ## Safety
+    /// This pointer **MUST** eventually be passed back to [`Frame::from_raw`] to avoid leaking
+    /// memory.
+    pub unsafe fn into_raw(mut self) -> *mut c_void {
+        let ptr = self.0;
+        self.0 = ptr::null_mut();
+
+        ptr as *mut c_void
+    }
+
+    /// Create a [`Frame`] from a raw pointer obtained from [`Frame::into_raw`].
+    ///
+    /// ## Safety
+    /// `ptr` **MUST** have been originally obtained from [`Frame::into_raw`]
+    pub unsafe fn from_raw(ptr: *mut c_void) -> Self {
+        assert!(!ptr.is_null());
+
+        Self(ptr.cast())
+    }
+
+    /// The presentation timestamp for this frame.
+    pub fn pts(&self) -> i64 {
+        // SAFETY: The pointer is valid while self is alive.
+        unsafe { (*self.0).pts }
+    }
+
+    /// The rotation of the frame.
+    pub fn rotation(&self) -> usize {
+        // SAFETY: The pointer is valid while self is alive.
+        unsafe { (*self.0).opaque as usize }
+    }
+}
+
+impl FrameRef for DecodedFrame {
+    fn width(&self) -> usize {
+        // SAFETY: The pointer is valid while self is alive.
+        unsafe { (*self.0).width as usize }
+    }
+
+    fn height(&self) -> usize {
+        // SAFETY: The pointer is valid while self is alive.
+        unsafe { (*self.0).height as usize }
+    }
+
+    fn plane_count(&self) -> usize {
+        // SAFETY: The pointer is valid while self is alive.
+        unsafe {
+            assert_eq!(
+                (*self.0).format,
+                PixelFormat::AV_PIX_FMT_YUV420P as i32,
+                "Only YUV420P is supported"
+            );
+            let fmt = sys::av_pix_fmt_desc_get(PixelFormat::AV_PIX_FMT_YUV420P);
+            (*fmt).nb_components.into()
+        }
+    }
+
+    fn get_plane(&self, i: usize) -> &[u8] {
+        // SAFETY:
+        // * The pointer is valid while self is alive.
+        // * The value calculated for `len` is correct
+        unsafe {
+            assert_eq!(
+                (*self.0).format,
+                PixelFormat::AV_PIX_FMT_YUV420P as i32,
+                "Only YUV420P is supported"
+            );
+            let ptr: *mut u8 = (*self.0).data[i];
+
+            let len = sys::av_image_get_linesize(
+                PixelFormat::AV_PIX_FMT_YUV420P,
+                self.width() as i32,
+                i as i32,
+            ) * self.height() as i32;
+
+            std::slice::from_raw_parts(ptr, len as usize)
+        }
+    }
+
+    fn get_stride(&self, i: usize) -> usize {
+        assert!(i < sys::AV_NUM_DATA_POINTERS as usize);
+
+        // SAFETY: The pointer is valid while self is alive.
+        unsafe { (*self.0).linesize[i] as usize }
+    }
+}
+
+impl PtsMap {
+    fn set(&mut self, pts: i64, value: usize) {
+        self.map[self.cur] = (pts, value);
+        self.cur = (self.cur + 1) % self.map.len();
+    }
+
+    fn get(&self, pts: i64) -> Option<usize> {
+        self.map
+            .iter()
+            .find(|(p, _)| *p == pts)
+            .copied()
+            .map(|(_, v)| v)
+    }
+}
+
+impl Default for DecodedFrame {
+    fn default() -> Self {
+        let ptr = unsafe { sys::av_frame_alloc() };
+        assert!(!ptr.is_null());
+
+        Self(ptr)
+    }
+}
+
+impl Drop for DecodedFrame {
+    fn drop(&mut self) {
+        if self.0.is_null() {
+            return;
+        }
+
+        unsafe {
+            sys::av_frame_free(&mut self.0);
+        }
+    }
+}
+
+impl Default for PtsMap {
+    fn default() -> Self {
+        Self {
+            map: [(-1, 0); 16],
+            cur: 0,
+        }
+    }
+}
+
+extern "C" fn free_boxed_slice(opaque: *mut c_void, data: *mut u8) {
+    let len = opaque as usize;
+    let ptr = std::ptr::slice_from_raw_parts_mut(data, len);
+
+    // SAFETY: The pointer was originally created from a Box<[u8]> and the length was that from
+    // said boxed slice.
+    let _ = unsafe { Box::from_raw(ptr) };
+}

src/error.rs

@@ -8,6 +8,9 @@ pub enum Error {
     #[error("Codec is not an encoder: {0}")]
     CodecIsNotEncoder(&'static str),
 
+    #[error("Codec is not a decoder: {0}")]
+    CodecIsNotDecoder(&'static str),
+
     #[error("Failed to avcodec_alloc_context3")]
     CreateContextFailed,
 
@@ -20,6 +23,15 @@ pub enum Error {
     #[error("Failed to encode frame: {0} {1}")]
     EncodeFrameFailed(i32, String),
 
+    #[error("Failed to decode packet: {0} {1}")]
+    DecodePacketFailed(i32, String),
+
     #[error("Failed to receive encoded packet: {0} {1}")]
     ReceivePacketFailed(i32, String),
+
+    #[error("Failed to receive decoded frame: {0} {1}")]
+    ReceiveFrameFailed(i32, String),
+
+    #[error("Failed to allocate memory: {0}")]
+    AlllocateFailed(&'static str),
 }

src/lib.rs

@@ -9,6 +9,8 @@ use sys::AVPixelFormat as PixelFormat;
 
 mod encoder;
 pub use encoder::{Encoder, EncoderConfig, EncoderProfile, PacketProducer};
+mod decoder;
+pub use decoder::{DecodedFrame, Decoder, DecoderPacket, PacketData};
 mod error;
 pub use error::Error;