Skip to main content

ustreamer_encode/
videotoolbox.rs

1//! macOS VideoToolbox HEVC encoder backend.
2//!
3//! The current implementation keeps VideoToolbox's native length-prefixed
4//! HEVC access units and extracts `hvcC` decoder configuration for browser-side
5//! WebCodecs setup.
6
7use std::ptr::{self, NonNull};
8use std::sync::{
9    Mutex,
10    mpsc::{self, Receiver, SyncSender},
11};
12use std::time::{Duration, Instant};
13
14use objc2_core_foundation::{
15    CFArray, CFBoolean, CFData, CFDictionary, CFNumber, CFRetained, CFString, CFType,
16};
17use objc2_core_media::{
18    CMBlockBuffer, CMFormatDescription, CMSampleBuffer, CMTime,
19    kCMFormatDescriptionExtension_SampleDescriptionExtensionAtoms, kCMSampleAttachmentKey_NotSync,
20    kCMTimeInvalid, kCMVideoCodecType_HEVC,
21};
22use objc2_core_video::{
23    CVImageBuffer, CVPixelBuffer, CVPixelBufferCreate, CVPixelBufferGetBaseAddress,
24    CVPixelBufferGetBytesPerRow, CVPixelBufferLockBaseAddress, CVPixelBufferLockFlags,
25    CVPixelBufferUnlockBaseAddress, kCVPixelBufferHeightKey, kCVPixelBufferPixelFormatTypeKey,
26    kCVPixelBufferWidthKey, kCVPixelFormatType_32BGRA, kCVPixelFormatType_64RGBAHalf,
27    kCVReturnSuccess,
28};
29use objc2_video_toolbox::{
30    VTCompressionSession, VTSessionSetProperty, kVTCompressionPropertyKey_AllowFrameReordering,
31    kVTCompressionPropertyKey_AllowTemporalCompression, kVTCompressionPropertyKey_AverageBitRate,
32    kVTCompressionPropertyKey_ExpectedFrameRate, kVTCompressionPropertyKey_MaxKeyFrameInterval,
33    kVTCompressionPropertyKey_ProfileLevel, kVTCompressionPropertyKey_Quality,
34    kVTCompressionPropertyKey_RealTime, kVTEncodeFrameOptionKey_ForceKeyFrame,
35    kVTProfileLevel_HEVC_Main_AutoLevel, kVTProfileLevel_HEVC_Main10_AutoLevel,
36    kVTVideoEncoderSpecification_RequireHardwareAcceleratedVideoEncoder,
37};
38use ustreamer_capture::CapturedFrame;
39use ustreamer_proto::quality::{EncodeMode, EncodeParams};
40
41use crate::{DecoderConfig, EncodeError, EncodedFrame, FrameEncoder};
42
43const CALLBACK_TIMEOUT: Duration = Duration::from_millis(750);
44const DEFAULT_MAIN_CODEC: &str = "hvc1.1.6.L153.B0";
45const DEFAULT_MAIN10_CODEC: &str = "hvc1.2.6.L153.B0";
46
47/// Configuration for the VideoToolbox encoder.
48#[derive(Debug, Clone)]
49pub struct VideoToolboxEncoderConfig {
50    /// RFC 6381 codec string used for 8-bit HEVC streams.
51    pub main_codec: String,
52    /// RFC 6381 codec string used for 10-bit HEVC streams.
53    pub main10_codec: String,
54    /// Timeout used while waiting for VideoToolbox callbacks.
55    pub callback_timeout: Duration,
56}
57
58impl Default for VideoToolboxEncoderConfig {
59    fn default() -> Self {
60        Self {
61            main_codec: DEFAULT_MAIN_CODEC.into(),
62            main10_codec: DEFAULT_MAIN10_CODEC.into(),
63            callback_timeout: CALLBACK_TIMEOUT,
64        }
65    }
66}
67
68/// Direct HEVC encoder backed by `VTCompressionSession`.
69#[derive(Debug)]
70pub struct VideoToolboxEncoder {
71    config: VideoToolboxEncoderConfig,
72    session: Option<SessionHandle>,
73    callback_state: Box<CallbackState>,
74    session_spec: Option<SessionSpec>,
75    runtime_config: Option<RuntimeConfig>,
76    decoder_config: Option<DecoderConfig>,
77    frame_index: u64,
78}
79
80#[derive(Debug)]
81struct CallbackState {
82    pending: Mutex<Option<PendingEncode>>,
83}
84
85#[derive(Debug)]
86struct PendingEncode {
87    sender: SyncSender<Result<CallbackResult, String>>,
88    started_at: Instant,
89    fallback_keyframe: bool,
90    fallback_refine: bool,
91    fallback_lossless: bool,
92}
93
94#[derive(Debug)]
95struct CallbackResult {
96    frame: EncodedFrame,
97    decoder_config: Option<DecoderConfig>,
98}
99
100#[derive(Debug, Clone, Copy, PartialEq, Eq)]
101struct SessionSpec {
102    width: u32,
103    height: u32,
104    pixel_format: u32,
105    profile: HevcProfile,
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109struct RuntimeConfig {
110    target_fps: u32,
111    bitrate_bps: u64,
112    max_bitrate_bps: u64,
113    mode: EncodeMode,
114}
115
116#[derive(Debug, Clone, Copy, PartialEq, Eq)]
117enum HevcProfile {
118    Main,
119    Main10,
120}
121
122#[derive(Debug)]
123struct SessionHandle(NonNull<VTCompressionSession>);
124
125unsafe impl Send for SessionHandle {}
126
127impl SessionHandle {
128    unsafe fn from_raw(ptr: NonNull<VTCompressionSession>) -> Self {
129        Self(ptr)
130    }
131
132    fn as_ref(&self) -> &VTCompressionSession {
133        unsafe { self.0.as_ref() }
134    }
135
136    fn invalidate(&self) {
137        unsafe { self.as_ref().invalidate() };
138    }
139}
140
141impl Drop for SessionHandle {
142    fn drop(&mut self) {
143        let retained = unsafe { CFRetained::<VTCompressionSession>::from_raw(self.0) };
144        drop(retained);
145    }
146}
147
148impl VideoToolboxEncoder {
149    /// Create a new VideoToolbox HEVC encoder with default browser codec strings.
150    pub fn new() -> Self {
151        Self::with_config(VideoToolboxEncoderConfig::default())
152    }
153
154    /// Create a new VideoToolbox encoder with explicit browser codec strings.
155    pub fn with_config(config: VideoToolboxEncoderConfig) -> Self {
156        Self {
157            config,
158            session: None,
159            callback_state: Box::new(CallbackState {
160                pending: Mutex::new(None),
161            }),
162            session_spec: None,
163            runtime_config: None,
164            decoder_config: None,
165            frame_index: 0,
166        }
167    }
168
169    fn encode_pixel_buffer(
170        &mut self,
171        pixel_buffer: &CVPixelBuffer,
172        width: u32,
173        height: u32,
174        pixel_format: u32,
175        params: &EncodeParams,
176    ) -> Result<EncodedFrame, EncodeError> {
177        if width != params.width || height != params.height {
178            return Err(EncodeError::UnsupportedConfig(format!(
179                "VideoToolboxEncoder does not scale frames yet; captured frame is {}x{} but EncodeParams requested {}x{}",
180                width, height, params.width, params.height
181            )));
182        }
183
184        let session_spec = SessionSpec {
185            width,
186            height,
187            pixel_format,
188            profile: profile_for_pixel_format(pixel_format)?,
189        };
190        let runtime_config = RuntimeConfig {
191            target_fps: params.target_fps.max(1),
192            bitrate_bps: params.bitrate_bps,
193            max_bitrate_bps: params.max_bitrate_bps.max(params.bitrate_bps),
194            mode: params.mode,
195        };
196
197        self.ensure_session(session_spec)?;
198        self.apply_runtime_config(runtime_config)?;
199
200        let session = self.session.as_ref().ok_or_else(|| {
201            EncodeError::InitFailed("VideoToolbox session was not created".into())
202        })?;
203
204        let force_keyframe = self.frame_index == 0 || params.force_keyframe;
205        let (tx, rx) = mpsc::sync_channel(1);
206        {
207            let mut pending =
208                self.callback_state.pending.lock().map_err(|_| {
209                    EncodeError::EncodeFailed("callback state lock poisoned".into())
210                })?;
211            *pending = Some(PendingEncode {
212                sender: tx,
213                started_at: Instant::now(),
214                fallback_keyframe: force_keyframe,
215                fallback_refine: matches!(params.mode, EncodeMode::LosslessRefine),
216                fallback_lossless: false,
217            });
218        }
219
220        let force_keyframe_dict = force_keyframe.then(force_keyframe_dictionary);
221        let frame_properties: Option<&CFDictionary> = force_keyframe_dict
222            .as_ref()
223            .map(|dict| (&**dict).as_opaque());
224        let presentation_time =
225            unsafe { CMTime::new(self.frame_index as i64, runtime_config.target_fps as i32) };
226        let duration = unsafe { CMTime::new(1, runtime_config.target_fps as i32) };
227        let image_buffer: &CVImageBuffer = pixel_buffer;
228
229        let status = unsafe {
230            session.as_ref().encode_frame(
231                image_buffer,
232                presentation_time,
233                duration,
234                frame_properties,
235                ptr::null_mut(),
236                ptr::null_mut(),
237            )
238        };
239        if status != 0 {
240            self.clear_pending();
241            return Err(EncodeError::EncodeFailed(format!(
242                "VTCompressionSessionEncodeFrame failed with status {status}"
243            )));
244        }
245
246        let complete_status = unsafe { session.as_ref().complete_frames(presentation_time) };
247        if complete_status != 0 {
248            self.clear_pending();
249            return Err(EncodeError::EncodeFailed(format!(
250                "VTCompressionSessionCompleteFrames failed with status {complete_status}"
251            )));
252        }
253
254        let result = wait_for_callback(&rx, self.config.callback_timeout)?;
255        if let Some(decoder_config) = result.decoder_config {
256            self.decoder_config = Some(DecoderConfig {
257                codec: self.codec_string(session_spec.profile),
258                ..decoder_config
259            });
260        }
261
262        self.frame_index += 1;
263        Ok(result.frame)
264    }
265
266    fn encode_cpu_buffer(
267        &mut self,
268        data: &[u8],
269        width: u32,
270        height: u32,
271        stride: u32,
272        format: wgpu::TextureFormat,
273        params: &EncodeParams,
274    ) -> Result<EncodedFrame, EncodeError> {
275        let pixel_format =
276            cv_pixel_format_for_texture(format).ok_or(EncodeError::UnsupportedConfig(format!(
277                "unsupported CPU buffer texture format {format:?}"
278            )))?;
279
280        let mut pixel_buffer = ptr::null_mut();
281        let status = unsafe {
282            CVPixelBufferCreate(
283                None,
284                width as usize,
285                height as usize,
286                pixel_format,
287                None,
288                NonNull::from(&mut pixel_buffer),
289            )
290        };
291        if status != kCVReturnSuccess {
292            return Err(EncodeError::EncodeFailed(format!(
293                "CVPixelBufferCreate failed with status {status}"
294            )));
295        }
296
297        let pixel_buffer_ptr = NonNull::new(pixel_buffer).ok_or_else(|| {
298            EncodeError::EncodeFailed("CVPixelBufferCreate returned a null pixel buffer".into())
299        })?;
300        let pixel_buffer = unsafe { CFRetained::from_raw(pixel_buffer_ptr) };
301        let lock_flags = CVPixelBufferLockFlags(0);
302        let lock_status = unsafe { CVPixelBufferLockBaseAddress(&pixel_buffer, lock_flags) };
303        if lock_status != kCVReturnSuccess {
304            return Err(EncodeError::EncodeFailed(format!(
305                "CVPixelBufferLockBaseAddress failed with status {lock_status}"
306            )));
307        }
308
309        let copy_result =
310            copy_cpu_frame_into_pixel_buffer(&pixel_buffer, data, width, height, stride, format);
311        let unlock_status = unsafe { CVPixelBufferUnlockBaseAddress(&pixel_buffer, lock_flags) };
312        if unlock_status != kCVReturnSuccess {
313            return Err(EncodeError::EncodeFailed(format!(
314                "CVPixelBufferUnlockBaseAddress failed with status {unlock_status}"
315            )));
316        }
317        copy_result?;
318
319        self.encode_pixel_buffer(&pixel_buffer, width, height, pixel_format, params)
320    }
321
322    fn ensure_session(&mut self, spec: SessionSpec) -> Result<(), EncodeError> {
323        if self.session_spec == Some(spec) && self.session.is_some() {
324            return Ok(());
325        }
326
327        self.invalidate_session();
328
329        let encoder_specification = hardware_encoder_specification();
330        let source_attributes = source_image_buffer_attributes(spec);
331        let mut session_ptr = ptr::null_mut();
332        let output_refcon = (&mut *self.callback_state) as *mut CallbackState as *mut _;
333        let status = unsafe {
334            VTCompressionSession::create(
335                None,
336                spec.width as i32,
337                spec.height as i32,
338                kCMVideoCodecType_HEVC,
339                Some((&*encoder_specification).as_opaque()),
340                Some((&*source_attributes).as_opaque()),
341                None,
342                Some(videotoolbox_output_callback),
343                output_refcon,
344                NonNull::from(&mut session_ptr),
345            )
346        };
347        if status != 0 {
348            return Err(EncodeError::InitFailed(format!(
349                "VTCompressionSessionCreate failed with status {status}"
350            )));
351        }
352
353        let session_ptr = NonNull::new(session_ptr).ok_or_else(|| {
354            EncodeError::InitFailed("VideoToolbox returned a null session".into())
355        })?;
356        let session = unsafe { SessionHandle::from_raw(session_ptr) };
357
358        set_property(
359            session.as_ref(),
360            unsafe { kVTCompressionPropertyKey_RealTime },
361            CFBoolean::new(true),
362        )?;
363        set_property(
364            session.as_ref(),
365            unsafe { kVTCompressionPropertyKey_AllowFrameReordering },
366            CFBoolean::new(false),
367        )?;
368        set_property(
369            session.as_ref(),
370            unsafe { kVTCompressionPropertyKey_ProfileLevel },
371            profile_level(spec.profile),
372        )?;
373
374        let prepare_status = unsafe { session.as_ref().prepare_to_encode_frames() };
375        if prepare_status != 0 {
376            return Err(EncodeError::InitFailed(format!(
377                "VTCompressionSessionPrepareToEncodeFrames failed with status {prepare_status}"
378            )));
379        }
380
381        self.session = Some(session);
382        self.session_spec = Some(spec);
383        self.runtime_config = None;
384        self.decoder_config = None;
385        self.frame_index = 0;
386        Ok(())
387    }
388
389    fn apply_runtime_config(&mut self, config: RuntimeConfig) -> Result<(), EncodeError> {
390        if self.runtime_config == Some(config) {
391            return Ok(());
392        }
393
394        let session = self
395            .session
396            .as_ref()
397            .ok_or_else(|| EncodeError::InitFailed("VideoToolbox session is missing".into()))?;
398        let session = session.as_ref();
399
400        set_property(
401            session,
402            unsafe { kVTCompressionPropertyKey_AllowTemporalCompression },
403            CFBoolean::new(config.mode != EncodeMode::LosslessRefine),
404        )?;
405        let expected_fps = CFNumber::new_i32(config.target_fps.min(i32::MAX as u32) as i32);
406        set_property(
407            session,
408            unsafe { kVTCompressionPropertyKey_ExpectedFrameRate },
409            &expected_fps,
410        )?;
411        let max_keyframe_interval =
412            CFNumber::new_i32(config.target_fps.min(i32::MAX as u32) as i32);
413        set_property(
414            session,
415            unsafe { kVTCompressionPropertyKey_MaxKeyFrameInterval },
416            &max_keyframe_interval,
417        )?;
418        let bitrate = CFNumber::new_i64(config.bitrate_bps.min(i64::MAX as u64) as i64);
419        set_property(
420            session,
421            unsafe { kVTCompressionPropertyKey_AverageBitRate },
422            &bitrate,
423        )?;
424
425        let quality = match config.mode {
426            EncodeMode::Interactive => 0.78,
427            EncodeMode::IdleLowFps => 0.9,
428            EncodeMode::LosslessRefine => 1.0,
429        };
430        let quality = CFNumber::new_f32(quality);
431        set_property(
432            session,
433            unsafe { kVTCompressionPropertyKey_Quality },
434            &quality,
435        )?;
436
437        self.runtime_config = Some(config);
438        Ok(())
439    }
440
441    fn clear_pending(&self) {
442        if let Ok(mut pending) = self.callback_state.pending.lock() {
443            pending.take();
444        }
445    }
446
447    fn invalidate_session(&mut self) {
448        self.clear_pending();
449        if let Some(session) = self.session.take() {
450            session.invalidate();
451        }
452        self.session_spec = None;
453        self.runtime_config = None;
454        self.decoder_config = None;
455        self.frame_index = 0;
456    }
457
458    fn codec_string(&self, profile: HevcProfile) -> String {
459        match profile {
460            HevcProfile::Main => self.config.main_codec.clone(),
461            HevcProfile::Main10 => self.config.main10_codec.clone(),
462        }
463    }
464}
465
466impl Default for VideoToolboxEncoder {
467    fn default() -> Self {
468        Self::new()
469    }
470}
471
472impl Drop for VideoToolboxEncoder {
473    fn drop(&mut self) {
474        self.invalidate_session();
475    }
476}
477
478impl FrameEncoder for VideoToolboxEncoder {
479    fn encode(
480        &mut self,
481        frame: &CapturedFrame,
482        params: &EncodeParams,
483    ) -> Result<EncodedFrame, EncodeError> {
484        match frame {
485            CapturedFrame::MetalPixelBuffer {
486                pixel_buffer,
487                width,
488                height,
489                pixel_format,
490                ..
491            } => self.encode_pixel_buffer(&*pixel_buffer, *width, *height, *pixel_format, params),
492            CapturedFrame::CpuBuffer {
493                data,
494                width,
495                height,
496                stride,
497                format,
498            } => self.encode_cpu_buffer(data, *width, *height, *stride, *format, params),
499        }
500    }
501
502    fn flush(&mut self) -> Result<Vec<EncodedFrame>, EncodeError> {
503        if let Some(session) = &self.session {
504            let status = unsafe { session.as_ref().complete_frames(kCMTimeInvalid) };
505            if status != 0 {
506                return Err(EncodeError::EncodeFailed(format!(
507                    "VTCompressionSessionCompleteFrames(flush) failed with status {status}"
508                )));
509            }
510        }
511
512        Ok(Vec::new())
513    }
514
515    fn decoder_config(&self) -> Option<DecoderConfig> {
516        self.decoder_config.clone()
517    }
518}
519
520fn profile_for_pixel_format(pixel_format: u32) -> Result<HevcProfile, EncodeError> {
521    if pixel_format == kCVPixelFormatType_32BGRA {
522        Ok(HevcProfile::Main)
523    } else if pixel_format == kCVPixelFormatType_64RGBAHalf {
524        Ok(HevcProfile::Main10)
525    } else {
526        Err(EncodeError::UnsupportedConfig(format!(
527            "unsupported VideoToolbox pixel format 0x{pixel_format:08x}"
528        )))
529    }
530}
531
532fn cv_pixel_format_for_texture(format: wgpu::TextureFormat) -> Option<u32> {
533    match format {
534        wgpu::TextureFormat::Bgra8Unorm | wgpu::TextureFormat::Bgra8UnormSrgb => {
535            Some(kCVPixelFormatType_32BGRA)
536        }
537        wgpu::TextureFormat::Rgba16Float => Some(kCVPixelFormatType_64RGBAHalf),
538        _ => None,
539    }
540}
541
542fn copy_cpu_frame_into_pixel_buffer(
543    pixel_buffer: &CVPixelBuffer,
544    data: &[u8],
545    width: u32,
546    height: u32,
547    stride: u32,
548    format: wgpu::TextureFormat,
549) -> Result<(), EncodeError> {
550    let row_bytes = row_bytes_for_texture(format, width)?;
551    let source_stride = stride as usize;
552    if source_stride < row_bytes {
553        return Err(EncodeError::EncodeFailed(format!(
554            "source stride {source_stride} is smaller than required row width {row_bytes}"
555        )));
556    }
557
558    let required_len = source_stride
559        .checked_mul(height as usize)
560        .ok_or_else(|| EncodeError::EncodeFailed("CPU frame size overflow".into()))?;
561    if data.len() < required_len {
562        return Err(EncodeError::EncodeFailed(format!(
563            "CPU frame buffer too small: {} bytes, expected at least {required_len}",
564            data.len()
565        )));
566    }
567
568    let base_address = NonNull::new(CVPixelBufferGetBaseAddress(pixel_buffer))
569        .ok_or_else(|| EncodeError::EncodeFailed("CVPixelBuffer base address was null".into()))?;
570    let destination_stride = CVPixelBufferGetBytesPerRow(pixel_buffer);
571    if destination_stride < row_bytes {
572        return Err(EncodeError::EncodeFailed(format!(
573            "pixel buffer stride {destination_stride} is smaller than row width {row_bytes}"
574        )));
575    }
576
577    for row in 0..height as usize {
578        let source_offset = row * source_stride;
579        let destination_offset = row * destination_stride;
580        let source = &data[source_offset..source_offset + row_bytes];
581        let destination = unsafe { (base_address.as_ptr() as *mut u8).add(destination_offset) };
582        unsafe {
583            ptr::copy_nonoverlapping(source.as_ptr(), destination, row_bytes);
584        }
585    }
586
587    Ok(())
588}
589
590fn row_bytes_for_texture(format: wgpu::TextureFormat, width: u32) -> Result<usize, EncodeError> {
591    let bytes_per_pixel = format
592        .block_copy_size(None)
593        .ok_or(EncodeError::UnsupportedConfig(format!(
594            "unsupported texture format {format:?}"
595        )))?;
596    (width as usize)
597        .checked_mul(bytes_per_pixel as usize)
598        .ok_or_else(|| EncodeError::EncodeFailed("row byte size overflow".into()))
599}
600
601fn profile_level(profile: HevcProfile) -> &'static CFString {
602    unsafe {
603        match profile {
604            HevcProfile::Main => kVTProfileLevel_HEVC_Main_AutoLevel,
605            HevcProfile::Main10 => kVTProfileLevel_HEVC_Main10_AutoLevel,
606        }
607    }
608}
609
610fn hardware_encoder_specification() -> CFRetained<CFDictionary<CFString, CFType>> {
611    CFDictionary::from_slices(
612        &[unsafe { kVTVideoEncoderSpecification_RequireHardwareAcceleratedVideoEncoder }],
613        &[CFBoolean::new(true).as_ref()],
614    )
615}
616
617fn source_image_buffer_attributes(spec: SessionSpec) -> CFRetained<CFDictionary<CFString, CFType>> {
618    let pixel_format = CFNumber::new_i32(spec.pixel_format as i32);
619    let width = CFNumber::new_i32(spec.width.min(i32::MAX as u32) as i32);
620    let height = CFNumber::new_i32(spec.height.min(i32::MAX as u32) as i32);
621
622    CFDictionary::from_slices(
623        &[
624            unsafe { kCVPixelBufferPixelFormatTypeKey },
625            unsafe { kCVPixelBufferWidthKey },
626            unsafe { kCVPixelBufferHeightKey },
627        ],
628        &[pixel_format.as_ref(), width.as_ref(), height.as_ref()],
629    )
630}
631
632fn force_keyframe_dictionary() -> CFRetained<CFDictionary<CFString, CFType>> {
633    CFDictionary::from_slices(
634        &[unsafe { kVTEncodeFrameOptionKey_ForceKeyFrame }],
635        &[CFBoolean::new(true).as_ref()],
636    )
637}
638
639fn set_property<V>(
640    session: &VTCompressionSession,
641    key: &CFString,
642    value: &V,
643) -> Result<(), EncodeError>
644where
645    V: AsRef<CFType> + ?Sized,
646{
647    let session_ref: &CFType = AsRef::<CFType>::as_ref(session);
648    let value_ref: &CFType = value.as_ref();
649    let status = unsafe { VTSessionSetProperty(session_ref, key, Some(value_ref)) };
650    if status == 0 {
651        Ok(())
652    } else {
653        Err(EncodeError::InitFailed(format!(
654            "VTSessionSetProperty failed with status {status}"
655        )))
656    }
657}
658
659fn wait_for_callback(
660    rx: &Receiver<Result<CallbackResult, String>>,
661    timeout: Duration,
662) -> Result<CallbackResult, EncodeError> {
663    match rx.recv_timeout(timeout) {
664        Ok(Ok(result)) => Ok(result),
665        Ok(Err(error)) => Err(EncodeError::EncodeFailed(error)),
666        Err(_) => Err(EncodeError::EncodeFailed(
667            "timed out waiting for VideoToolbox output callback".into(),
668        )),
669    }
670}
671
672unsafe extern "C-unwind" fn videotoolbox_output_callback(
673    output_callback_ref_con: *mut std::ffi::c_void,
674    _source_frame_ref_con: *mut std::ffi::c_void,
675    status: i32,
676    _info_flags: objc2_video_toolbox::VTEncodeInfoFlags,
677    sample_buffer: *mut CMSampleBuffer,
678) {
679    let Some(state_ptr) = NonNull::new(output_callback_ref_con.cast::<CallbackState>()) else {
680        return;
681    };
682    let state = unsafe { state_ptr.as_ref() };
683    let pending = match state.pending.lock() {
684        Ok(mut pending) => pending.take(),
685        Err(_) => None,
686    };
687    let Some(pending) = pending else {
688        return;
689    };
690
691    if status != 0 {
692        let _ = pending.sender.send(Err(format!(
693            "VideoToolbox callback returned status {status}"
694        )));
695        return;
696    }
697
698    let Some(sample_ptr) = NonNull::new(sample_buffer) else {
699        let _ = pending
700            .sender
701            .send(Err("VideoToolbox did not return a CMSampleBuffer".into()));
702        return;
703    };
704    let sample = unsafe { sample_ptr.as_ref() };
705
706    let result = sample_buffer_to_result(
707        sample,
708        pending.started_at,
709        pending.fallback_keyframe,
710        pending.fallback_refine,
711        pending.fallback_lossless,
712    );
713    let _ = pending.sender.send(result);
714}
715
716fn sample_buffer_to_result(
717    sample: &CMSampleBuffer,
718    started_at: Instant,
719    fallback_keyframe: bool,
720    fallback_refine: bool,
721    fallback_lossless: bool,
722) -> Result<CallbackResult, String> {
723    let data_buffer = unsafe { sample.data_buffer() }
724        .ok_or_else(|| "VideoToolbox sample buffer had no CMBlockBuffer".to_string())?;
725    let data = block_buffer_bytes(&data_buffer)?;
726    let format_description = unsafe { sample.format_description() }
727        .ok_or_else(|| "VideoToolbox sample buffer had no format description".to_string())?;
728    let is_keyframe = sample_is_keyframe(sample).unwrap_or(fallback_keyframe);
729    let decoder_config = extract_decoder_config(&format_description)?;
730
731    Ok(CallbackResult {
732        frame: EncodedFrame {
733            data,
734            is_keyframe,
735            is_refine: fallback_refine,
736            is_lossless: fallback_lossless,
737            encode_time_us: started_at.elapsed().as_micros().min(u64::MAX as u128) as u64,
738        },
739        decoder_config,
740    })
741}
742
743fn block_buffer_bytes(block: &CMBlockBuffer) -> Result<Vec<u8>, String> {
744    let length = unsafe { block.data_length() };
745    let mut data = vec![0u8; length];
746    if length == 0 {
747        return Ok(data);
748    }
749
750    let status = unsafe {
751        block.copy_data_bytes(0, length, NonNull::new(data.as_mut_ptr().cast()).unwrap())
752    };
753    if status == 0 {
754        Ok(data)
755    } else {
756        Err(format!(
757            "CMBlockBufferCopyDataBytes failed with status {status}"
758        ))
759    }
760}
761
762fn sample_is_keyframe(sample: &CMSampleBuffer) -> Option<bool> {
763    let attachments = unsafe { sample.sample_attachments_array(false) }?;
764    let attachments: &CFArray<CFDictionary> = unsafe { (&*attachments).cast_unchecked() };
765    let first = attachments.get(0)?;
766    let first: &CFDictionary<CFString, CFType> = unsafe { (&*first).cast_unchecked() };
767    let not_sync = first.get(unsafe { kCMSampleAttachmentKey_NotSync })?;
768    let flag: CFRetained<CFBoolean> = not_sync.downcast().ok()?;
769    Some(!(*flag).as_bool())
770}
771
772fn extract_decoder_config(
773    format_description: &CMFormatDescription,
774) -> Result<Option<DecoderConfig>, String> {
775    let Some(description) = extract_hvcc_description(format_description)? else {
776        return Ok(None);
777    };
778
779    let dimensions =
780        unsafe { objc2_core_media::CMVideoFormatDescriptionGetDimensions(format_description) };
781    let profile = if description_main10_hint(&description) {
782        HevcProfile::Main10
783    } else {
784        HevcProfile::Main
785    };
786
787    Ok(Some(DecoderConfig {
788        codec: match profile {
789            HevcProfile::Main => DEFAULT_MAIN_CODEC.into(),
790            HevcProfile::Main10 => DEFAULT_MAIN10_CODEC.into(),
791        },
792        description: Some(description),
793        coded_width: dimensions.width.max(0) as u32,
794        coded_height: dimensions.height.max(0) as u32,
795    }))
796}
797
798fn extract_hvcc_description(
799    format_description: &CMFormatDescription,
800) -> Result<Option<Vec<u8>>, String> {
801    let Some(extensions) = (unsafe { format_description.extensions() }) else {
802        return Ok(None);
803    };
804    let extensions: &CFDictionary<CFString, CFType> = unsafe { (&*extensions).cast_unchecked() };
805    let Some(sample_atoms) =
806        extensions.get(unsafe { kCMFormatDescriptionExtension_SampleDescriptionExtensionAtoms })
807    else {
808        return Ok(None);
809    };
810    let sample_atoms: CFRetained<CFDictionary> = sample_atoms
811        .downcast()
812        .map_err(|_| "sample description atoms were not a CFDictionary".to_string())?;
813    let sample_atoms: &CFDictionary<CFString, CFType> =
814        unsafe { (&*sample_atoms).cast_unchecked() };
815    let hvcc_key = CFString::from_str("hvcC");
816    let Some(hvcc) = sample_atoms.get(&hvcc_key) else {
817        return Ok(None);
818    };
819    let hvcc: CFRetained<CFData> = hvcc
820        .downcast()
821        .map_err(|_| "hvcC atom payload was not CFData".to_string())?;
822
823    let length = hvcc.length().max(0) as usize;
824    let bytes = unsafe { std::slice::from_raw_parts(hvcc.byte_ptr(), length) };
825    Ok(Some(bytes.to_vec()))
826}
827
828fn description_main10_hint(description: &[u8]) -> bool {
829    description.get(1).is_some_and(|byte| (*byte & 0x1f) == 2)
830}
831
832#[cfg(test)]
833mod tests {
834    use super::{HevcProfile, description_main10_hint, profile_for_pixel_format};
835    use objc2_core_video::{kCVPixelFormatType_32BGRA, kCVPixelFormatType_64RGBAHalf};
836
837    #[test]
838    fn maps_bgra_to_main_profile() {
839        assert_eq!(
840            profile_for_pixel_format(kCVPixelFormatType_32BGRA).unwrap(),
841            HevcProfile::Main
842        );
843    }
844
845    #[test]
846    fn maps_half_float_to_main10_profile() {
847        assert_eq!(
848            profile_for_pixel_format(kCVPixelFormatType_64RGBAHalf).unwrap(),
849            HevcProfile::Main10
850        );
851    }
852
853    #[test]
854    fn detects_main10_from_hvcc_profile_bits() {
855        assert!(!description_main10_hint(&[1, 1]));
856        assert!(description_main10_hint(&[1, 2]));
857    }
858}