From 0c19ea7167eca71ffae351c7c889ec33e2fe76a7 Mon Sep 17 00:00:00 2001
From: Ian Douglas Scott <ian@iandouglasscott.com>
Date: Thu, 16 Nov 2023 19:38:42 -0800
Subject: [PATCH] Add `CaptureImage` type

---
 src/main.rs               |  4 ++--
 src/view/mod.rs           | 28 ++++++++++++----------------
 src/wayland/mod.rs        | 10 +++++++---
 src/wayland/screencopy.rs |  5 +++--
 4 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index daa3211..e66a662 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -82,7 +82,7 @@ enum Msg {
 #[derive(Debug)]
 struct Workspace {
     name: String,
-    img_for_output: HashMap<wl_output::WlOutput, iced::widget::image::Handle>,
+    img_for_output: HashMap<wl_output::WlOutput, wayland::CaptureImage>,
     handle: zcosmic_workspace_handle_v1::ZcosmicWorkspaceHandleV1,
     outputs: HashSet<wl_output::WlOutput>,
     is_active: bool,
@@ -92,7 +92,7 @@ struct Workspace {
 struct Toplevel {
     handle: zcosmic_toplevel_handle_v1::ZcosmicToplevelHandleV1,
     info: ToplevelInfo,
-    img: Option<iced::widget::image::Handle>,
+    img: Option<wayland::CaptureImage>,
 }
 
 #[derive(Clone)]
diff --git a/src/view/mod.rs b/src/view/mod.rs
index f3c35aa..eeda07b 100644
--- a/src/view/mod.rs
+++ b/src/view/mod.rs
@@ -1,7 +1,7 @@
 use cctk::wayland_client::protocol::wl_output;
 use cosmic::iced::{self, widget};
 
-use crate::{App, DragSurface, LayerSurface, Msg, Toplevel, Workspace};
+use crate::{wayland::CaptureImage, App, DragSurface, LayerSurface, Msg, Toplevel, Workspace};
 
 pub(crate) fn layer_surface<'a>(
     app: &'a App,
@@ -51,17 +51,7 @@ pub(crate) fn workspace_item<'a>(
     widget::column![
         close_button(Msg::CloseWorkspace(workspace.handle.clone())),
         cosmic::widget::button(widget::column![
-            widget::Image::new(
-                workspace
-                    .img_for_output
-                    .get(output)
-                    .cloned()
-                    .unwrap_or_else(|| widget::image::Handle::from_pixels(
-                        1,
-                        1,
-                        vec![0, 0, 0, 255]
-                    ))
-            ),
+            capture_image(workspace.img_for_output.get(output)),
             widget::text(&workspace.name)
         ])
         .style(theme)
@@ -115,10 +105,8 @@ fn workspaces_sidebar<'a>(
 fn toplevel_preview(toplevel: &Toplevel) -> cosmic::Element<Msg> {
     widget::column![
         close_button(Msg::CloseToplevel(toplevel.handle.clone())),
-        widget::button(widget::Image::new(toplevel.img.clone().unwrap_or_else(
-            || widget::image::Handle::from_pixels(1, 1, vec![0, 0, 0, 255]),
-        )))
-        .on_press(Msg::ActivateToplevel(toplevel.handle.clone())),
+        widget::button(capture_image(toplevel.img.as_ref()))
+            .on_press(Msg::ActivateToplevel(toplevel.handle.clone())),
         widget::text(&toplevel.info.title)
             .horizontal_alignment(iced::alignment::Horizontal::Center)
     ]
@@ -136,3 +124,11 @@ fn toplevel_previews<'a>(
         .align_items(iced::Alignment::Center)
         .into()
 }
+
+fn capture_image(image: Option<&CaptureImage>) -> cosmic::Element<'_, Msg> {
+    if let Some(image) = image {
+        widget::Image::new(image.img.clone()).into()
+    } else {
+        widget::Image::new(widget::image::Handle::from_pixels(1, 1, vec![0, 0, 0, 255])).into()
+    }
+}
diff --git a/src/wayland/mod.rs b/src/wayland/mod.rs
index fb19859..3370d2c 100644
--- a/src/wayland/mod.rs
+++ b/src/wayland/mod.rs
@@ -33,7 +33,6 @@ use cctk::{
 use cosmic::iced::{
     self,
     futures::{executor::block_on, FutureExt, SinkExt},
-    widget::image,
 };
 use futures_channel::mpsc;
 use std::{
@@ -67,7 +66,7 @@ pub enum Event {
     WorkspaceCapture(
         zcosmic_workspace_handle_v1::ZcosmicWorkspaceHandleV1,
         wl_output::WlOutput,
-        image::Handle,
+        CaptureImage,
     ),
     NewToplevel(
         zcosmic_toplevel_handle_v1::ZcosmicToplevelHandleV1,
@@ -80,11 +79,16 @@ pub enum Event {
     CloseToplevel(zcosmic_toplevel_handle_v1::ZcosmicToplevelHandleV1),
     ToplevelCapture(
         zcosmic_toplevel_handle_v1::ZcosmicToplevelHandleV1,
-        image::Handle,
+        CaptureImage,
     ),
     Seats(Vec<wl_seat::WlSeat>),
 }
 
+#[derive(Clone, Debug)]
+pub struct CaptureImage {
+    pub img: iced::widget::image::Handle,
+}
+
 pub fn subscription(conn: Connection) -> iced::Subscription<Event> {
     iced::subscription::run_with_id("wayland-sub", async { start(conn) }.flatten_stream())
 }
diff --git a/src/wayland/screencopy.rs b/src/wayland/screencopy.rs
index 37e65d8..ca45641 100644
--- a/src/wayland/screencopy.rs
+++ b/src/wayland/screencopy.rs
@@ -4,7 +4,7 @@ use cctk::{
     wayland_client::{Connection, QueueHandle, WEnum},
 };
 
-use super::{AppData, Capture, CaptureSource, Event};
+use super::{AppData, Capture, CaptureImage, CaptureSource, Event};
 
 impl ScreencopyHandler for AppData {
     fn screencopy_state(&mut self) -> &mut ScreencopyState {
@@ -59,7 +59,8 @@ impl ScreencopyHandler for AppData {
         }
 
         let mut buffer = capture.buffer.lock().unwrap();
-        let image = unsafe { buffer.as_mut().unwrap().to_image() };
+        let img = unsafe { buffer.as_mut().unwrap().to_image() };
+        let image = CaptureImage { img };
         match &capture.source {
             CaptureSource::Toplevel(toplevel) => {
                 self.send_event(Event::ToplevelCapture(toplevel.clone(), image))