lemur/src/images.rs

use std::{
    collections::HashMap,
    ops::Deref,
    sync::{Arc, Mutex, Weak},
    thread,
    time::Duration,
};

use egui::{
    Color32, ColorImage, TextureHandle, TextureOptions,
    epaint::ImageDelta,
    load::{LoadError, SizedTexture, TextureLoader, TexturePoll},
};
use tokio::{sync::mpsc, time::timeout};

pub struct ImageProcessor {
    sender: mpsc::UnboundedSender<Box<dyn ImageRendererImpl>>,
}

impl ImageProcessor {
    pub fn new() -> Self {
        let (sender, receiver) = mpsc::unbounded_channel();
        thread::spawn(move || {
            tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .unwrap()
                .block_on(async move {
                    let mut worker = ImageProcessorWorker {
                        receiver,
                        renderers: vec![],
                    };
                    worker.run().await
                })
        });
        Self { sender }
    }

    pub fn add<const N: usize, R: ImageRenderer<N> + 'static>(
        &self,
        renderer: R,
        params: R::Params,
    ) -> ([ImageHandle; N], 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 sink = Arc::new(Mutex::new(params.clone()));
        let _ = self.sender.send(Box::new(ImageRendererWrapper {
            renderer,
            params: Arc::downgrade(&sink),
            images,
            states,
        }));
        let params = ImageParams {
            value: params,
            sink,
        };
        (handles, params)
    }
}

struct ImageProcessorWorker {
    receiver: mpsc::UnboundedReceiver<Box<dyn ImageRendererImpl>>,
    renderers: Vec<Box<dyn ImageRendererImpl>>,
}

impl ImageProcessorWorker {
    async fn run(&mut self) {
        loop {
            if self.renderers.is_empty() {
                // if we have nothing to do, block until we have something to do
                if self.receiver.recv_many(&mut self.renderers, 64).await == 0 {
                    // shutdown
                    return;
                }
                while let Ok(renderer) = self.receiver.try_recv() {
                    self.renderers.push(renderer);
                }
            }
            self.renderers
                .retain_mut(|renderer| renderer.try_update().is_ok());
            // wait up to 10 ms for more renderers
            if timeout(
                Duration::from_millis(10),
                self.receiver.recv_many(&mut self.renderers, 64),
            )
            .await
            .is_ok()
            {
                while let Ok(renderer) = self.receiver.try_recv() {
                    self.renderers.push(renderer);
                }
            }
        }
    }
}

pub struct ImageBuffer {
    pub size: [usize; 2],
    pub pixels: Vec<Color32>,
}

impl ImageBuffer {
    pub fn new(width: usize, height: usize) -> Self {
        Self {
            size: [width, height],
            pixels: vec![Color32::BLACK; width * height],
        }
    }

    pub fn clear(&mut self) {
        for pixel in self.pixels.iter_mut() {
            *pixel = Color32::BLACK;
        }
    }

    pub fn write(&mut self, coords: (usize, usize), pixel: Color32) {
        self.pixels[coords.1 * self.size[0] + coords.0] = pixel;
    }

    pub fn add(&mut self, coords: (usize, usize), pixel: Color32) {
        let index = coords.1 * self.size[0] + coords.0;
        let old = self.pixels[index];
        self.pixels[index] = Color32::from_rgb(
            old.r() + pixel.r(),
            old.g() + pixel.g(),
            old.b() + pixel.b(),
        );
    }

    pub fn changed(&self, image: &ColorImage) -> bool {
        image.pixels.iter().zip(&self.pixels).any(|(a, b)| a != b)
    }

    pub fn read(&self, image: &mut ColorImage) {
        image.pixels.copy_from_slice(&self.pixels);
    }
}

#[derive(Clone)]
pub struct ImageHandle {
    size: [f32; 2],
    data: Arc<Mutex<Option<Arc<ColorImage>>>>,
}

impl ImageHandle {
    fn pull(&mut self) -> Option<Arc<ColorImage>> {
        self.data.lock().unwrap().take()
    }
}

pub struct ImageParams<T> {
    value: T,
    sink: Arc<Mutex<T>>,
}

impl<T> Deref for ImageParams<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

impl<T: Clone + Eq> ImageParams<T> {
    pub fn write(&mut self, value: T) {
        if self.value != value {
            self.value = value.clone();
            *self.sink.lock().unwrap() = value;
        }
    }
}

pub trait ImageRenderer<const N: usize>: Send {
    type Params: Clone + Send;
    fn sizes(&self) -> [[usize; 2]; N];
    fn render(&mut self, params: &Self::Params, images: &mut [ImageBuffer; N]);
}

#[derive(Clone)]
struct ImageState {
    size: [usize; 2],
    buffers: [Arc<ColorImage>; 2],
    last_buffer: usize,
    sink: Arc<Mutex<Option<Arc<ColorImage>>>>,
}

impl ImageState {
    fn new(size: [usize; 2]) -> Self {
        let buffers = [
            Arc::new(ColorImage::new(size, Color32::BLACK)),
            Arc::new(ColorImage::new(size, Color32::BLACK)),
        ];
        let sink = buffers[0].clone();
        Self {
            size,
            buffers,
            last_buffer: 0,
            sink: Arc::new(Mutex::new(Some(sink))),
        }
    }
    fn try_send_update(&mut self, image: &ImageBuffer) {
        let last = &self.buffers[self.last_buffer];
        if !image.changed(last) {
            return;
        }

        let next_buffer = (self.last_buffer + 1) % self.buffers.len();
        let next = &mut self.buffers[next_buffer];
        image.read(Arc::make_mut(next));
        self.last_buffer = next_buffer;
        self.sink.lock().unwrap().replace(next.clone());
    }
}

struct ImageRendererWrapper<const N: usize, R: ImageRenderer<N>> {
    renderer: R,
    params: Weak<Mutex<R::Params>>,
    images: [ImageBuffer; N],
    states: [ImageState; N],
}

trait ImageRendererImpl: Send {
    fn try_update(&mut self) -> Result<(), ()>;
}

impl<const N: usize, R: ImageRenderer<N> + Send> ImageRendererImpl for ImageRendererWrapper<N, R> {
    fn try_update(&mut self) -> Result<(), ()> {
        let params = match self.params.upgrade() {
            Some(params) => params.lock().unwrap().clone(),
            None => {
                // the UI isn't using this anymore
                return Err(());
            }
        };
        self.renderer.render(&params, &mut self.images);

        for (state, image) in self.states.iter_mut().zip(&self.images) {
            state.try_send_update(image);
        }
        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()
    }
}