From 73e343c23c2b5f8fd3f5f44ec79ea25c74b79c9f Mon Sep 17 00:00:00 2001
From: "mgsloan@gmail.com" <michael@zed.dev>
Date: Sun, 1 Dec 2024 00:33:24 -0700
Subject: [PATCH] Improvements to screen share rendering code

---
 crates/call/src/call.rs                       |  2 -
 crates/gpui/src/platform/linux.rs             |  4 +-
 crates/gpui/src/platform/windows.rs           |  4 +-
 crates/live_kit_client/src/live_kit_client.rs | 39 ++++++++++---------
 .../src/remote_video_track_view.rs            | 13 ++-----
 5 files changed, 27 insertions(+), 35 deletions(-)

diff --git a/crates/call/src/call.rs b/crates/call/src/call.rs
index eb91ff08851cb..60bc70d0a3ec2 100644
--- a/crates/call/src/call.rs
+++ b/crates/call/src/call.rs
@@ -18,8 +18,6 @@ use room::Event;
 use settings::Settings;
 use std::sync::Arc;
 
-#[cfg(not(target_os = "windows"))]
-pub use live_kit_client::play_remote_video_track;
 pub use live_kit_client::{
     track::RemoteVideoTrack, RemoteVideoTrackView, RemoteVideoTrackViewEvent,
 };
diff --git a/crates/gpui/src/platform/linux.rs b/crates/gpui/src/platform/linux.rs
index 08afde4d6dba6..089b52cf1e730 100644
--- a/crates/gpui/src/platform/linux.rs
+++ b/crates/gpui/src/platform/linux.rs
@@ -21,6 +21,4 @@ pub(crate) use wayland::*;
 #[cfg(feature = "x11")]
 pub(crate) use x11::*;
 
-// TODO(mgsloan): This type won't make sense for frame capture. A `type VideoFrame` with this type
-// should be added to `live_kit_client`.
-pub(crate) type PlatformScreenCaptureFrame = std::sync::Arc<crate::RenderImage>;
+pub(crate) type PlatformScreenCaptureFrame = ();
diff --git a/crates/gpui/src/platform/windows.rs b/crates/gpui/src/platform/windows.rs
index 933f0255faef1..51d09f0013f96 100644
--- a/crates/gpui/src/platform/windows.rs
+++ b/crates/gpui/src/platform/windows.rs
@@ -22,6 +22,4 @@ pub(crate) use wrapper::*;
 
 pub(crate) use windows::Win32::Foundation::HWND;
 
-// TODO(mgsloan): This type won't make sense for frame capture. A `type VideoFrame` with this type
-// should be added to `live_kit_client`.
-pub(crate) type PlatformScreenCaptureFrame = std::sync::Arc<crate::RenderImage>;
+pub(crate) type PlatformScreenCaptureFrame = ();
diff --git a/crates/live_kit_client/src/live_kit_client.rs b/crates/live_kit_client/src/live_kit_client.rs
index 361da47d7c579..41d9c1fd41f2f 100644
--- a/crates/live_kit_client/src/live_kit_client.rs
+++ b/crates/live_kit_client/src/live_kit_client.rs
@@ -429,13 +429,16 @@ fn start_output_stream(
 
 pub fn play_remote_video_track(
     track: &track::RemoteVideoTrack,
-) -> impl Stream<Item = ScreenCaptureFrame> {
+) -> impl Stream<Item = RemoteVideoFrame> {
     NativeVideoStream::new(track.rtc_track())
         .filter_map(|frame| async move { video_frame_buffer_from_webrtc(frame.buffer) })
 }
 
 #[cfg(target_os = "macos")]
-fn video_frame_buffer_from_webrtc(buffer: Box<dyn VideoBuffer>) -> Option<ScreenCaptureFrame> {
+pub type RemoteVideoFrame = media::core_video::CVImageBuffer;
+
+#[cfg(target_os = "macos")]
+fn video_frame_buffer_from_webrtc(buffer: Box<dyn VideoBuffer>) -> Option<RemoteVideoFrame> {
     use core_foundation::base::TCFType as _;
     use media::core_video::CVImageBuffer;
 
@@ -445,15 +448,14 @@ fn video_frame_buffer_from_webrtc(buffer: Box<dyn VideoBuffer>) -> Option<Screen
         return None;
     }
 
-    unsafe {
-        Some(ScreenCaptureFrame(CVImageBuffer::wrap_under_get_rule(
-            pixel_buffer as _,
-        )))
-    }
+    unsafe { Some(CVImageBuffer::wrap_under_get_rule(pixel_buffer as _)) }
 }
 
-#[cfg(any(target_os = "linux", target_os = "windows"))]
-fn video_frame_buffer_from_webrtc(buffer: Box<dyn VideoBuffer>) -> Option<ScreenCaptureFrame> {
+#[cfg(not(target_os = "macos"))]
+pub type RemoteVideoFrame = Arc<gpui::RenderImage>;
+
+#[cfg(not(any(target_os = "macos", target_os = "windows")))]
+fn video_frame_buffer_from_webrtc(buffer: Box<dyn VideoBuffer>) -> Option<RemoteVideoFrame> {
     use gpui::RenderImage;
     use image::{Frame, RgbaImage};
     use livekit::webrtc::prelude::VideoFormatType;
@@ -467,7 +469,10 @@ fn video_frame_buffer_from_webrtc(buffer: Box<dyn VideoBuffer>) -> Option<Screen
     let bgra_frame_vec = unsafe {
         // Motivation for this unsafe code is to avoid initializing the frame data, since to_argb
         // will write all bytes anyway.
-        let start_ptr = alloc(Layout::array::<u8>(byte_len).unwrap());
+        let start_ptr = alloc(Layout::array::<u8>(byte_len).log_err()?);
+        if start_ptr.is_null() {
+            return None;
+        }
         let bgra_frame_slice = std::slice::from_raw_parts_mut(start_ptr, byte_len);
         buffer.to_argb(
             VideoFormatType::BGRA,
@@ -479,15 +484,13 @@ fn video_frame_buffer_from_webrtc(buffer: Box<dyn VideoBuffer>) -> Option<Screen
         Vec::from_raw_parts(start_ptr, byte_len, byte_len)
     };
 
-    Some(ScreenCaptureFrame(Arc::new(RenderImage::new(
-        SmallVec::from_elem(
-            Frame::new(
-                RgbaImage::from_raw(width, height, bgra_frame_vec)
-                    .with_context(|| "Bug: not enough bytes allocated for image.")
-                    .unwrap(),
-            ),
-            1,
+    Some(Arc::new(RenderImage::new(SmallVec::from_elem(
+        Frame::new(
+            RgbaImage::from_raw(width, height, bgra_frame_vec)
+                .with_context(|| "Bug: not enough bytes allocated for image.")
+                .log_err()?,
         ),
+        1,
     ))))
 }
 
diff --git a/crates/live_kit_client/src/remote_video_track_view.rs b/crates/live_kit_client/src/remote_video_track_view.rs
index 9137ba2f235cc..6cc4b25ac6f35 100644
--- a/crates/live_kit_client/src/remote_video_track_view.rs
+++ b/crates/live_kit_client/src/remote_video_track_view.rs
@@ -1,14 +1,11 @@
 use crate::track::RemoteVideoTrack;
 use anyhow::Result;
 use futures::StreamExt as _;
-use gpui::{
-    img, Empty, EventEmitter, IntoElement, Render, ScreenCaptureFrame, Task, View, ViewContext,
-    VisualContext as _,
-};
+use gpui::{Empty, EventEmitter, IntoElement, Render, Task, View, ViewContext, VisualContext as _};
 
 pub struct RemoteVideoTrackView {
     track: RemoteVideoTrack,
-    frame: Option<ScreenCaptureFrame>,
+    frame: Option<crate::RemoteVideoFrame>,
     _maintain_frame: Task<Result<()>>,
 }
 
@@ -51,14 +48,12 @@ impl Render for RemoteVideoTrackView {
         #[cfg(target_os = "macos")]
         if let Some(frame) = &self.frame {
             use gpui::Styled as _;
-            return gpui::surface(frame.0.clone())
-                .size_full()
-                .into_any_element();
+            return gpui::surface(frame.clone()).size_full().into_any_element();
         }
 
         #[cfg(not(target_os = "macos"))]
         if let Some(frame) = &self.frame {
-            return img(frame.0.clone()).into_any_element();
+            return gpui::img(frame.clone()).into_any_element();
         }
 
         Empty.into_any_element()