From 1cdea8eff42e4d26caf257edc23b3aac2f96a96b Mon Sep 17 00:00:00 2001
From: Fl1tzi <git@fl1tzi.com>
Date: Sun, 26 Nov 2023 02:54:52 +0100
Subject: [PATCH] add image cache for faster image display

---
 Cargo.toml             |  1 +
 src/device.rs          | 24 ++++++++++---
 src/modules.rs         | 81 +++++++++++++++++++++++++++++++++++++++---
 src/modules/blank.rs   |  6 +++-
 src/modules/counter.rs |  8 +++--
 src/modules/image.rs   | 32 ++++++++++-------
 src/modules/space.rs   |  6 +++-
 7 files changed, 133 insertions(+), 25 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index fdf3996..217fcac 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,3 +21,4 @@ tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 lazy_static = "1.4.0"
 font-loader = "0.11.0"
+clru = "0.6.1"
diff --git a/src/device.rs b/src/device.rs
index c3c1efb..3df00d1 100644
--- a/src/device.rs
+++ b/src/device.rs
@@ -3,9 +3,11 @@ use crate::{
     modules::{retrieve_module_from_name, start_module, HostEvent},
     unwrap_or_error,
 };
+use clru::CLruCache;
 use deck_driver as streamdeck;
 use hidapi::HidApi;
-use std::{collections::HashMap, fmt::Display, sync::Arc};
+use image::DynamicImage;
+use std::{collections::HashMap, fmt::Display, num::NonZeroUsize, sync::Arc};
 use streamdeck::{
     asynchronous::{AsyncStreamDeck, ButtonStateUpdate},
     info::Kind,
@@ -14,7 +16,10 @@ use streamdeck::{
 use tokio::{
     process::Command,
     runtime::Runtime,
-    sync::mpsc::{self, error::TrySendError},
+    sync::{
+        mpsc::{self, error::TrySendError},
+        Mutex,
+    },
 };
 use tracing::{debug, error, info_span, trace, warn};
 
@@ -41,6 +46,11 @@ impl Display for Device {
     }
 }
 
+/// image cache
+///
+/// (button, key), image data
+pub type ImageCache = CLruCache<(u8, u32), Arc<DynamicImage>>;
+
 /// Handles everything related to a single device
 pub struct Device {
     modules: HashMap<u8, ModuleController>,
@@ -50,6 +60,7 @@ pub struct Device {
     spaces: Arc<HashMap<String, Space>>,
     selected_space: Option<String>,
     serial: String,
+    image_cache: Arc<Mutex<ImageCache>>,
 }
 
 impl Device {
@@ -90,6 +101,9 @@ impl Device {
             spaces,
             selected_space: None,
             serial,
+            image_cache: Arc::new(Mutex::new(CLruCache::new(
+                NonZeroUsize::new(button_count.into()).unwrap(),
+            ))),
         })
     }
 
@@ -130,9 +144,11 @@ impl Device {
                 let ser = self.serial.clone();
                 let dev = self.device.clone();
                 let b = btn.clone();
+                let image_cache = self.image_cache.clone();
 
-                runtime
-                    .spawn(async move { start_module(ser, b, module, dev, module_receiver).await });
+                runtime.spawn(async move {
+                    start_module(ser, b, module, dev, module_receiver, image_cache).await
+                });
             }
             // if the receiver already dropped the listener then just directly insert none.
             // Optimizes performance because the key_listener just does not try to send the event.
diff --git a/src/modules.rs b/src/modules.rs
index d807345..a651fa6 100644
--- a/src/modules.rs
+++ b/src/modules.rs
@@ -10,6 +10,9 @@ use self::space::Space;
 
 // other things
 use crate::config::{Button, ButtonConfigError};
+use crate::device::ImageCache;
+use crate::image_rendering::load_image;
+use ::image::imageops::FilterType;
 use ::image::DynamicImage;
 use async_trait::async_trait;
 pub use deck_driver as streamdeck;
@@ -19,8 +22,8 @@ use streamdeck::info::ImageFormat;
 use streamdeck::info::Kind;
 use streamdeck::AsyncStreamDeck;
 use streamdeck::StreamDeckError;
-use tokio::sync::mpsc;
-use tracing::{debug, error};
+use tokio::sync::{mpsc, Mutex};
+use tracing::{debug, error, trace};
 
 /// Events that are coming from the host
 #[derive(Clone, Copy, Debug)]
@@ -34,7 +37,7 @@ pub enum HostEvent {
 pub type ModuleObject = Box<dyn Module + Send + Sync>;
 pub type ModuleFuture =
     Pin<Box<dyn Future<Output = Result<ModuleObject, ButtonConfigError>> + Send>>;
-pub type ModuleInitFunction = fn(Arc<Button>) -> ModuleFuture;
+pub type ModuleInitFunction = fn(Arc<Button>, ModuleCache) -> ModuleFuture;
 
 pub fn retrieve_module_from_name(name: &str) -> Option<ModuleInitFunction> {
     match name {
@@ -54,14 +57,20 @@ pub async fn start_module(
     module_init_function: ModuleInitFunction,
     device: Arc<AsyncStreamDeck>,
     br: ChannelReceiver,
+    image_cache: Arc<Mutex<ImageCache>>,
 ) {
     debug!("STARTED");
+    let mc = ModuleCache::new(
+        image_cache,
+        button.index,
+        device.kind().key_image_format().size,
+    );
     let da = DeviceAccess::new(device, button.index).await;
 
     // run init first
     //
     // panic should be prevented by the config being checked before running
-    let mut module = match module_init_function(button).await {
+    let mut module = match module_init_function(button, mc).await {
         Ok(m) => m,
         Err(e) => panic!("{}", e),
     };
@@ -73,6 +82,65 @@ pub async fn start_module(
     }
 }
 
+/// A wrapper around [ImageCache] to provide easy access to values in the device cache
+pub struct ModuleCache {
+    image_cache: Arc<Mutex<ImageCache>>,
+    button_index: u8,
+    /// Resolution of the deck (required for optimization of storage space)
+    resolution: (usize, usize),
+}
+
+impl ModuleCache {
+    pub fn new(
+        image_cache: Arc<Mutex<ImageCache>>,
+        button_index: u8,
+        resolution: (usize, usize),
+    ) -> Self {
+        ModuleCache {
+            image_cache,
+            button_index,
+            resolution,
+        }
+    }
+
+    /// Load an image from the [ImageCache] or create a new one and insert it into the [ImageCache].
+    /// Returns None if no image was found.
+    ///
+    /// index: Provide an index where your data is cached. With this number the value can be
+    /// accessed again. Use [DeviceAccess::get_image_cached()] for just getting the data.
+    #[allow(dead_code)]
+    pub async fn load_image(&mut self, path: String, index: u32) -> Option<Arc<DynamicImage>> {
+        if let Some(image) = self.get_image(index).await {
+            Some(image)
+        } else {
+            trace!("Decoding image");
+            let mut image = tokio::task::spawn_blocking(move || load_image(path))
+                .await
+                .unwrap()
+                .ok()?;
+            image = image.resize_exact(
+                self.resolution.0 as u32,
+                self.resolution.1 as u32,
+                FilterType::Nearest,
+            );
+            trace!("Decoding finished");
+            let image = Arc::new(image);
+            let mut data = self.image_cache.lock().await;
+            data.put((self.button_index, index), image.clone());
+            trace!("Wrote data into cache (new size: {})", data.len());
+            drop(data);
+            Some(image.into())
+        }
+    }
+
+    /// Just try to retrieve a value from the key (index) in the [ImageCache].
+    #[allow(dead_code)]
+    pub async fn get_image(&self, index: u32) -> Option<Arc<DynamicImage>> {
+        let mut data = self.image_cache.lock().await;
+        data.get(&(self.button_index, index)).cloned()
+    }
+}
+
 /// Wrapper to provide easier access to the Deck
 pub struct DeviceAccess {
     streamdeck: Arc<AsyncStreamDeck>,
@@ -130,7 +198,10 @@ pub type ChannelReceiver = mpsc::Receiver<HostEvent>;
 /// - init() -> function for checking config and creating module
 /// - run() -> function that happens when the device actually runs
 pub trait Module: Sync + Send {
-    async fn init(config: Arc<Button>) -> Result<ModuleObject, ButtonConfigError>
+    async fn init(
+        config: Arc<Button>,
+        mut cache: ModuleCache,
+    ) -> Result<ModuleObject, ButtonConfigError>
     where
         Self: Sized;
     async fn run(
diff --git a/src/modules/blank.rs b/src/modules/blank.rs
index 0528a8d..71adc54 100644
--- a/src/modules/blank.rs
+++ b/src/modules/blank.rs
@@ -3,6 +3,7 @@ use super::ButtonConfigError;
 use super::ChannelReceiver;
 use super::DeviceAccess;
 use super::Module;
+use super::ModuleCache;
 use super::ModuleObject;
 use super::ReturnError;
 use async_trait::async_trait;
@@ -12,7 +13,10 @@ pub struct Blank;
 
 #[async_trait]
 impl Module for Blank {
-    async fn init(_config: Arc<Button>) -> Result<ModuleObject, ButtonConfigError> {
+    async fn init(
+        _config: Arc<Button>,
+        _cache: ModuleCache,
+    ) -> Result<ModuleObject, ButtonConfigError> {
         Ok(Box::new(Blank {}))
     }
 
diff --git a/src/modules/counter.rs b/src/modules/counter.rs
index 465a5f3..a903d1e 100644
--- a/src/modules/counter.rs
+++ b/src/modules/counter.rs
@@ -3,7 +3,8 @@ use std::sync::Arc;
 use crate::config::Button;
 
 use super::{
-    ButtonConfigError, ChannelReceiver, DeviceAccess, HostEvent, Module, ModuleObject, ReturnError,
+    ButtonConfigError, ChannelReceiver, DeviceAccess, HostEvent, Module, ModuleCache, ModuleObject,
+    ReturnError,
 };
 
 use crate::GLOBAL_FONT;
@@ -22,7 +23,10 @@ pub struct Counter {
 
 #[async_trait]
 impl Module for Counter {
-    async fn init(config: Arc<Button>) -> Result<ModuleObject, ButtonConfigError> {
+    async fn init(
+        config: Arc<Button>,
+        _cache: ModuleCache,
+    ) -> Result<ModuleObject, ButtonConfigError> {
         let title = config.parse_module("TITLE", " ".to_string()).res()?;
         let title_size = config.parse_module("TITLE_SIZE", 15.0).res()?;
         let number_size = config.parse_module("NUMBER_SIZE", 25.0).res()?;
diff --git a/src/modules/image.rs b/src/modules/image.rs
index 3f730e4..b434a42 100644
--- a/src/modules/image.rs
+++ b/src/modules/image.rs
@@ -3,32 +3,34 @@ use super::ButtonConfigError;
 use super::ChannelReceiver;
 use super::DeviceAccess;
 use super::Module;
+use super::ModuleCache;
 use super::ModuleObject;
 use super::ReturnError;
-use crate::image_rendering::{load_image, ImageBuilder};
+use crate::image_rendering::ImageBuilder;
 use async_trait::async_trait;
 use image::DynamicImage;
 use std::sync::Arc;
-use tokio::task;
 
 pub struct Image {
-    image: DynamicImage,
+    image: Arc<DynamicImage>,
     scale: f32,
 }
 
 #[async_trait]
 impl Module for Image {
-    async fn init(config: Arc<Button>) -> Result<ModuleObject, ButtonConfigError> {
+    async fn init(
+        config: Arc<Button>,
+        mut cache: ModuleCache,
+    ) -> Result<ModuleObject, ButtonConfigError> {
         let path = config.parse_module("PATH", String::new()).required()?;
         let scale = config.parse_module("SCALE", 100.0).res()?;
 
-        // TODO: decoding takes really long sometimes. Maybe this can be cached?
-        let image = task::spawn_blocking(move || {
-            load_image(path)
-                .map_err(|_| ButtonConfigError::General("Image was not found.".to_string()))
-        })
-        .await
-        .unwrap()?;
+        let image = cache
+            .load_image(path, 1)
+            .await
+            .ok_or(ButtonConfigError::General(
+                "Image was not found".to_string(),
+            ))?;
 
         Ok(Box::new(Image { image, scale }))
     }
@@ -38,7 +40,13 @@ impl Module for Image {
         streamdeck: DeviceAccess,
         _button_receiver: ChannelReceiver,
     ) -> Result<(), ReturnError> {
-        streamdeck.write_img(self.image.clone()).await.unwrap();
+        let (h, w) = streamdeck.resolution();
+        let img = (*self.image).clone();
+        let img = ImageBuilder::new(h, w)
+            .set_image(img)
+            .set_image_scale(self.scale)
+            .build();
+        streamdeck.write_img(img).await.unwrap();
         Ok(())
     }
 }
diff --git a/src/modules/space.rs b/src/modules/space.rs
index 640f7da..a55bea0 100644
--- a/src/modules/space.rs
+++ b/src/modules/space.rs
@@ -3,6 +3,7 @@ use super::ButtonConfigError;
 use super::ChannelReceiver;
 use super::DeviceAccess;
 use super::Module;
+use super::ModuleCache;
 use super::ModuleObject;
 use super::ReturnError;
 use crate::image_rendering::{create_error_image, ImageBuilder};
@@ -16,7 +17,10 @@ pub struct Space {
 
 #[async_trait]
 impl Module for Space {
-    async fn init(config: Arc<Button>) -> Result<ModuleObject, ButtonConfigError> {
+    async fn init(
+        config: Arc<Button>,
+        _cache: ModuleCache,
+    ) -> Result<ModuleObject, ButtonConfigError> {
         let name = config.parse_module("NAME", "Unknown".to_string()).res()?;
         Ok(Box::new(Space { name }))
     }