ffmpeg/
avcodec.rs

1// Copyright 2022 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! This module implements a lightweight and safe decoder interface over `libavcodec`. It is
6//! designed to concentrate all calls to unsafe methods in one place, while providing the same
7//! low-level access as the libavcodec functions do.
8
9use std::ffi::CStr;
10use std::fmt::Debug;
11use std::fmt::Display;
12use std::marker::PhantomData;
13use std::mem::ManuallyDrop;
14use std::ops::Deref;
15
16use libc::c_char;
17use libc::c_int;
18use libc::c_void;
19use thiserror::Error as ThisError;
20
21use super::*;
22use crate::ffi::AVPictureType;
23
24/// An error returned by a low-level libavcodec function.
25#[derive(Debug, ThisError)]
26pub struct AvError(pub libc::c_int);
27
28impl AvError {
29    pub fn result(ret: c_int) -> Result<(), Self> {
30        if ret >= 0 {
31            Ok(())
32        } else {
33            Err(AvError(ret))
34        }
35    }
36}
37
38impl Display for AvError {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        let mut buffer = [0u8; 255];
41        let ret =
42            // SAFETY:
43            // Safe because we are passing valid bounds for the buffer.
44            unsafe { ffi::av_strerror(self.0, buffer.as_mut_ptr() as *mut c_char, buffer.len()) };
45        match ret {
46            ret if ret >= 0 => {
47                let end_of_string = buffer.iter().position(|i| *i == 0).unwrap_or(buffer.len());
48                let error_string = std::string::String::from_utf8_lossy(&buffer[..end_of_string]);
49                f.write_str(&error_string)
50            }
51            _ => f.write_fmt(format_args!("Unknown avcodec error {}", self.0)),
52        }
53    }
54}
55
56/// Lightweight abstraction over libavcodec's `AVCodec` struct, allowing the query the capabilities
57/// of supported codecs and opening a session to work with them.
58///
59/// `AVCodec` instances in libavcodec are all static, hence we can safely use a static reference
60/// lifetime here.
61pub struct AvCodec(&'static ffi::AVCodec);
62
63#[derive(Debug, ThisError)]
64pub enum AvCodecOpenError {
65    #[error("failed to allocate AVContext object")]
66    ContextAllocation,
67    #[error("failed to open AVContext object")]
68    ContextOpen,
69    #[error("ContextBuilder variant does not match codec type")]
70    UnexpectedCodecType,
71}
72
73/// Dimensions of a frame, used in AvCodecContext and AvFrame.
74#[derive(Copy, Clone, Debug, PartialEq, Eq)]
75pub struct Dimensions {
76    pub width: u32,
77    pub height: u32,
78}
79
80impl AvCodec {
81    /// Returns whether the codec is a decoder.
82    pub fn is_decoder(&self) -> bool {
83        // SAFETY:
84        // Safe because `av_codec_is_decoder` is called on a valid static `AVCodec` reference.
85        (unsafe { ffi::av_codec_is_decoder(self.0) } != 0)
86    }
87
88    /// Returns whether the codec is an encoder.
89    pub fn is_encoder(&self) -> bool {
90        // SAFETY:
91        // Safe because `av_codec_is_encoder` is called on a valid static `AVCodec` reference.
92        (unsafe { ffi::av_codec_is_encoder(self.0) } != 0)
93    }
94
95    /// Returns the name of the codec.
96    pub fn name(&self) -> &'static str {
97        const INVALID_CODEC_STR: &str = "invalid codec";
98
99        // SAFETY:
100        // Safe because `CStr::from_ptr` is called on a valid zero-terminated C string.
101        unsafe { CStr::from_ptr(self.0.name).to_str() }.unwrap_or(INVALID_CODEC_STR)
102    }
103
104    /// Returns the capabilities of the codec, as a mask of AV_CODEC_CAP_* bits.
105    pub fn capabilities(&self) -> u32 {
106        self.0.capabilities as u32
107    }
108
109    /// Returns an iterator over the profiles supported by this codec.
110    pub fn profile_iter(&self) -> AvProfileIterator {
111        AvProfileIterator(self.0.profiles)
112    }
113
114    /// Returns an iterator over the pixel formats supported by this codec.
115    ///
116    /// For a decoder, the returned array will likely be empty. This means that ffmpeg's native
117    /// pixel format (YUV420) will be used.
118    pub fn pixel_format_iter(&self) -> AvPixelFormatIterator {
119        AvPixelFormatIterator(self.0.pix_fmts)
120    }
121
122    /// Get a builder for a encoder [`AvCodecContext`] using this codec.
123    pub fn build_encoder(&self) -> Result<EncoderContextBuilder, AvCodecOpenError> {
124        if !self.is_encoder() {
125            return Err(AvCodecOpenError::UnexpectedCodecType);
126        }
127
128        Ok(EncoderContextBuilder {
129            codec: self.0,
130            context: self.alloc_context()?,
131        })
132    }
133
134    /// Get a builder for a decoder [`AvCodecContext`] using this codec.
135    pub fn build_decoder(&self) -> Result<DecoderContextBuilder, AvCodecOpenError> {
136        if !self.is_decoder() {
137            return Err(AvCodecOpenError::UnexpectedCodecType);
138        }
139
140        Ok(DecoderContextBuilder {
141            codec: self.0,
142            context: self.alloc_context()?,
143        })
144    }
145
146    /// Internal helper for `build_decoder` to allocate an [`AvCodecContext`]. This needs to be
147    /// paired with a later call to [`AvCodecContext::init`].
148    fn alloc_context(&self) -> Result<AvCodecContext, AvCodecOpenError> {
149        // TODO(b:315859322): add safety doc string
150        #[allow(clippy::undocumented_unsafe_blocks)]
151        let context = unsafe { ffi::avcodec_alloc_context3(self.0).as_mut() }
152            .ok_or(AvCodecOpenError::ContextAllocation)?;
153
154        Ok(AvCodecContext(context))
155    }
156}
157
158/// A builder to create a [`AvCodecContext`] suitable for decoding.
159// This struct wraps an AvCodecContext directly, but the only way it can be taken out is to call
160// `build()`, which finalizes the context and prevent further modification to the callback, etc.
161pub struct DecoderContextBuilder {
162    codec: *const ffi::AVCodec,
163    context: AvCodecContext,
164}
165
166impl DecoderContextBuilder {
167    /// Set a custom callback that provides output buffers.
168    ///
169    /// `get_buffer2` is a function that decides which buffer is used to render a frame (see
170    /// libavcodec's documentation for `get_buffer2` for more details). If provided, this function
171    /// must be thread-safe.
172    /// `opaque` is a pointer that will be passed as first argument to `get_buffer2` when it is
173    /// called.
174    pub fn set_get_buffer_2(
175        &mut self,
176        get_buffer2: unsafe extern "C" fn(*mut ffi::AVCodecContext, *mut ffi::AVFrame, i32) -> i32,
177        opaque: *mut libc::c_void,
178    ) {
179        // SAFETY:
180        // Safe because self.context.0 is a pointer to a live AVCodecContext allocation.
181        let context = unsafe { &mut *(self.context.0) };
182        context.get_buffer2 = Some(get_buffer2);
183        context.opaque = opaque;
184    }
185
186    /// Build a decoder AvCodecContext from the configured options.
187    pub fn build(mut self) -> Result<AvCodecContext, AvCodecOpenError> {
188        self.context.init(self.codec)?;
189        Ok(self.context)
190    }
191}
192
193/// A builder to create a [`AvCodecContext`] suitable for encoding.
194// This struct wraps an AvCodecContext directly, but the only way it can be taken out is to call
195// `build()`, which finalizes the context and prevent further modification to the callback, etc.
196pub struct EncoderContextBuilder {
197    codec: *const ffi::AVCodec,
198    context: AvCodecContext,
199}
200
201impl EncoderContextBuilder {
202    /// Set the width of input frames for this encoding context.
203    pub fn set_dimensions(&mut self, dimensions: Dimensions) {
204        // TODO(b:315859322): add safety doc string
205        #[allow(clippy::undocumented_unsafe_blocks)]
206        let context = unsafe { &mut *(self.context.0) };
207        context.width = dimensions.width as _;
208        context.height = dimensions.height as _;
209    }
210
211    /// Set the time base for this encoding context.
212    pub fn set_time_base(&mut self, time_base: ffi::AVRational) {
213        // TODO(b:315859322): add safety doc string
214        #[allow(clippy::undocumented_unsafe_blocks)]
215        let context = unsafe { &mut *(self.context.0) };
216        context.time_base = time_base;
217    }
218
219    /// Set the input pixel format for this encoding context.
220    pub fn set_pix_fmt(&mut self, fmt: AvPixelFormat) {
221        // TODO(b:315859322): add safety doc string
222        #[allow(clippy::undocumented_unsafe_blocks)]
223        let context = unsafe { &mut *(self.context.0) };
224        context.pix_fmt = fmt.pix_fmt();
225    }
226
227    /// Build a encoder AvCodecContext from the configured options.
228    pub fn build(mut self) -> Result<AvCodecContext, AvCodecOpenError> {
229        self.context.init(self.codec)?;
230        Ok(self.context)
231    }
232}
233
234impl Default for AvCodecIterator {
235    fn default() -> Self {
236        Self::new()
237    }
238}
239
240/// Lightweight abstraction over libavcodec's `av_codec_iterate` function that can be used to
241/// enumerate all the supported codecs.
242pub struct AvCodecIterator(*mut libc::c_void);
243
244impl AvCodecIterator {
245    pub fn new() -> Self {
246        Self(std::ptr::null_mut())
247    }
248}
249
250impl Iterator for AvCodecIterator {
251    type Item = AvCodec;
252
253    fn next(&mut self) -> Option<Self::Item> {
254        // SAFETY:
255        // Safe because our pointer was initialized to `NULL` and we only use it with
256        // `av_codec_iterate`, which will update it to a valid value.
257        unsafe { ffi::av_codec_iterate(&mut self.0 as *mut *mut libc::c_void).as_ref() }
258            .map(AvCodec)
259    }
260}
261
262/// Simple wrapper over `AVProfile` that provides helpful methods.
263pub struct AvProfile(&'static ffi::AVProfile);
264
265impl AvProfile {
266    /// Return the profile id, which can be matched against FF_PROFILE_*.
267    pub fn profile(&self) -> u32 {
268        self.0.profile as u32
269    }
270
271    /// Return the name of this profile.
272    pub fn name(&self) -> &'static str {
273        const INVALID_PROFILE_STR: &str = "invalid profile";
274
275        // SAFETY:
276        // Safe because `CStr::from_ptr` is called on a valid zero-terminated C string.
277        unsafe { CStr::from_ptr(self.0.name).to_str() }.unwrap_or(INVALID_PROFILE_STR)
278    }
279}
280
281impl Display for AvProfile {
282    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
283        f.write_str(self.name())
284    }
285}
286
287impl Debug for AvProfile {
288    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
289        Display::fmt(self, f)
290    }
291}
292
293/// Lightweight abstraction over the array of supported profiles for a given codec.
294pub struct AvProfileIterator(*const ffi::AVProfile);
295
296impl Iterator for AvProfileIterator {
297    type Item = AvProfile;
298
299    fn next(&mut self) -> Option<Self::Item> {
300        // SAFETY:
301        // Safe because the contract of `new` stipulates we have received a valid `AVCodec`
302        // reference, thus the `profiles` pointer must either be NULL or point to a valid array
303        // or `VAProfile`s.
304        match unsafe { self.0.as_ref() } {
305            None => None,
306            Some(profile) => {
307                match profile.profile {
308                    ffi::FF_PROFILE_UNKNOWN => None,
309                    _ => {
310                        // SAFETY:
311                        // Safe because we have been initialized to a static, valid profiles array
312                        // which is terminated by FF_PROFILE_UNKNOWN.
313                        self.0 = unsafe { self.0.offset(1) };
314                        Some(AvProfile(profile))
315                    }
316                }
317            }
318        }
319    }
320}
321
322#[derive(Clone, Copy)]
323/// Simple wrapper over `AVPixelFormat` that provides helpful methods.
324pub struct AvPixelFormat(ffi::AVPixelFormat);
325
326impl AvPixelFormat {
327    /// Return the name of this pixel format.
328    pub fn name(&self) -> &'static str {
329        const INVALID_FORMAT_STR: &str = "invalid pixel format";
330
331        // SAFETY:
332        // Safe because `av_get_pix_fmt_name` returns either NULL or a valid C string.
333        let pix_fmt_name = unsafe { ffi::av_get_pix_fmt_name(self.0) };
334        // SAFETY:
335        // Safe because `pix_fmt_name` is a valid pointer to a C string.
336        match unsafe {
337            pix_fmt_name
338                .as_ref()
339                .and_then(|s| CStr::from_ptr(s).to_str().ok())
340        } {
341            None => INVALID_FORMAT_STR,
342            Some(string) => string,
343        }
344    }
345
346    /// Return the avcodec profile id, which can be matched against AV_PIX_FMT_*.
347    ///
348    /// Note that this is **not** the same as a fourcc.
349    pub fn pix_fmt(&self) -> ffi::AVPixelFormat {
350        self.0
351    }
352
353    /// Return the fourcc of the pixel format, or a series of zeros if its fourcc is unknown.
354    pub fn fourcc(&self) -> [u8; 4] {
355        // SAFETY:
356        // Safe because `avcodec_pix_fmt_to_codec_tag` does not take any pointer as input and
357        // handles any value passed as argument.
358        unsafe { ffi::avcodec_pix_fmt_to_codec_tag(self.0) }.to_le_bytes()
359    }
360
361    /// Given the width and plane index, returns the line size (data pointer increment per row) in
362    /// bytes.
363    pub fn line_size(&self, width: u32, plane: usize) -> Result<usize, AvError> {
364        av_image_line_size(*self, width, plane)
365    }
366
367    /// Given an iterator of line sizes and height, return the size required for each plane's buffer
368    /// in bytes.
369    pub fn plane_sizes<I: IntoIterator<Item = u32>>(
370        &self,
371        linesizes: I,
372        height: u32,
373    ) -> Result<Vec<usize>, AvError> {
374        av_image_plane_sizes(*self, linesizes, height)
375    }
376}
377
378#[derive(Debug)]
379pub struct FromAVPixelFormatError(());
380
381impl TryFrom<ffi::AVPixelFormat> for AvPixelFormat {
382    type Error = FromAVPixelFormatError;
383
384    fn try_from(value: ffi::AVPixelFormat) -> Result<Self, Self::Error> {
385        if value > ffi::AVPixelFormat_AV_PIX_FMT_NONE && value < ffi::AVPixelFormat_AV_PIX_FMT_NB {
386            Ok(AvPixelFormat(value))
387        } else {
388            Err(FromAVPixelFormatError(()))
389        }
390    }
391}
392
393impl Display for AvPixelFormat {
394    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
395        f.write_str(self.name())
396    }
397}
398
399impl Debug for AvPixelFormat {
400    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
401        let fourcc = self.fourcc();
402        f.write_fmt(format_args!(
403            "{}{}{}{}",
404            fourcc[0] as char, fourcc[1] as char, fourcc[2] as char, fourcc[3] as char
405        ))
406    }
407}
408
409/// Lightweight abstraction over the array of supported pixel formats for a given codec.
410pub struct AvPixelFormatIterator(*const ffi::AVPixelFormat);
411
412impl Iterator for AvPixelFormatIterator {
413    type Item = AvPixelFormat;
414
415    fn next(&mut self) -> Option<Self::Item> {
416        // SAFETY:
417        // Safe because the contract of `AvCodec::new` and `AvCodec::pixel_format_iter` guarantees
418        // that we have been built from a valid `AVCodec` reference, which `pix_fmts` pointer
419        // must either be NULL or point to a valid array or `VAPixelFormat`s.
420        match unsafe { self.0.as_ref() } {
421            None => None,
422            Some(&pixfmt) => {
423                match pixfmt {
424                    // Array of pixel formats is terminated by AV_PIX_FMT_NONE.
425                    ffi::AVPixelFormat_AV_PIX_FMT_NONE => None,
426                    _ => {
427                        // SAFETY:
428                        // Safe because we have been initialized to a static, valid profiles array
429                        // which is terminated by AV_PIX_FMT_NONE.
430                        self.0 = unsafe { self.0.offset(1) };
431                        Some(AvPixelFormat(pixfmt))
432                    }
433                }
434            }
435        }
436    }
437}
438
439/// A codec context from which decoding can be performed.
440pub struct AvCodecContext(*mut ffi::AVCodecContext);
441
442impl Drop for AvCodecContext {
443    fn drop(&mut self) {
444        // SAFETY:
445        // Safe because our context member is properly allocated and owned by us.
446        // Note: `avcodec_open2` might not have been called in case we're wrapped by a
447        //       `DecoderContextBuilder` but avcodec_free_context works on both opened and closed
448        //       contexts.
449        unsafe { ffi::avcodec_free_context(&mut self.0) };
450    }
451}
452
453impl AsRef<ffi::AVCodecContext> for AvCodecContext {
454    fn as_ref(&self) -> &ffi::AVCodecContext {
455        // SAFETY:
456        // Safe because our context member is properly initialized and fully owned by us.
457        unsafe { &*self.0 }
458    }
459}
460
461pub enum TryReceiveResult {
462    Received,
463    TryAgain,
464    FlushCompleted,
465}
466
467impl AvCodecContext {
468    /// Internal helper for [`DecoderContextBuilder`] to initialize the context.
469    fn init(&mut self, codec: *const ffi::AVCodec) -> Result<(), AvCodecOpenError> {
470        // SAFETY:
471        // Safe because `codec` is a valid static AVCodec reference, and `self.0` is a valid
472        // AVCodecContext allocation.
473        if unsafe { ffi::avcodec_open2(self.0, codec, std::ptr::null_mut()) } < 0 {
474            return Err(AvCodecOpenError::ContextOpen);
475        }
476
477        Ok(())
478    }
479
480    /// Send a packet to be decoded by the codec.
481    ///
482    /// Returns `true` if the packet has been accepted and will be decoded, `false` if the codec can
483    /// not accept frames at the moment - in this case `try_receive_frame` must be called before
484    /// the packet can be submitted again.
485    ///
486    /// Error codes are the same as those returned by `avcodec_send_packet` with the exception of
487    /// EAGAIN which is converted into `Ok(false)` as it is not actually an error.
488    pub fn try_send_packet(&mut self, packet: &AvPacket) -> Result<bool, AvError> {
489        // SAFETY:
490        // Safe because the context is valid through the life of this object, and `packet`'s
491        // lifetime properties ensures its memory area is readable.
492        match unsafe { ffi::avcodec_send_packet(self.0, &packet.packet) } {
493            AVERROR_EAGAIN => Ok(false),
494            ret if ret >= 0 => Ok(true),
495            err => Err(AvError(err)),
496        }
497    }
498
499    /// Attempt to write a decoded frame in `frame` if the codec has enough data to do so.
500    ///
501    /// Returned `Received` if `frame` has been filled with the next decoded frame, `TryAgain` if
502    /// no frame could be returned at that time (in which case `try_send_packet` should be called to
503    /// submit more input to decode), or `FlushCompleted` to signal that a previous flush triggered
504    /// by calling the `flush` method has completed.
505    ///
506    /// Error codes are the same as those returned by `avcodec_receive_frame` with the exception of
507    /// EAGAIN and EOF which are handled as `TryAgain` and `FlushCompleted` respectively.
508    pub fn try_receive_frame(&mut self, frame: &mut AvFrame) -> Result<TryReceiveResult, AvError> {
509        // SAFETY:
510        // Safe because the context is valid through the life of this object, and `avframe` is
511        // guaranteed to contain a properly initialized frame.
512        match unsafe { ffi::avcodec_receive_frame(self.0, frame.0) } {
513            AVERROR_EAGAIN => Ok(TryReceiveResult::TryAgain),
514            AVERROR_EOF => Ok(TryReceiveResult::FlushCompleted),
515            ret if ret >= 0 => Ok(TryReceiveResult::Received),
516            err => Err(AvError(err)),
517        }
518    }
519
520    /// Send a frame to be encoded by the codec.
521    ///
522    /// Returns `true` if the frame has been accepted and will be encoded, `false` if the codec can
523    /// not accept input at the moment - in this case `try_receive_frame` must be called before
524    /// the frame can be submitted again.
525    ///
526    /// Error codes are the same as those returned by `avcodec_send_frame` with the exception of
527    /// EAGAIN which is converted into `Ok(false)` as it is not actually an error.
528    pub fn try_send_frame(&mut self, frame: &AvFrame) -> Result<bool, AvError> {
529        // TODO(b:315859322): add safety doc string
530        #[allow(clippy::undocumented_unsafe_blocks)]
531        match unsafe { ffi::avcodec_send_frame(self.0, frame.0 as *const _) } {
532            AVERROR_EAGAIN => Ok(false),
533            ret if ret >= 0 => Ok(true),
534            err => Err(AvError(err)),
535        }
536    }
537
538    /// Attempt to write an encoded frame in `packet` if the codec has enough data to do so.
539    ///
540    /// Returned `Received` if `packet` has been filled with encoded data, `TryAgain` if
541    /// no packet could be returned at that time (in which case `try_send_frame` should be called to
542    /// submit more input to decode), or `FlushCompleted` to signal that a previous flush triggered
543    /// by calling the `flush` method has completed.
544    ///
545    /// Error codes are the same as those returned by `avcodec_receive_packet` with the exception of
546    /// EAGAIN and EOF which are handled as `TryAgain` and `FlushCompleted` respectively.
547    pub fn try_receive_packet(
548        &mut self,
549        packet: &mut AvPacket,
550    ) -> Result<TryReceiveResult, AvError> {
551        // SAFETY:
552        // Safe because the context is valid through the life of this object, and `avframe` is
553        // guaranteed to contain a properly initialized frame.
554        match unsafe { ffi::avcodec_receive_packet(self.0, &mut packet.packet) } {
555            AVERROR_EAGAIN => Ok(TryReceiveResult::TryAgain),
556            AVERROR_EOF => Ok(TryReceiveResult::FlushCompleted),
557            ret if ret >= 0 => Ok(TryReceiveResult::Received),
558            err => Err(AvError(err)),
559        }
560    }
561
562    /// Reset the internal codec state/flush internal buffers.
563    /// Should be called e.g. when seeking or switching to a different stream.
564    pub fn reset(&mut self) {
565        // SAFETY:
566        // Safe because the context is valid through the life of this object.
567        unsafe { ffi::avcodec_flush_buffers(self.0) }
568    }
569
570    /// Ask the context to start flushing, i.e. to process all pending input packets and produce
571    /// frames for them.
572    ///
573    /// The flush process is complete when `try_receive_frame` returns `FlushCompleted`,
574    pub fn flush_decoder(&mut self) -> Result<(), AvError> {
575        // SAFETY:
576        // Safe because the context is valid through the life of this object.
577        AvError::result(unsafe { ffi::avcodec_send_packet(self.0, std::ptr::null()) })
578    }
579
580    /// Ask the context to start flushing, i.e. to process all pending input frames and produce
581    /// packets for them.
582    ///
583    /// The flush process is complete when `try_receive_packet` returns `FlushCompleted`,
584    pub fn flush_encoder(&mut self) -> Result<(), AvError> {
585        // SAFETY:
586        // Safe because the context is valid through the life of this object.
587        AvError::result(unsafe { ffi::avcodec_send_frame(self.0, std::ptr::null()) })
588    }
589
590    /// Set the time base for this context.
591    pub fn set_time_base(&mut self, time_base: AVRational) {
592        // TODO(b:315859322): add safety doc string
593        #[allow(clippy::undocumented_unsafe_blocks)]
594        let context = unsafe { &mut *(self.0) };
595        context.time_base = time_base;
596    }
597
598    /// Set the bit rate for this context.
599    pub fn set_bit_rate(&mut self, bit_rate: u64) {
600        // TODO(b:315859322): add safety doc string
601        #[allow(clippy::undocumented_unsafe_blocks)]
602        let context = unsafe { &mut *(self.0) };
603        context.bit_rate = bit_rate as _;
604    }
605
606    /// Set the max bit rate (rc_max_rate) for this context.
607    pub fn set_max_bit_rate(&mut self, bit_rate: u64) {
608        // TODO(b:315859322): add safety doc string
609        #[allow(clippy::undocumented_unsafe_blocks)]
610        let context = unsafe { &mut *(self.0) };
611        context.rc_max_rate = bit_rate as _;
612    }
613}
614
615/// Trait for types that can be used as data provider for a `AVBuffer`.
616///
617/// `AVBuffer` is an owned buffer type, so all the type needs to do is being able to provide a
618/// stable pointer to its own data as well as its length. Implementors need to be sendable across
619/// threads because avcodec is allowed to use threads in its codec implementations.
620pub trait AvBufferSource: Send {
621    fn as_ptr(&self) -> *const u8;
622    fn as_mut_ptr(&mut self) -> *mut u8 {
623        self.as_ptr() as *mut u8
624    }
625    fn len(&self) -> usize;
626    fn is_empty(&self) -> bool;
627}
628
629/// Wrapper around `AVBuffer` and `AVBufferRef`.
630///
631/// libavcodec can manage its own memory for input and output data. Doing so implies a transparent
632/// copy of user-provided data (packets or frames) from and to this memory, which is wasteful.
633///
634/// This copy can be avoided by explicitly providing our own buffers to libavcodec using
635/// `AVBufferRef`. Doing so means that the lifetime of these buffers becomes managed by avcodec.
636/// This struct helps make this process safe by taking full ownership of an `AvBufferSource` and
637/// dropping it when libavcodec is done with it.
638pub struct AvBuffer(*mut ffi::AVBufferRef);
639
640impl AvBuffer {
641    /// Create a new `AvBuffer` from an `AvBufferSource`.
642    ///
643    /// Ownership of `source` is transferred to libavcodec, which will drop it when the number of
644    /// references to this buffer reaches zero.
645    ///
646    /// Returns `None` if the buffer could not be created due to an error in libavcodec.
647    pub fn new<D: AvBufferSource + 'static>(source: D) -> Option<Self> {
648        // Move storage to the heap so we find it at the same place in `avbuffer_free`
649        let mut storage = Box::new(source);
650
651        extern "C" fn avbuffer_free<D>(opaque: *mut c_void, _data: *mut u8) {
652            // SAFETY:
653            // Safe because `opaque` has been created from `Box::into_raw`. `storage` will be
654            // dropped immediately which will release any resources held by the storage.
655            let _ = unsafe { Box::from_raw(opaque as *mut D) };
656        }
657
658        // SAFETY:
659        // Safe because storage points to valid data throughout the lifetime of AVBuffer and we are
660        // checking the return value against NULL, which signals an error.
661        Some(Self(unsafe {
662            ffi::av_buffer_create(
663                storage.as_mut_ptr(),
664                storage.len(),
665                Some(avbuffer_free::<D>),
666                Box::into_raw(storage) as *mut c_void,
667                0,
668            )
669            .as_mut()?
670        }))
671    }
672
673    /// Return a slice to the data contained in this buffer.
674    pub fn as_mut_slice(&mut self) -> &mut [u8] {
675        // SAFETY:
676        // Safe because the data has been initialized from valid storage in the constructor.
677        unsafe { std::slice::from_raw_parts_mut((*self.0).data, (*self.0).size) }
678    }
679
680    /// Consumes the `AVBuffer`, returning a `AVBufferRef` that can be used in `AVFrame`, `AVPacket`
681    /// and others.
682    ///
683    /// After calling, the caller is responsible for unref-ing the returned AVBufferRef, either
684    /// directly or through one of the automatic management facilities in `AVFrame`, `AVPacket` or
685    /// others.
686    pub fn into_raw(self) -> *mut ffi::AVBufferRef {
687        ManuallyDrop::new(self).0
688    }
689}
690
691impl Drop for AvBuffer {
692    fn drop(&mut self) {
693        // SAFETY:
694        // Safe because `self.0` is a valid pointer to an AVBufferRef.
695        unsafe { ffi::av_buffer_unref(&mut self.0) };
696    }
697}
698
699/// An encoded input packet that can be submitted to `AvCodecContext::try_send_packet`.
700pub struct AvPacket<'a> {
701    packet: ffi::AVPacket,
702    _buffer_data: PhantomData<&'a ()>,
703}
704
705impl Drop for AvPacket<'_> {
706    fn drop(&mut self) {
707        // SAFETY:
708        // Safe because `self.packet` is a valid `AVPacket` instance.
709        unsafe {
710            ffi::av_packet_unref(&mut self.packet);
711        }
712    }
713}
714
715impl AsRef<ffi::AVPacket> for AvPacket<'_> {
716    fn as_ref(&self) -> &ffi::AVPacket {
717        &self.packet
718    }
719}
720
721impl<'a> AvPacket<'a> {
722    /// Create an empty AvPacket without buffers.
723    ///
724    /// This packet should be only used with an encoder; in which case the encoder will
725    /// automatically allocate a buffer of appropriate size and store it inside this `AvPacket`.
726    pub fn empty() -> Self {
727        Self {
728            packet: ffi::AVPacket {
729                pts: AV_NOPTS_VALUE as i64,
730                dts: AV_NOPTS_VALUE as i64,
731                pos: -1,
732                // SAFETY:
733                // Safe because all the other elements of this struct can be zeroed.
734                ..unsafe { std::mem::zeroed() }
735            },
736            _buffer_data: PhantomData,
737        }
738    }
739
740    /// Create a new AvPacket that borrows the `input_data`.
741    ///
742    /// The returned `AvPacket` will hold a reference to `input_data`, meaning that libavcodec might
743    /// perform a copy from/to it.
744    pub fn new<T: AvBufferSource>(pts: i64, input_data: &'a mut T) -> Self {
745        Self {
746            packet: ffi::AVPacket {
747                buf: std::ptr::null_mut(),
748                pts,
749                dts: AV_NOPTS_VALUE as i64,
750                data: input_data.as_mut_ptr(),
751                size: input_data.len() as c_int,
752                side_data: std::ptr::null_mut(),
753                pos: -1,
754                // SAFETY:
755                // Safe because all the other elements of this struct can be zeroed.
756                ..unsafe { std::mem::zeroed() }
757            },
758            _buffer_data: PhantomData,
759        }
760    }
761
762    /// Create a new AvPacket that owns the `av_buffer`.
763    ///
764    /// The returned `AvPacket` will have a `'static` lifetime and will keep `input_data` alive for
765    /// as long as libavcodec needs it.
766    pub fn new_owned(pts: i64, mut av_buffer: AvBuffer) -> Self {
767        let data_slice = av_buffer.as_mut_slice();
768        let data = data_slice.as_mut_ptr();
769        let size = data_slice.len() as i32;
770
771        Self {
772            packet: ffi::AVPacket {
773                buf: av_buffer.into_raw(),
774                pts,
775                dts: AV_NOPTS_VALUE as i64,
776                data,
777                size,
778                side_data: std::ptr::null_mut(),
779                pos: -1,
780                // SAFETY:
781                // Safe because all the other elements of this struct can be zeroed.
782                ..unsafe { std::mem::zeroed() }
783            },
784            _buffer_data: PhantomData,
785        }
786    }
787}
788
789/// An owned AVFrame, i.e. one decoded frame from libavcodec that can be converted into a
790/// destination buffer.
791pub struct AvFrame(*mut ffi::AVFrame);
792
793/// A builder for AVFrame that allows specifying buffers and image metadata.
794pub struct AvFrameBuilder(AvFrame);
795
796/// A descriptor describing a subslice of `buffers` in [`AvFrameBuilder::build_owned`] that
797/// represents a plane's image data.
798pub struct PlaneDescriptor {
799    /// The index within `buffers`.
800    pub buffer_index: usize,
801    /// The offset from the start of `buffers[buffer_index]`.
802    pub offset: usize,
803    /// The increment of data pointer in bytes per row of the plane.
804    pub stride: usize,
805}
806
807#[derive(Debug, ThisError)]
808pub enum AvFrameError {
809    #[error("failed to allocate AVFrame object")]
810    FrameAllocationFailed,
811    #[error("dimension is negative or too large")]
812    DimensionOverflow,
813    #[error("a row does not fit in the specified stride")]
814    InvalidStride,
815    #[error("buffer index out of range")]
816    BufferOutOfRange,
817    #[error("specified dimensions overflow the buffer size")]
818    BufferTooSmall,
819    #[error("plane reference to buffer alias each other")]
820    BufferAlias,
821    #[error("error while calling libavcodec")]
822    AvError(#[from] AvError),
823}
824
825impl AvFrame {
826    /// Create a new AvFrame. The frame's parameters and backing memory will be assigned when it is
827    /// decoded into.
828    pub fn new() -> Result<Self, AvFrameError> {
829        Ok(Self(
830            // SAFETY:
831            // Safe because `av_frame_alloc` does not take any input.
832            unsafe { ffi::av_frame_alloc().as_mut() }.ok_or(AvFrameError::FrameAllocationFailed)?,
833        ))
834    }
835
836    /// Create a new AvFrame builder that allows setting the frame's parameters and backing memory
837    /// through its methods.
838    pub fn builder() -> Result<AvFrameBuilder, AvFrameError> {
839        AvFrame::new().map(AvFrameBuilder)
840    }
841
842    /// Return the frame's width and height.
843    pub fn dimensions(&self) -> Dimensions {
844        Dimensions {
845            width: self.as_ref().width as _,
846            height: self.as_ref().height as _,
847        }
848    }
849
850    /// Return the frame's pixel format.
851    pub fn format(&self) -> AvPixelFormat {
852        AvPixelFormat(self.as_ref().format)
853    }
854
855    /// Set the picture type (I-frame, P-frame etc.) on this frame.
856    pub fn set_pict_type(&mut self, ty: AVPictureType) {
857        // SAFETY:
858        // Safe because self.0 is a valid AVFrame reference.
859        unsafe {
860            (*self.0).pict_type = ty;
861        }
862    }
863
864    /// Set the presentation timestamp (PTS) of this frame.
865    pub fn set_pts(&mut self, ts: i64) {
866        // SAFETY:
867        // Safe because self.0 is a valid AVFrame reference.
868        unsafe {
869            (*self.0).pts = ts;
870        }
871    }
872
873    /// Query if this AvFrame is writable, i.e. it is refcounted and the refcounts are 1.
874    pub fn is_writable(&self) -> bool {
875        // SAFETY:
876        // Safe because self.0 is a valid AVFrame reference.
877        unsafe { ffi::av_frame_is_writable(self.0) != 0 }
878    }
879
880    /// If the frame is not writable already (see [`is_writable`]), make a copy of its buffer to
881    /// make it writable.
882    ///
883    /// [`is_writable`]: AvFrame::is_writable
884    pub fn make_writable(&mut self) -> Result<(), AvFrameError> {
885        // SAFETY:
886        // Safe because self.0 is a valid AVFrame reference.
887        AvError::result(unsafe { ffi::av_frame_make_writable(self.0) }).map_err(Into::into)
888    }
889}
890
891impl AvFrameBuilder {
892    /// Set the frame's width and height.
893    ///
894    /// The dimensions must not be greater than `i32::MAX`.
895    pub fn set_dimensions(&mut self, dimensions: Dimensions) -> Result<(), AvFrameError> {
896        // SAFETY:
897        // Safe because self.0 is a valid AVFrame instance and width and height are in range.
898        unsafe {
899            (*self.0 .0).width = dimensions
900                .width
901                .try_into()
902                .map_err(|_| AvFrameError::DimensionOverflow)?;
903            (*self.0 .0).height = dimensions
904                .height
905                .try_into()
906                .map_err(|_| AvFrameError::DimensionOverflow)?;
907        }
908        Ok(())
909    }
910
911    /// Set the frame's format.
912    pub fn set_format(&mut self, format: AvPixelFormat) -> Result<(), AvFrameError> {
913        // SAFETY:
914        // Safe because self.0 is a valid AVFrame instance and format is a valid pixel format.
915        unsafe {
916            (*self.0 .0).format = format.pix_fmt();
917        }
918        Ok(())
919    }
920
921    /// Build an AvFrame from iterators of [`AvBuffer`]s and subslice of buffers describing the
922    /// planes.
923    ///
924    /// The frame will own the `buffers`.
925    ///
926    /// This function checks that:
927    /// - Each plane fits inside the bounds of the associated buffer.
928    /// - Different planes do not overlap each other's buffer slice. In this check, all planes are
929    ///   assumed to be potentially mutable, regardless of whether the AvFrame is actually used for
930    ///   read or write access. Aliasing reference to the same buffer will be rejected, since it can
931    ///   potentially allow routines to overwrite each
932    //    other's result.
933    ///   An exception to this is when the same buffer is passed multiple times in `buffers`. In
934    ///   this case, each buffer is treated as a different buffer. Since clones have to be made to
935    ///   be passed multiple times in `buffers`, the frame will not be considered [writable]. Hence
936    ///   aliasing is safe in this case, but the caller is required to explicit opt-in to this
937    ///   read-only handling by passing clones of the buffer into `buffers` and have a different
938    ///   buffer index for each plane combination that could overlap in their range.
939    ///
940    /// [writable]: AvFrame::is_writable
941    pub fn build_owned<
942        BI: IntoIterator<Item = AvBuffer>,
943        PI: IntoIterator<Item = PlaneDescriptor>,
944    >(
945        self,
946        buffers: BI,
947        planes: PI,
948    ) -> Result<AvFrame, AvFrameError> {
949        let mut buffers: Vec<_> = buffers.into_iter().collect();
950        let planes: Vec<_> = planes.into_iter().collect();
951        let format = self.0.format();
952        let plane_sizes = format.plane_sizes(
953            planes.iter().map(|x| x.stride as u32),
954            self.0.dimensions().height,
955        )?;
956        let mut ranges = vec![];
957
958        for (
959            plane,
960            PlaneDescriptor {
961                buffer_index,
962                offset,
963                stride,
964            },
965        ) in planes.into_iter().enumerate()
966        {
967            if buffer_index > buffers.len() {
968                return Err(AvFrameError::BufferOutOfRange);
969            }
970            let end = offset + plane_sizes[plane];
971            if end > buffers[buffer_index].as_mut_slice().len() {
972                return Err(AvFrameError::BufferTooSmall);
973            }
974            if stride < format.line_size(self.0.dimensions().width, plane)? {
975                return Err(AvFrameError::InvalidStride);
976            }
977            // TODO(b:315859322): add safety doc string
978            #[allow(clippy::undocumented_unsafe_blocks)]
979            unsafe {
980                (*self.0 .0).data[plane] =
981                    buffers[buffer_index].as_mut_slice()[offset..].as_mut_ptr();
982                (*self.0 .0).linesize[plane] = stride as c_int;
983            }
984            ranges.push((buffer_index, offset, end));
985        }
986
987        // Check for range overlaps.
988        // See function documentation for the exact rule and reasoning.
989        ranges.sort_unstable();
990        for pair in ranges.windows(2) {
991            // (buffer_index, start, end)
992            let (b0, _s0, e0) = pair[0];
993            let (b1, s1, _e1) = pair[1];
994
995            if b0 != b1 {
996                continue;
997            }
998            // Note that s0 <= s1 always holds, so we only need to check
999            // that the start of the second range is before the end of the first range.
1000            if s1 < e0 {
1001                return Err(AvFrameError::BufferAlias);
1002            }
1003        }
1004
1005        for (i, buf) in buffers.into_iter().enumerate() {
1006            // SAFETY:
1007            // Safe because self.0 is a valid AVFrame instance and buffers contains valid AvBuffers.
1008            unsafe {
1009                (*self.0 .0).buf[i] = buf.into_raw();
1010            }
1011        }
1012        Ok(self.0)
1013    }
1014}
1015
1016impl AsRef<ffi::AVFrame> for AvFrame {
1017    fn as_ref(&self) -> &ffi::AVFrame {
1018        // SAFETY:
1019        // Safe because the AVFrame has been properly initialized during construction.
1020        unsafe { &*self.0 }
1021    }
1022}
1023
1024impl Deref for AvFrame {
1025    type Target = ffi::AVFrame;
1026
1027    fn deref(&self) -> &Self::Target {
1028        // SAFETY:
1029        // Safe because the AVFrame has been properly initialized during construction.
1030        unsafe { self.0.as_ref().unwrap() }
1031    }
1032}
1033
1034impl Drop for AvFrame {
1035    fn drop(&mut self) {
1036        // SAFETY:
1037        // Safe because the AVFrame is valid through the life of this object and fully owned by us.
1038        unsafe { ffi::av_frame_free(&mut self.0) };
1039    }
1040}
1041
1042#[cfg(test)]
1043mod tests {
1044    use std::sync::atomic::AtomicBool;
1045    use std::sync::atomic::Ordering;
1046    use std::sync::Arc;
1047
1048    use super::*;
1049
1050    #[test]
1051    fn test_averror() {
1052        // Just test that the error is wrapper properly. The bindings test module already checks
1053        // that the error bindings correspond to the right ffmpeg errors.
1054        let averror = AvError(AVERROR_EOF);
1055        let msg = format!("{averror}");
1056        assert_eq!(msg, "End of file");
1057
1058        let averror = AvError(0);
1059        let msg = format!("{averror}");
1060        assert_eq!(msg, "Success");
1061
1062        let averror = AvError(10);
1063        let msg = format!("{averror}");
1064        assert_eq!(msg, "Unknown avcodec error 10");
1065    }
1066
1067    // Test that the AVPacket wrapper frees the owned AVBuffer on drop.
1068    #[test]
1069    fn test_avpacket_drop() {
1070        struct DropTestBufferSource {
1071            dropped: Arc<AtomicBool>,
1072        }
1073        impl Drop for DropTestBufferSource {
1074            fn drop(&mut self) {
1075                self.dropped.store(true, Ordering::SeqCst);
1076            }
1077        }
1078        impl AvBufferSource for DropTestBufferSource {
1079            fn as_ptr(&self) -> *const u8 {
1080                [].as_ptr()
1081            }
1082
1083            fn len(&self) -> usize {
1084                0
1085            }
1086
1087            fn is_empty(&self) -> bool {
1088                true
1089            }
1090        }
1091
1092        let dropped = Arc::new(AtomicBool::new(false));
1093
1094        let pkt = AvPacket::new_owned(
1095            0,
1096            AvBuffer::new(DropTestBufferSource {
1097                dropped: dropped.clone(),
1098            })
1099            .unwrap(),
1100        );
1101        assert!(!dropped.load(Ordering::SeqCst));
1102        drop(pkt);
1103        assert!(dropped.load(Ordering::SeqCst));
1104    }
1105}