1mod bindings;
2
3use crate::bindings::SCStreamOutputType;
4use block::ConcreteBlock;
5use cocoa::{
6 base::{id, nil, YES},
7 foundation::{NSArray, NSString, NSUInteger},
8};
9use core_foundation::{base::TCFType, number::CFNumberRef, string::CFStringRef};
10use core_media::{CMSampleBuffer, CMSampleBufferRef};
11use futures::StreamExt;
12use gpui::{
13 actions,
14 elements::{Canvas, *},
15 keymap::Binding,
16 platform::current::Surface,
17 Menu, MenuItem, ViewContext,
18};
19use io_surface::IOSurface;
20use log::LevelFilter;
21use objc::{
22 class,
23 declare::ClassDecl,
24 msg_send,
25 runtime::{Object, Sel},
26 sel, sel_impl,
27};
28use parking_lot::Mutex;
29use simplelog::SimpleLogger;
30use std::{ffi::c_void, ptr, slice, str, sync::Arc};
31
32#[allow(non_upper_case_globals)]
33const NSUTF8StringEncoding: NSUInteger = 4;
34
35actions!(capture, [Quit]);
36
37fn main() {
38 SimpleLogger::init(LevelFilter::Info, Default::default()).expect("could not initialize logger");
39
40 gpui::App::new(()).unwrap().run(|cx| {
41 cx.platform().activate(true);
42 cx.add_global_action(quit);
43
44 cx.add_bindings([Binding::new("cmd-q", Quit, None)]);
45 cx.set_menus(vec![Menu {
46 name: "Zed",
47 items: vec![MenuItem::Action {
48 name: "Quit",
49 action: Box::new(Quit),
50 }],
51 }]);
52
53 cx.add_window(Default::default(), |cx| ScreenCaptureView::new(cx));
54 });
55}
56
57struct ScreenCaptureView {
58 surface: Option<io_surface::IOSurface>,
59}
60
61impl gpui::Entity for ScreenCaptureView {
62 type Event = ();
63}
64
65impl ScreenCaptureView {
66 pub fn new(cx: &mut ViewContext<Self>) -> Self {
67 let (surface_tx, mut surface_rx) = postage::watch::channel::<Option<IOSurface>>();
68 let surface_tx = Arc::new(Mutex::new(surface_tx));
69
70 unsafe {
71 let block = ConcreteBlock::new(move |content: id, error: id| {
72 if !error.is_null() {
73 println!(
74 "ERROR {}",
75 string_from_objc(msg_send![error, localizedDescription])
76 );
77 return;
78 }
79
80 let applications: id = msg_send![content, applications];
81 let displays: id = msg_send![content, displays];
82 let display: id = displays.objectAtIndex(0);
83 let display_width: usize = msg_send![display, width];
84 let display_height: usize = msg_send![display, height];
85
86 let mut decl = ClassDecl::new("CaptureOutput", class!(NSObject)).unwrap();
87 decl.add_ivar::<*mut c_void>("callback");
88 decl.add_method(
89 sel!(stream:didOutputSampleBuffer:ofType:),
90 sample_output as extern "C" fn(&Object, Sel, id, id, SCStreamOutputType),
91 );
92 let capture_output_class = decl.register();
93
94 let output: id = msg_send![capture_output_class, alloc];
95 let output: id = msg_send![output, init];
96 let surface_tx = surface_tx.clone();
97
98 let callback = Box::new(move |buffer: CMSampleBufferRef| {
99 let buffer = CMSampleBuffer::wrap_under_get_rule(buffer);
100 let attachments = buffer.attachments();
101 let attachments = attachments.first().expect("no attachments for sample");
102 let string = bindings::SCStreamFrameInfoStatus.0 as CFStringRef;
103 let status = core_foundation::number::CFNumber::wrap_under_get_rule(
104 *attachments.get(string) as CFNumberRef,
105 )
106 .to_i64()
107 .expect("invalid frame info status");
108
109 if status != bindings::SCFrameStatus_SCFrameStatusComplete {
110 println!("received incomplete frame");
111 return;
112 }
113
114 let image_buffer = buffer.image_buffer();
115 let io_surface = image_buffer.io_surface();
116 *surface_tx.lock().borrow_mut() = Some(io_surface);
117 }) as Box<dyn FnMut(CMSampleBufferRef)>;
118 let callback = Box::into_raw(Box::new(callback));
119 (*output).set_ivar("callback", callback as *mut c_void);
120
121 let filter: id = msg_send![class!(SCContentFilter), alloc];
122 let filter: id = msg_send![filter, initWithDisplay: display includingApplications: applications exceptingWindows: nil];
123 // let filter: id = msg_send![filter, initWithDesktopIndependentWindow: window];
124 let config: id = msg_send![class!(SCStreamConfiguration), alloc];
125 let config: id = msg_send![config, init];
126 let _: () = msg_send![config, setWidth: display_width];
127 let _: () = msg_send![config, setHeight: display_height];
128 let _: () = msg_send![config, setMinimumFrameInterval: bindings::CMTimeMake(1, 60)];
129 let _: () = msg_send![config, setQueueDepth: 6];
130 let _: () = msg_send![config, setShowsCursor: YES];
131
132 let stream: id = msg_send![class!(SCStream), alloc];
133 let stream: id = msg_send![stream, initWithFilter: filter configuration: config delegate: output];
134 let error: id = nil;
135 let queue = bindings::dispatch_queue_create(
136 ptr::null(),
137 bindings::NSObject(ptr::null_mut()),
138 );
139
140 let _: () = msg_send![stream,
141 addStreamOutput: output type: bindings::SCStreamOutputType_SCStreamOutputTypeScreen
142 sampleHandlerQueue: queue
143 error: &error
144 ];
145
146 let start_capture_completion = ConcreteBlock::new(move |error: id| {
147 if !error.is_null() {
148 println!(
149 "error starting capture... error? {}",
150 string_from_objc(msg_send![error, localizedDescription])
151 );
152 return;
153 }
154
155 println!("starting capture");
156 });
157
158 assert!(!stream.is_null());
159 let _: () = msg_send![
160 stream,
161 startCaptureWithCompletionHandler: start_capture_completion
162 ];
163 });
164
165 let _: id = msg_send![
166 class!(SCShareableContent),
167 getShareableContentWithCompletionHandler: block
168 ];
169 }
170
171 cx.spawn_weak(|this, mut cx| async move {
172 while let Some(surface) = surface_rx.next().await {
173 if let Some(this) = this.upgrade(&cx) {
174 this.update(&mut cx, |this, cx| {
175 this.surface = surface;
176 println!("NEW SURFACE!");
177 cx.notify();
178 })
179 } else {
180 break;
181 }
182 }
183 })
184 .detach();
185
186 Self { surface: None }
187 }
188}
189
190impl gpui::View for ScreenCaptureView {
191 fn ui_name() -> &'static str {
192 "View"
193 }
194
195 fn render(&mut self, _: &mut gpui::RenderContext<Self>) -> gpui::ElementBox {
196 let surface = self.surface.clone();
197 Canvas::new(move |bounds, _, cx| {
198 if let Some(native_surface) = surface.clone() {
199 cx.scene.push_surface(Surface {
200 bounds,
201 native_surface,
202 });
203 }
204 })
205 .boxed()
206 }
207}
208
209pub unsafe fn string_from_objc(string: id) -> String {
210 if string.is_null() {
211 Default::default()
212 } else {
213 let len = msg_send![string, lengthOfBytesUsingEncoding: NSUTF8StringEncoding];
214 let bytes = string.UTF8String() as *const u8;
215 str::from_utf8(slice::from_raw_parts(bytes, len))
216 .unwrap()
217 .to_string()
218 }
219}
220
221extern "C" fn sample_output(
222 this: &Object,
223 _: Sel,
224 _stream: id,
225 buffer: id,
226 _kind: SCStreamOutputType,
227) {
228 unsafe {
229 let callback = *this.get_ivar::<*mut c_void>("callback");
230 let callback = &mut *(callback as *mut Box<dyn FnMut(CMSampleBufferRef)>);
231 (*callback)(buffer as CMSampleBufferRef);
232 }
233}
234
235fn quit(_: &Quit, cx: &mut gpui::MutableAppContext) {
236 cx.platform().quit();
237}
238
239mod core_media {
240 #![allow(non_snake_case)]
241
242 use crate::core_video::{CVImageBuffer, CVImageBufferRef};
243 use core_foundation::{
244 array::{CFArray, CFArrayRef},
245 base::{CFTypeID, TCFType},
246 declare_TCFType,
247 dictionary::CFDictionary,
248 impl_CFTypeDescription, impl_TCFType,
249 string::CFString,
250 };
251 use std::ffi::c_void;
252
253 #[repr(C)]
254 pub struct __CMSampleBuffer(c_void);
255 // The ref type must be a pointer to the underlying struct.
256 pub type CMSampleBufferRef = *const __CMSampleBuffer;
257
258 declare_TCFType!(CMSampleBuffer, CMSampleBufferRef);
259 impl_TCFType!(CMSampleBuffer, CMSampleBufferRef, CMSampleBufferGetTypeID);
260 impl_CFTypeDescription!(CMSampleBuffer);
261
262 impl CMSampleBuffer {
263 pub fn attachments(&self) -> Vec<CFDictionary<CFString>> {
264 unsafe {
265 let attachments =
266 CMSampleBufferGetSampleAttachmentsArray(self.as_concrete_TypeRef(), true);
267 CFArray::<CFDictionary>::wrap_under_get_rule(attachments)
268 .into_iter()
269 .map(|attachments| {
270 CFDictionary::wrap_under_get_rule(attachments.as_concrete_TypeRef())
271 })
272 .collect()
273 }
274 }
275
276 pub fn image_buffer(&self) -> CVImageBuffer {
277 unsafe {
278 CVImageBuffer::wrap_under_get_rule(CMSampleBufferGetImageBuffer(
279 self.as_concrete_TypeRef(),
280 ))
281 }
282 }
283 }
284
285 extern "C" {
286 fn CMSampleBufferGetTypeID() -> CFTypeID;
287 fn CMSampleBufferGetSampleAttachmentsArray(
288 buffer: CMSampleBufferRef,
289 create_if_necessary: bool,
290 ) -> CFArrayRef;
291 fn CMSampleBufferGetImageBuffer(buffer: CMSampleBufferRef) -> CVImageBufferRef;
292 }
293}
294
295mod core_video {
296 #![allow(non_snake_case)]
297
298 use core_foundation::{
299 base::{CFTypeID, TCFType},
300 declare_TCFType, impl_CFTypeDescription, impl_TCFType,
301 };
302 use io_surface::{IOSurface, IOSurfaceRef};
303 use std::ffi::c_void;
304
305 #[repr(C)]
306 pub struct __CVImageBuffer(c_void);
307 // The ref type must be a pointer to the underlying struct.
308 pub type CVImageBufferRef = *const __CVImageBuffer;
309
310 declare_TCFType!(CVImageBuffer, CVImageBufferRef);
311 impl_TCFType!(CVImageBuffer, CVImageBufferRef, CVImageBufferGetTypeID);
312 impl_CFTypeDescription!(CVImageBuffer);
313
314 impl CVImageBuffer {
315 pub fn io_surface(&self) -> IOSurface {
316 unsafe {
317 IOSurface::wrap_under_get_rule(CVPixelBufferGetIOSurface(
318 self.as_concrete_TypeRef(),
319 ))
320 }
321 }
322
323 pub fn width(&self) -> usize {
324 unsafe { CVPixelBufferGetWidth(self.as_concrete_TypeRef()) }
325 }
326
327 pub fn height(&self) -> usize {
328 unsafe { CVPixelBufferGetHeight(self.as_concrete_TypeRef()) }
329 }
330 }
331
332 extern "C" {
333 fn CVImageBufferGetTypeID() -> CFTypeID;
334 fn CVPixelBufferGetIOSurface(buffer: CVImageBufferRef) -> IOSurfaceRef;
335 fn CVPixelBufferGetWidth(buffer: CVImageBufferRef) -> usize;
336 fn CVPixelBufferGetHeight(buffer: CVImageBufferRef) -> usize;
337 }
338}