Skip to main content

video_toolbox/
compression_session.rs

1use std::ptr::{null, null_mut};
2
3use bitflags::bitflags;
4use block::{Block, ConcreteBlock};
5use core_foundation::{
6    base::{Boolean, CFAllocator, CFAllocatorRef, CFType, CFTypeID, OSStatus, TCFType},
7    declare_TCFType,
8    dictionary::{CFDictionary, CFDictionaryRef},
9    impl_CFTypeDescription, impl_TCFType,
10    string::{CFString, CFStringRef},
11};
12use core_media::{base::CMItemCount, format_description::CMVideoCodecType, sample_buffer::CMSampleBufferRef, time::CMTime, time_range::CMTimeRange};
13use core_video::{
14    image_buffer::{CVImageBuffer, CVImageBufferRef},
15    pixel_buffer_pool::{CVPixelBufferPool, CVPixelBufferPoolRef},
16};
17use libc::c_void;
18
19use crate::{
20    base::status_to_result,
21    errors::VTEncodeInfoFlags,
22    session::{TVTSession, VTSessionRef},
23};
24
25pub type VTCompressionSessionRef = VTSessionRef;
26
27pub type VTCompressionOutputCallback = extern "C" fn(
28    outputCallbackRefCon: *mut c_void,
29    sourceFrameRefCon: *mut c_void,
30    status: OSStatus,
31    infoFlags: VTEncodeInfoFlags,
32    sampleBuffer: CMSampleBufferRef,
33);
34
35pub type VTCompressionOutputHandler = *const Block<(OSStatus, VTEncodeInfoFlags, CMSampleBufferRef), ()>;
36
37bitflags! {
38    #[repr(transparent)]
39    #[derive(Clone, Copy, Debug, Default, PartialEq)]
40    pub struct VTCompressionSessionOptionFlags: u32 {
41        #[doc(alias = "kVTCompressionSessionBeginFinalPass")]
42        const BeginFinalPass = 1 << 0;
43    }
44}
45
46extern "C" {
47    pub static kVTVideoEncoderSpecification_EncoderID: CFStringRef;
48
49    pub fn VTCompressionSessionCreate(
50        allocator: CFAllocatorRef,
51        width: i32,
52        height: i32,
53        codecType: CMVideoCodecType,
54        encoderSpecification: CFDictionaryRef,
55        sourceImageBufferAttributes: CFDictionaryRef,
56        compressedDataAllocator: CFAllocatorRef,
57        outputCallback: *const VTCompressionOutputCallback,
58        outputCallbackRefCon: *mut c_void,
59        compressionSessionOut: *mut VTCompressionSessionRef,
60    ) -> OSStatus;
61    pub fn VTCompressionSessionInvalidate(session: VTCompressionSessionRef) -> ();
62    pub fn VTCompressionSessionGetTypeID() -> CFTypeID;
63    pub fn VTCompressionSessionGetPixelBufferPool(session: VTCompressionSessionRef) -> CVPixelBufferPoolRef;
64    pub fn VTCompressionSessionPrepareToEncodeFrames(session: VTCompressionSessionRef) -> OSStatus;
65    pub fn VTCompressionSessionEncodeFrame(
66        session: VTCompressionSessionRef,
67        imageBuffer: CVImageBufferRef,
68        presentationTimeStamp: CMTime,
69        duration: CMTime,
70        frameProperties: CFDictionaryRef,
71        sourceFrameRefcon: *mut c_void,
72        infoFlagsOut: *mut VTEncodeInfoFlags,
73    ) -> OSStatus;
74    pub fn VTCompressionSessionEncodeFrameWithOutputHandler(
75        session: VTCompressionSessionRef,
76        imageBuffer: CVImageBufferRef,
77        presentationTimeStamp: CMTime,
78        duration: CMTime,
79        frameProperties: CFDictionaryRef,
80        infoFlagsOut: *mut VTEncodeInfoFlags,
81        outputHandler: VTCompressionOutputHandler,
82    ) -> OSStatus;
83    pub fn VTCompressionSessionCompleteFrames(session: VTCompressionSessionRef, completeUntilPresentationTimeStamp: CMTime) -> OSStatus;
84    pub fn VTCompressionSessionBeginPass(
85        session: VTCompressionSessionRef,
86        beginPassFlags: VTCompressionSessionOptionFlags,
87        reserved: *mut u32,
88    ) -> OSStatus;
89    pub fn VTCompressionSessionEndPass(session: VTCompressionSessionRef, furtherPassesRequestedOut: *mut Boolean, reserved: *mut u32) -> OSStatus;
90    pub fn VTCompressionSessionGetTimeRangesForNextPass(
91        session: VTCompressionSessionRef,
92        timeRangeCountOut: *mut CMItemCount,
93        timeRangeArrayOut: *const CMTimeRange,
94    ) -> OSStatus;
95}
96
97declare_TCFType!(VTCompressionSession, VTCompressionSessionRef);
98impl_TCFType!(VTCompressionSession, VTCompressionSessionRef, VTCompressionSessionGetTypeID);
99impl_CFTypeDescription!(VTCompressionSession);
100
101impl TVTSession for VTCompressionSession {}
102
103impl VTCompressionSession {
104    pub fn new(
105        width: i32,
106        height: i32,
107        codec_type: CMVideoCodecType,
108        encoder_specification: Option<&CFDictionary<CFString, CFType>>,
109        source_image_buffer_attributes: Option<&CFDictionary<CFString, CFType>>,
110        compressed_data_allocator: CFAllocator,
111    ) -> Result<VTCompressionSession, OSStatus> {
112        unsafe {
113            Self::new_with_callback(
114                width,
115                height,
116                codec_type,
117                encoder_specification,
118                source_image_buffer_attributes,
119                compressed_data_allocator,
120                None,
121                None,
122            )
123        }
124    }
125    pub unsafe fn new_with_callback(
126        width: i32,
127        height: i32,
128        codec_type: CMVideoCodecType,
129        encoder_specification: Option<&CFDictionary<CFString, CFType>>,
130        source_image_buffer_attributes: Option<&CFDictionary<CFString, CFType>>,
131        compressed_data_allocator: CFAllocator,
132        output_callback: Option<*const VTCompressionOutputCallback>,
133        output_callback_ref_con: Option<*mut c_void>,
134    ) -> Result<VTCompressionSession, OSStatus> {
135        let mut session: VTCompressionSessionRef = null_mut();
136        let status = VTCompressionSessionCreate(
137            compressed_data_allocator.as_concrete_TypeRef(),
138            width,
139            height,
140            codec_type,
141            encoder_specification.map_or(null(), |s| s.as_concrete_TypeRef()),
142            source_image_buffer_attributes.map_or(null(), |a| a.as_concrete_TypeRef()),
143            compressed_data_allocator.as_concrete_TypeRef(),
144            output_callback.unwrap_or(null()),
145            output_callback_ref_con.unwrap_or(null_mut()),
146            &mut session,
147        );
148        status_to_result(status).map(|_| unsafe { VTCompressionSession::wrap_under_create_rule(session) })
149    }
150
151    pub fn invalidate(&self) {
152        unsafe { VTCompressionSessionInvalidate(self.as_concrete_TypeRef()) }
153    }
154
155    pub fn get_pixel_buffer_pool(&self) -> CVPixelBufferPool {
156        unsafe {
157            let pool_ref = VTCompressionSessionGetPixelBufferPool(self.as_concrete_TypeRef());
158            CVPixelBufferPool::wrap_under_get_rule(pool_ref)
159        }
160    }
161
162    pub fn prepare_to_encode_frames(&self) -> Result<(), OSStatus> {
163        let status = unsafe { VTCompressionSessionPrepareToEncodeFrames(self.as_concrete_TypeRef()) };
164        status_to_result(status)
165    }
166
167    pub unsafe fn encode_frame(
168        &self,
169        image_buffer: CVImageBuffer,
170        presentation_time_stamp: CMTime,
171        duration: CMTime,
172        frame_properties: Option<&CFDictionary<CFString, CFType>>,
173        source_frame_refcon: Option<*mut c_void>,
174    ) -> Result<VTEncodeInfoFlags, OSStatus> {
175        let mut info_flags_out: VTEncodeInfoFlags = VTEncodeInfoFlags::empty();
176
177        let status = VTCompressionSessionEncodeFrame(
178            self.as_concrete_TypeRef(),
179            image_buffer.as_concrete_TypeRef(),
180            presentation_time_stamp,
181            duration,
182            frame_properties.map_or(null(), |p| p.as_concrete_TypeRef()),
183            source_frame_refcon.unwrap_or(null_mut()),
184            &mut info_flags_out,
185        );
186
187        status_to_result(status).map(|_| info_flags_out)
188    }
189
190    pub fn encode_frame_with_closure<F>(
191        &self,
192        image_buffer: CVImageBuffer,
193        presentation_time_stamp: CMTime,
194        duration: CMTime,
195        frame_properties: Option<&CFDictionary<CFString, CFType>>,
196        closure: F,
197    ) -> Result<VTEncodeInfoFlags, OSStatus>
198    where
199        F: Fn(OSStatus, VTEncodeInfoFlags, CMSampleBufferRef) + 'static,
200    {
201        let handler = ConcreteBlock::new(move |status: OSStatus, info_flags: VTEncodeInfoFlags, sample_buffer: CMSampleBufferRef| {
202            closure(status, info_flags, sample_buffer)
203        })
204        .copy();
205
206        let mut info_flags_out: VTEncodeInfoFlags = VTEncodeInfoFlags::empty();
207
208        let status = unsafe {
209            VTCompressionSessionEncodeFrameWithOutputHandler(
210                self.as_concrete_TypeRef(),
211                image_buffer.as_concrete_TypeRef(),
212                presentation_time_stamp,
213                duration,
214                frame_properties.map_or(null(), |p| p.as_concrete_TypeRef()),
215                &mut info_flags_out,
216                &*handler,
217            )
218        };
219
220        status_to_result(status).map(|_| info_flags_out)
221    }
222
223    pub fn complete_frames(&self, complete_until_presentation_time_stamp: CMTime) -> Result<(), OSStatus> {
224        let status = unsafe { VTCompressionSessionCompleteFrames(self.as_concrete_TypeRef(), complete_until_presentation_time_stamp) };
225        status_to_result(status)
226    }
227
228    pub fn begin_pass(&self, begin_pass_flags: VTCompressionSessionOptionFlags) -> Result<(), OSStatus> {
229        let status = unsafe { VTCompressionSessionBeginPass(self.as_concrete_TypeRef(), begin_pass_flags, null_mut()) };
230        status_to_result(status)
231    }
232
233    pub fn end_pass(&self) -> Result<Boolean, OSStatus> {
234        let mut further_passes_requested_out: Boolean = 0;
235        let status = unsafe { VTCompressionSessionEndPass(self.as_concrete_TypeRef(), &mut further_passes_requested_out, null_mut()) };
236        status_to_result(status).map(|_| further_passes_requested_out)
237    }
238
239    pub fn get_time_ranges_for_next_pass(&self) -> Result<Vec<CMTimeRange>, OSStatus> {
240        let mut time_range_count_out: CMItemCount = 0;
241
242        let status = unsafe { VTCompressionSessionGetTimeRangesForNextPass(self.as_concrete_TypeRef(), &mut time_range_count_out, null()) };
243        status_to_result(status)?;
244
245        let mut time_range_array: Vec<CMTimeRange> = Vec::with_capacity(time_range_count_out as usize);
246        let status = unsafe {
247            let result =
248                VTCompressionSessionGetTimeRangesForNextPass(self.as_concrete_TypeRef(), &mut time_range_count_out, time_range_array.as_mut_ptr());
249            if result == 0 {
250                time_range_array.set_len(time_range_count_out as usize);
251            }
252            result
253        };
254
255        status_to_result(status).map(|_| time_range_array)
256    }
257}