1 files changed, 70 insertions, 50 deletions

diff --git a/src/cursor.rs b/src/cursor.rs
index 3a827d0..1ca63ac 100644
--- a/src/cursor.rs
+++ b/src/cursor.rs
@@ -1,11 +1,8 @@
+use std::collections::HashMap;
 use std::io;
 
 use wayrs_client::connection::Connection;
-use wayrs_client::protocol::wl_compositor::WlCompositor;
-use wayrs_client::protocol::wl_pointer::WlPointer;
-use wayrs_client::protocol::wl_shm::Format;
-use wayrs_client::protocol::wl_surface::WlSurface;
-
+use wayrs_client::protocol::*;
 use wayrs_shm_alloc::ShmAlloc;
 
 use xcursor::parser::Image;
@@ -14,65 +11,67 @@ use crate::state::State;
 
 #[derive(Debug, thiserror::Error)]
 pub enum CursorError {
-    #[error("theme not found")]
-    ThemeNotFound,
+    #[error("cursor '{0}' not found")]
+    CursorNotFound(String),
     #[error("theme could not be parsed")]
     ThemeParseError,
     #[error(transparent)]
     ReadError(#[from] io::Error),
 }
 
-pub struct Cursor {
-    pub surface: WlSurface,
-    pub images: Vec<Image>,
+pub struct CursorTheme {
+    xcursor_size: u32,
+    theme: xcursor::CursorTheme,
+    cursors: HashMap<String, Vec<Image>>,
 }
 
-impl Cursor {
-    pub fn new(
-        conn: &mut Connection<State>,
-        compositor: WlCompositor,
-    ) -> Result<Self, CursorError> {
+impl CursorTheme {
+    pub fn new() -> Self {
         let theme_name = std::env::var("XCURSOR_THEME").ok();
         let theme_name = theme_name.as_deref().unwrap_or("default");
 
-        let theme_path = xcursor::CursorTheme::load(theme_name)
-            .load_icon("default")
-            .ok_or(CursorError::ThemeNotFound)?;
-        let raw_theme = std::fs::read(theme_path)?;
+        let xcursor_size = std::env::var("XCURSOR_SIZE")
+            .ok()
+            .and_then(|x| x.parse().ok())
+            .unwrap_or(24);
 
-        let mut images =
-            xcursor::parser::parse_xcursor(&raw_theme).ok_or(CursorError::ThemeParseError)?;
-        images.sort_unstable_by_key(|img| img.size);
+        let theme = xcursor::CursorTheme::load(theme_name);
 
-        let surface = compositor.create_surface(conn);
+        CursorTheme {
+            xcursor_size,
+            theme,
+            cursors: HashMap::new(),
+        }
+    }
 
-        Ok(Cursor { surface, images })
+    /// Use this to handle errors ahead of time.
+    pub fn ensure_cursor_is_loaded(&mut self, cursor: &str) -> Result<(), CursorError> {
+        let _ = self.get_images(cursor)?;
+        Ok(())
     }
 
-    pub fn set(
-        &self,
+    #[allow(clippy::too_many_arguments)]
+    pub fn set_cursor(
+        &mut self,
         conn: &mut Connection<State>,
-        serial: u32,
-        pointer: WlPointer,
+        cursor: &str,
         scale: u32,
+        serial: u32,
         shm: &mut ShmAlloc,
-    ) {
-        let target_size = std::env::var("XCURSOR_SIZE")
-            .ok()
-            .and_then(|x| x.parse().ok())
-            .unwrap_or(24)
-            * scale;
+        surface: WlSurface,
+        pointer: WlPointer,
+    ) -> Result<(), CursorError> {
+        let target_size = self.xcursor_size * scale;
+
+        let images = self.get_images(cursor)?;
 
-        let image = match self
-            .images
-            .binary_search_by_key(&target_size, |img| img.size)
-        {
-            Ok(indx) => &self.images[indx],
-            Err(indx) if indx == 0 => self.images.first().unwrap(),
-            Err(indx) if indx >= self.images.len() => self.images.last().unwrap(),
+        let image = match images.binary_search_by_key(&target_size, |img| img.size) {
+            Ok(indx) => &images[indx],
+            Err(indx) if indx == 0 => images.first().unwrap(),
+            Err(indx) if indx >= images.len() => images.last().unwrap(),
             Err(indx) => {
-                let a = &self.images[indx - 1];
-                let b = &self.images[indx];
+                let a = &images[indx - 1];
+                let b = &images[indx];
                 if target_size - a.size < b.size - target_size {
                     a
                 } else {
@@ -86,24 +85,45 @@ impl Cursor {
             image.width as i32,
             image.height as i32,
             (image.width * 4) as i32,
-            Format::Argb8888,
+            wl_shm::Format::Argb8888,
         );
 
         assert_eq!(image.pixels_argb.len(), canvas.len(),);
         canvas.copy_from_slice(&image.pixels_rgba);
 
-        self.surface.attach(conn, buffer.wl, 0, 0);
-        self.surface
-            .damage_buffer(conn, 0, 0, image.width as i32, image.height as i32);
-        self.surface.set_buffer_scale(conn, scale as i32);
-        self.surface.commit(conn);
+        surface.attach(conn, buffer.wl, 0, 0);
+        surface.damage(conn, 0, 0, i32::MAX, i32::MAX);
+        surface.set_buffer_scale(conn, scale as i32);
+        surface.commit(conn);
 
         pointer.set_cursor(
             conn,
             serial,
-            self.surface,
+            surface,
             (image.xhot / scale) as i32,
             (image.yhot / scale) as i32,
         );
+
+        Ok(())
+    }
+
+    fn get_images(&mut self, cursor: &'_ str) -> Result<&[Image], CursorError> {
+        // Borrow checker does't allow `if let Some(...) = ...` here for some reason :(
+        if self.cursors.get(cursor).is_some() {
+            return Ok(self.cursors.get(cursor).unwrap());
+        }
+
+        let theme_path = self
+            .theme
+            .load_icon(cursor)
+            .ok_or_else(|| CursorError::CursorNotFound(cursor.to_owned()))?;
+
+        let raw_theme = std::fs::read(theme_path)?;
+
+        let mut images =
+            xcursor::parser::parse_xcursor(&raw_theme).ok_or(CursorError::ThemeParseError)?;
+        images.sort_unstable_by_key(|img| img.size);
+
+        Ok(self.cursors.entry(cursor.to_owned()).or_insert(images))
     }
 }