1use 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#[derive(Debug, Clone)]
49pub struct VideoToolboxEncoderConfig {
50 pub main_codec: String,
52 pub main10_codec: String,
54 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#[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 pub fn new() -> Self {
151 Self::with_config(VideoToolboxEncoderConfig::default())
152 }
153
154 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}