Rewrite VIP tools to not mutate context

src/images.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::HashMap,
+    collections::{HashMap, hash_map::Entry},
     ops::Deref,
     sync::{Arc, Mutex, Weak},
     thread,
@@ -7,17 +7,21 @@ use std::{
 };
 
 use egui::{
-    Color32, ColorImage, TextureHandle, TextureOptions,
+    Color32, ColorImage, TextureHandle,
     epaint::ImageDelta,
+    generate_loader_id,
     load::{LoadError, SizedTexture, TextureLoader, TexturePoll},
 };
 use tokio::{sync::mpsc, time::timeout};
 
-pub struct ImageProcessor {
+use crate::emulator::SimId;
+
+pub struct ImageTextureLoader {
+    cache: Mutex<HashMap<String, ImageEntry>>,
     sender: mpsc::UnboundedSender<Box<dyn ImageRendererImpl>>,
 }
 
-impl ImageProcessor {
+impl ImageTextureLoader {
     pub fn new() -> Self {
         let (sender, receiver) = mpsc::unbounded_channel();
         thread::spawn(move || {
@@ -26,29 +30,52 @@ impl ImageProcessor {
                 .build()
                 .unwrap()
                 .block_on(async move {
-                    let mut worker = ImageProcessorWorker {
+                    let mut worker = ImageTextureLoaderWorker {
                         receiver,
                         renderers: vec![],
                     };
                     worker.run().await
                 })
         });
-        Self { sender }
+        Self {
+            cache: Mutex::new(HashMap::new()),
+            sender,
+        }
     }
 
     pub fn add<const N: usize, R: ImageRenderer<N> + 'static>(
         &self,
+        sim_id: SimId,
         renderer: R,
         params: R::Params,
-    ) -> ([ImageHandle; N], ImageParams<R::Params>) {
+    ) -> ImageParams<R::Params> {
         let states = renderer.sizes().map(ImageState::new);
-        let handles = states.clone().map(|state| ImageHandle {
-            size: state.size.map(|i| i as f32),
-            data: state.sink,
-        });
         let images = renderer
             .sizes()
             .map(|[width, height]| ImageBuffer::new(width, height));
+        {
+            let mut cache = self.cache.lock().unwrap();
+            for (name, state) in renderer.names().into_iter().zip(&states) {
+                let url = format!("vip://{sim_id}/{name}");
+                let entry = ImageEntry {
+                    size: state.size.map(|i| i as f32),
+                    data: Arc::downgrade(&state.sink),
+                    texture: None,
+                };
+                match cache.entry(url) {
+                    Entry::Vacant(e) => {
+                        e.insert(entry);
+                    }
+                    Entry::Occupied(mut e) => {
+                        // Only overwrite an old entry if that entry is dead.
+                        // Otherwise, we clear out the reference to a texture actually in use.
+                        if e.get().data.strong_count() == 0 {
+                            e.insert(entry);
+                        }
+                    }
+                }
+            }
+        }
         let sink = Arc::new(Mutex::new(params.clone()));
         let _ = self.sender.send(Box::new(ImageRendererWrapper {
             renderer,
@@ -56,20 +83,87 @@ impl ImageProcessor {
             images,
             states,
         }));
-        let params = ImageParams {
+        ImageParams {
             value: params,
             sink,
-        };
-        (handles, params)
+        }
     }
 }
 
-struct ImageProcessorWorker {
+impl TextureLoader for ImageTextureLoader {
+    fn id(&self) -> &str {
+        generate_loader_id!(ImageTextureLoader2)
+    }
+
+    fn load(
+        &self,
+        ctx: &egui::Context,
+        uri: &str,
+        texture_options: egui::TextureOptions,
+        _size_hint: egui::SizeHint,
+    ) -> Result<TexturePoll, LoadError> {
+        let mut cache = self.cache.lock().unwrap();
+        let Some(entry) = cache.get_mut(uri) else {
+            return Err(LoadError::NotSupported);
+        };
+        if texture_options != egui::TextureOptions::NEAREST {
+            return Err(LoadError::Loading(
+                "Only TextureOptions::NEAREST are supported".into(),
+            ));
+        }
+        let Some(data) = entry.data.upgrade() else {
+            cache.remove(uri);
+            return Err(LoadError::Loading("Backing loader no longer exists".into()));
+        };
+        match (data.lock().unwrap().take(), entry.texture.as_ref()) {
+            (Some(update), Some(handle)) => {
+                let delta = ImageDelta::full(update, texture_options);
+                ctx.tex_manager().write().set(handle.id(), delta);
+                let texture = SizedTexture::new(handle, entry.size);
+                Ok(TexturePoll::Ready { texture })
+            }
+            (Some(update), None) => {
+                let handle = ctx.load_texture(uri, update, texture_options);
+                let texture = SizedTexture::new(&handle, entry.size);
+                entry.texture.replace(handle);
+                Ok(TexturePoll::Ready { texture })
+            }
+            (None, Some(handle)) => {
+                let texture = SizedTexture::new(handle, entry.size);
+                Ok(TexturePoll::Ready { texture })
+            }
+            (None, None) => {
+                let size = entry.size.into();
+                Ok(TexturePoll::Pending { size: Some(size) })
+            }
+        }
+    }
+
+    fn forget(&self, uri: &str) {
+        let _ = uri;
+    }
+
+    fn forget_all(&self) {}
+
+    fn byte_size(&self) -> usize {
+        self.cache
+            .lock()
+            .unwrap()
+            .values()
+            .map(|entry| {
+                let [width, height] = entry.size;
+                width as usize * height as usize * 4
+            })
+            .sum()
+    }
+}
+
+struct ImageTextureLoaderWorker {
     receiver: mpsc::UnboundedReceiver<Box<dyn ImageRendererImpl>>,
     renderers: Vec<Box<dyn ImageRendererImpl>>,
 }
 
-impl ImageProcessorWorker {
+impl ImageTextureLoaderWorker {
     async fn run(&mut self) {
         loop {
             if self.renderers.is_empty() {
@@ -142,16 +236,10 @@ impl ImageBuffer {
     }
 }
 
-#[derive(Clone)]
-pub struct ImageHandle {
+struct ImageEntry {
     size: [f32; 2],
-    data: Arc<Mutex<Option<Arc<ColorImage>>>>,
-}
-
-impl ImageHandle {
-    fn pull(&mut self) -> Option<Arc<ColorImage>> {
-        self.data.lock().unwrap().take()
-    }
+    data: Weak<Mutex<Option<Arc<ColorImage>>>>,
+    texture: Option<TextureHandle>,
 }
 
 pub struct ImageParams<T> {
@@ -178,6 +266,7 @@ impl<T: Clone + Eq> ImageParams<T> {
 
 pub trait ImageRenderer<const N: usize>: Send {
     type Params: Clone + Send;
+    fn names(&self) -> [&str; N];
     fn sizes(&self) -> [[usize; 2]; N];
     fn render(&mut self, params: &Self::Params, images: &mut [ImageBuffer; N]);
 }
@@ -246,83 +335,3 @@ impl<const N: usize, R: ImageRenderer<N> + Send> ImageRendererImpl for ImageRend
         Ok(())
     }
 }
-
-pub struct ImageTextureLoader {
-    cache: Mutex<HashMap<String, (ImageHandle, Option<TextureHandle>)>>,
-}
-
-impl ImageTextureLoader {
-    pub fn new(renderers: impl IntoIterator<Item = (String, ImageHandle)>) -> Self {
-        let mut cache = HashMap::new();
-        for (key, image) in renderers {
-            cache.insert(key, (image, None));
-        }
-        Self {
-            cache: Mutex::new(cache),
-        }
-    }
-}
-
-impl TextureLoader for ImageTextureLoader {
-    fn id(&self) -> &str {
-        concat!(module_path!(), "ImageTextureLoader")
-    }
-
-    fn load(
-        &self,
-        ctx: &egui::Context,
-        uri: &str,
-        texture_options: TextureOptions,
-        _size_hint: egui::SizeHint,
-    ) -> Result<TexturePoll, LoadError> {
-        let mut cache = self.cache.lock().unwrap();
-        let Some((image, maybe_handle)) = cache.get_mut(uri) else {
-            return Err(LoadError::NotSupported);
-        };
-        if texture_options != TextureOptions::NEAREST {
-            return Err(LoadError::Loading(
-                "Only TextureOptions::NEAREST are supported".into(),
-            ));
-        }
-        match (image.pull(), maybe_handle.as_ref()) {
-            (Some(update), Some(handle)) => {
-                let delta = ImageDelta::full(update, texture_options);
-                ctx.tex_manager().write().set(handle.id(), delta);
-                let texture = SizedTexture::new(handle, image.size);
-                Ok(TexturePoll::Ready { texture })
-            }
-            (Some(update), None) => {
-                let handle = ctx.load_texture(uri, update, texture_options);
-                let texture = SizedTexture::new(&handle, image.size);
-                maybe_handle.replace(handle);
-                Ok(TexturePoll::Ready { texture })
-            }
-            (None, Some(handle)) => {
-                let texture = SizedTexture::new(handle, image.size);
-                Ok(TexturePoll::Ready { texture })
-            }
-            (None, None) => {
-                let size = image.size.into();
-                Ok(TexturePoll::Pending { size: Some(size) })
-            }
-        }
-    }
-
-    fn forget(&self, uri: &str) {
-        let _ = uri;
-    }
-
-    fn forget_all(&self) {}
-
-    fn byte_size(&self) -> usize {
-        self.cache
-            .lock()
-            .unwrap()
-            .values()
-            .map(|(image, _)| {
-                let [width, height] = image.size;
-                width as usize * height as usize * 4
-            })
-            .sum()
-    }
-}