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 log::LevelFilter;
20use media::core_video::{self, CVImageBuffer};
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 image_buffer: Option<core_video::CVImageBuffer>,
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 (image_buffer_tx, mut image_buffer_rx) =
68 postage::watch::channel::<Option<CVImageBuffer>>();
69 let image_buffer_tx = Arc::new(Mutex::new(image_buffer_tx));
70
71 unsafe {
72 let block = ConcreteBlock::new(move |content: id, error: id| {
73 if !error.is_null() {
74 println!(
75 "ERROR {}",
76 string_from_objc(msg_send![error, localizedDescription])
77 );
78 return;
79 }
80
81 let applications: id = msg_send![content, applications];
82 let displays: id = msg_send![content, displays];
83 let display: id = displays.objectAtIndex(0);
84 let display_width: usize = msg_send![display, width];
85 let display_height: usize = msg_send![display, height];
86
87 let mut decl = ClassDecl::new("CaptureOutput", class!(NSObject)).unwrap();
88 decl.add_ivar::<*mut c_void>("callback");
89 decl.add_method(
90 sel!(stream:didOutputSampleBuffer:ofType:),
91 sample_output as extern "C" fn(&Object, Sel, id, id, SCStreamOutputType),
92 );
93 let capture_output_class = decl.register();
94
95 let output: id = msg_send![capture_output_class, alloc];
96 let output: id = msg_send![output, init];
97 let surface_tx = image_buffer_tx.clone();
98
99 let callback = Box::new(move |buffer: CMSampleBufferRef| {
100 let buffer = CMSampleBuffer::wrap_under_get_rule(buffer);
101 let attachments = buffer.attachments();
102 let attachments = attachments.first().expect("no attachments for sample");
103 let string = bindings::SCStreamFrameInfoStatus.0 as CFStringRef;
104 let status = core_foundation::number::CFNumber::wrap_under_get_rule(
105 *attachments.get(string) as CFNumberRef,
106 )
107 .to_i64()
108 .expect("invalid frame info status");
109
110 if status != bindings::SCFrameStatus_SCFrameStatusComplete {
111 println!("received incomplete frame");
112 return;
113 }
114
115 let image_buffer = buffer.image_buffer();
116 *surface_tx.lock().borrow_mut() = Some(image_buffer);
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(image_buffer) = image_buffer_rx.next().await {
173 if let Some(this) = this.upgrade(&cx) {
174 this.update(&mut cx, |this, cx| {
175 this.image_buffer = image_buffer;
176 println!("NEW SURFACE!");
177 cx.notify();
178 })
179 } else {
180 break;
181 }
182 }
183 })
184 .detach();
185
186 Self { image_buffer: 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 image_buffer = self.image_buffer.clone();
197 Canvas::new(move |bounds, _, cx| {
198 if let Some(image_buffer) = image_buffer.clone() {
199 cx.scene.push_surface(Surface {
200 bounds,
201 image_buffer,
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}