diff --git a/Cargo.lock b/Cargo.lock index 669fedf659..6b34c02f0d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -663,9 +663,9 @@ checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" [[package]] name = "bytes" -version = "1.1.0" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4872d67bab6358e59559027aa3b9157c53d9358c51423c17554809a8858e0f8" +checksum = "ec8a7b6a70fde80372154c65702f00a0f56f3e1c36abbc6c440484be248856db" [[package]] name = "bzip2-sys" @@ -754,8 +754,10 @@ dependencies = [ name = "capture" version = "0.1.0" dependencies = [ + "anyhow", "bindgen", "block", + "bytes", "cocoa", "core-foundation", "core-graphics", @@ -3036,6 +3038,7 @@ dependencies = [ "anyhow", "bindgen", "block", + "bytes", "core-foundation", "foreign-types", "metal", diff --git a/crates/capture/Cargo.toml b/crates/capture/Cargo.toml index 54b1d0990e..c2bee8312a 100644 --- a/crates/capture/Cargo.toml +++ b/crates/capture/Cargo.toml @@ -12,7 +12,9 @@ identifier = "dev.zed.Capture" gpui = { path = "../gpui" } media = { path = "../media" } +anyhow = "1.0.38" block = "0.1" +bytes = "1.2" cocoa = "0.24" core-foundation = "0.9.3" core-graphics = "0.22.3" diff --git a/crates/capture/src/compression_session.rs b/crates/capture/src/compression_session.rs new file mode 100644 index 0000000000..1a771d057b --- /dev/null +++ b/crates/capture/src/compression_session.rs @@ -0,0 +1,178 @@ +use anyhow::Result; +use core_foundation::base::{OSStatus, TCFType}; +use media::{ + core_media::{CMSampleBufferRef, CMSampleTimingInfo, CMVideoCodecType}, + core_video::CVImageBuffer, + video_toolbox::{VTCompressionSession, VTEncodeInfoFlags}, +}; +use std::ffi::c_void; + +pub struct CompressionSession { + session: VTCompressionSession, + output_callback: Box, +} + +impl + CompressionSession +{ + pub fn new(width: usize, height: usize, codec: CMVideoCodecType, callback: F) -> Result { + let callback = Box::new(callback); + let session = VTCompressionSession::new( + width, + height, + codec, + Some(Self::output_callback), + callback.as_ref() as *const _ as *const c_void, + )?; + Ok(Self { + session, + output_callback: callback, + }) + } + + pub fn encode_frame(&self, buffer: &CVImageBuffer, timing: CMSampleTimingInfo) -> Result<()> { + self.session.encode_frame( + buffer.as_concrete_TypeRef(), + timing.presentationTimeStamp, + timing.duration, + ) + } + + extern "C" fn output_callback( + output_callback_ref_con: *mut c_void, + _: *mut c_void, + status: OSStatus, + flags: VTEncodeInfoFlags, + sample_buffer: CMSampleBufferRef, + ) { + let callback = unsafe { &mut *(output_callback_ref_con as *mut F) }; + callback(status, flags, sample_buffer); + } +} + +// unsafe extern "C" fn output( +// output_callback_ref_con: *mut c_void, +// source_frame_ref_con: *mut c_void, +// status: OSStatus, +// info_flags: VTEncodeInfoFlags, +// sample_buffer: CMSampleBufferRef, +// ) { +// if status != 0 { +// println!("error encoding frame, code: {}", status); +// return; +// } +// let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer); + +// let mut is_iframe = false; +// let attachments = sample_buffer.attachments(); +// if let Some(attachments) = attachments.first() { +// is_iframe = attachments +// .find(bindings::kCMSampleAttachmentKey_NotSync as CFStringRef) +// .map_or(true, |not_sync| { +// CFBooleanGetValue(*not_sync as CFBooleanRef) +// }); +// } + +// const START_CODE: [u8; 4] = [0x00, 0x00, 0x00, 0x01]; +// if is_iframe { +// let format_description = sample_buffer.format_description(); +// for ix in 0..format_description.h264_parameter_set_count() { +// let parameter_set = format_description.h264_parameter_set_at_index(ix); +// stream.extend(START_CODE); +// stream.extend(parameter_set); +// } +// } + +// println!("YO!"); +// } + +// static void videoFrameFinishedEncoding(void *outputCallbackRefCon, +// void *sourceFrameRefCon, +// OSStatus status, +// VTEncodeInfoFlags infoFlags, +// CMSampleBufferRef sampleBuffer) { +// // Check if there were any errors encoding +// if (status != noErr) { +// NSLog(@"Error encoding video, err=%lld", (int64_t)status); +// return; +// } + +// // In this example we will use a NSMutableData object to store the +// // elementary stream. +// NSMutableData *elementaryStream = [NSMutableData data]; + +// // Find out if the sample buffer contains an I-Frame. +// // If so we will write the SPS and PPS NAL units to the elementary stream. +// BOOL isIFrame = NO; +// CFArrayRef attachmentsArray = CMSampleBufferGetSampleAttachmentsArray(sampleBuffer, 0); +// if (CFArrayGetCount(attachmentsArray)) { +// CFBooleanRef notSync; +// CFDictionaryRef dict = CFArrayGetValueAtIndex(attachmentsArray, 0); +// BOOL keyExists = CFDictionaryGetValueIfPresent(dict, +// kCMSampleAttachmentKey_NotSync, +// (const void **)¬Sync); +// // An I-Frame is a sync frame +// isIFrame = !keyExists || !CFBooleanGetValue(notSync); +// } + +// // This is the start code that we will write to +// // the elementary stream before every NAL unit +// static const size_t startCodeLength = 4; +// static const uint8_t startCode[] = {0x00, 0x00, 0x00, 0x01}; + +// // Write the SPS and PPS NAL units to the elementary stream before every I-Frame +// if (isIFrame) { +// CMFormatDescriptionRef description = CMSampleBufferGetFormatDescription(sampleBuffer); + +// // Find out how many parameter sets there are +// size_t numberOfParameterSets; +// CMVideoFormatDescriptionGetH264ParameterSetAtIndex(description, +// 0, NULL, NULL, +// &numberOfParameterSets, +// NULL); + +// // Write each parameter set to the elementary stream +// for (int i = 0; i < numberOfParameterSets; i++) { +// const uint8_t *parameterSetPointer; +// size_t parameterSetLength; +// CMVideoFormatDescriptionGetH264ParameterSetAtIndex(description, +// i, +// ¶meterSetPointer, +// ¶meterSetLength, +// NULL, NULL); + +// // Write the parameter set to the elementary stream +// [elementaryStream appendBytes:startCode length:startCodeLength]; +// [elementaryStream appendBytes:parameterSetPointer length:parameterSetLength]; +// } +// } + +// // Get a pointer to the raw AVCC NAL unit data in the sample buffer +// size_t blockBufferLength; +// uint8_t *bufferDataPointer = NULL; +// CMBlockBufferGetDataPointer(CMSampleBufferGetDataBuffer(sampleBuffer), +// 0, +// NULL, +// &blockBufferLength, +// (char **)&bufferDataPointer); + +// // Loop through all the NAL units in the block buffer +// // and write them to the elementary stream with +// // start codes instead of AVCC length headers +// size_t bufferOffset = 0; +// static const int AVCCHeaderLength = 4; +// while (bufferOffset < blockBufferLength - AVCCHeaderLength) { +// // Read the NAL unit length +// uint32_t NALUnitLength = 0; +// memcpy(&NALUnitLength, bufferDataPointer + bufferOffset, AVCCHeaderLength); +// // Convert the length value from Big-endian to Little-endian +// NALUnitLength = CFSwapInt32BigToHost(NALUnitLength); +// // Write start code to the elementary stream +// [elementaryStream appendBytes:startCode length:startCodeLength]; +// // Write the NAL unit without the AVCC length header to the elementary stream +// [elementaryStream appendBytes:bufferDataPointer + bufferOffset + AVCCHeaderLength +// length:NALUnitLength]; +// // Move to the next NAL unit in the block buffer +// bufferOffset += AVCCHeaderLength + NALUnitLength; +// } +// } diff --git a/crates/capture/src/main.rs b/crates/capture/src/main.rs index a727dff37c..b1d8f11795 100644 --- a/crates/capture/src/main.rs +++ b/crates/capture/src/main.rs @@ -1,12 +1,18 @@ mod bindings; +mod compression_session; -use crate::bindings::SCStreamOutputType; +use crate::{bindings::SCStreamOutputType, compression_session::CompressionSession}; use block::ConcreteBlock; +use bytes::BytesMut; use cocoa::{ base::{id, nil, YES}, foundation::{NSArray, NSString, NSUInteger}, }; -use core_foundation::{base::TCFType, number::CFNumberRef, string::CFStringRef}; +use core_foundation::{ + base::TCFType, + number::{CFBooleanGetValue, CFBooleanRef, CFNumberRef}, + string::CFStringRef, +}; use futures::StreamExt; use gpui::{ actions, @@ -17,7 +23,10 @@ use gpui::{ }; use log::LevelFilter; use media::{ - core_media::{kCMVideoCodecType_H264, CMSampleBuffer, CMSampleBufferRef, CMTimeMake}, + core_media::{ + kCMSampleAttachmentKey_NotSync, kCMVideoCodecType_H264, CMSampleBuffer, CMSampleBufferRef, + CMTimeMake, + }, core_video::{self, CVImageBuffer}, video_toolbox::VTCompressionSession, }; @@ -86,12 +95,40 @@ impl ScreenCaptureView { let display: id = displays.objectAtIndex(0); let display_width: usize = msg_send![display, width]; let display_height: usize = msg_send![display, height]; - let compression_session = VTCompressionSession::new( + let mut compression_buffer = BytesMut::new(); + let compression_session = CompressionSession::new( display_width, display_height, kCMVideoCodecType_H264, - None, - ptr::null(), + move |status, flags, sample_buffer| { + if status != 0 { + println!("error encoding frame, code: {}", status); + return; + } + let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer); + + let mut is_iframe = false; + let attachments = sample_buffer.attachments(); + if let Some(attachments) = attachments.first() { + is_iframe = attachments + .find(kCMSampleAttachmentKey_NotSync as CFStringRef) + .map_or(true, |not_sync| { + CFBooleanGetValue(*not_sync as CFBooleanRef) + }); + } + + const START_CODE: [u8; 4] = [0x00, 0x00, 0x00, 0x01]; + if is_iframe { + let format_description = sample_buffer.format_description(); + for ix in 0..format_description.h264_parameter_set_count() { + let parameter_set = + format_description.h264_parameter_set_at_index(ix).unwrap(); + compression_buffer.extend_from_slice(&START_CODE); + compression_buffer.extend_from_slice(parameter_set); + let nal_unit = compression_buffer.split(); + } + } + }, ) .unwrap(); @@ -126,11 +163,7 @@ impl ScreenCaptureView { let timing_info = buffer.sample_timing_info(0).unwrap(); let image_buffer = buffer.image_buffer(); compression_session - .encode_frame( - image_buffer.as_concrete_TypeRef(), - timing_info.presentationTimeStamp, - timing_info.duration, - ) + .encode_frame(&image_buffer, timing_info) .unwrap(); *surface_tx.lock().borrow_mut() = Some(image_buffer); }) as Box; diff --git a/crates/media/Cargo.toml b/crates/media/Cargo.toml index 8773fbbe63..aad2b74c02 100644 --- a/crates/media/Cargo.toml +++ b/crates/media/Cargo.toml @@ -10,6 +10,7 @@ doctest = false [dependencies] anyhow = "1.0" block = "0.1" +bytes = "1.2" core-foundation = "0.9.3" foreign-types = "0.3" metal = "0.21.0" diff --git a/crates/media/src/media.rs b/crates/media/src/media.rs index c6aa1e3d0a..5c5f43af9d 100644 --- a/crates/media/src/media.rs +++ b/crates/media/src/media.rs @@ -195,8 +195,8 @@ pub mod core_media { #![allow(non_snake_case)] pub use crate::bindings::{ - kCMTimeInvalid, kCMVideoCodecType_H264, CMItemIndex, CMSampleTimingInfo, CMTime, - CMTimeMake, CMVideoCodecType, + kCMSampleAttachmentKey_NotSync, kCMTimeInvalid, kCMVideoCodecType_H264, CMItemIndex, + CMSampleTimingInfo, CMTime, CMTimeMake, CMVideoCodecType, }; use crate::core_video::{CVImageBuffer, CVImageBufferRef}; use anyhow::{anyhow, Result}; @@ -357,18 +357,12 @@ pub mod video_toolbox { use super::*; use crate::{ - core_media::{CMSampleBuffer, CMSampleBufferRef, CMTime, CMVideoCodecType}, + core_media::{CMSampleBufferRef, CMTime, CMVideoCodecType}, core_video::CVImageBufferRef, }; use anyhow::{anyhow, Result}; - use bindings::VTEncodeInfoFlags; - use core_foundation::{ - base::OSStatus, - dictionary::CFDictionaryRef, - mach_port::CFAllocatorRef, - number::{CFBooleanGetValue, CFBooleanRef}, - string::CFStringRef, - }; + pub use bindings::VTEncodeInfoFlags; + use core_foundation::{base::OSStatus, dictionary::CFDictionaryRef, mach_port::CFAllocatorRef}; use std::ptr; #[repr(C)] @@ -402,7 +396,7 @@ pub mod video_toolbox { ptr::null(), ptr::null(), ptr::null(), - Some(Self::output), + callback, callback_data, &mut this, ); @@ -418,42 +412,6 @@ pub mod video_toolbox { } } - unsafe extern "C" fn output( - output_callback_ref_con: *mut c_void, - source_frame_ref_con: *mut c_void, - status: OSStatus, - info_flags: VTEncodeInfoFlags, - sample_buffer: CMSampleBufferRef, - ) { - if status != 0 { - println!("error encoding frame, code: {}", status); - return; - } - let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer); - - let mut is_iframe = false; - let attachments = sample_buffer.attachments(); - if let Some(attachments) = attachments.first() { - is_iframe = attachments - .find(bindings::kCMSampleAttachmentKey_NotSync as CFStringRef) - .map_or(true, |not_sync| { - CFBooleanGetValue(*not_sync as CFBooleanRef) - }); - } - - const START_CODE: [u8; 4] = [0x00, 0x00, 0x00, 0x01]; - if is_iframe { - let format_description = sample_buffer.format_description(); - for ix in 0..format_description.h264_parameter_set_count() { - let parameter_set = format_description.h264_parameter_set_at_index(ix); - stream.extend(START_CODE); - stream.extend(parameter_set); - } - } - - println!("YO!"); - } - pub fn encode_frame( &self, buffer: CVImageBufferRef,