lemur/src/vram.rs

use std::{
    collections::{hash_map::Entry, HashMap},
    hash::Hash,
    sync::{Arc, Mutex},
};

use egui::{
    epaint::ImageDelta,
    load::{LoadError, SizedTexture, TextureLoader, TexturePoll},
    Color32, ColorImage, Context, TextureHandle, TextureOptions,
};
use serde::{Deserialize, Serialize};

pub trait VramResource: Sized + PartialEq + Eq + Hash {
    fn to_uri(&self) -> String;
    fn from_uri(uri: &str) -> Option<Self>;
}

impl<T: Serialize + for<'a> Deserialize<'a> + PartialEq + Eq + Hash> VramResource for T {
    fn to_uri(&self) -> String {
        format!("vram://{}", serde_json::to_string(self).unwrap())
    }

    fn from_uri(uri: &str) -> Option<Self> {
        let content = uri.strip_prefix("vram://")?;
        serde_json::from_str(content).ok()
    }
}

pub enum VramImage {
    Unchanged(Arc<ColorImage>),
    Changed(ColorImage),
}

impl VramImage {
    pub fn new(width: usize, height: usize) -> Self {
        Self::Changed(ColorImage::new([width, height], Color32::BLACK))
    }

    pub fn write(&mut self, coords: (usize, usize), shade: u8) {
        match self {
            Self::Unchanged(image) => {
                let value = image[coords];
                if value.r() == shade {
                    return;
                };
                let mut new_image = ColorImage::clone(image);
                new_image[coords] = Color32::from_gray(shade);
                *self = Self::Changed(new_image);
            }
            Self::Changed(image) => {
                image[coords] = Color32::from_gray(shade);
            }
        }
    }

    pub fn take(&mut self) -> Option<Arc<ColorImage>> {
        match self {
            Self::Unchanged(_) => None,
            Self::Changed(image) => {
                let arced = Arc::new(std::mem::take(image));
                *self = Self::Unchanged(arced.clone());
                Some(arced)
            }
        }
    }
}

pub trait VramImageLoader {
    type Resource: VramResource;

    fn id(&self) -> &str;
    fn add(&self, resource: &Self::Resource) -> Option<VramImage>;
    fn update<'a>(
        &'a self,
        resources: impl Iterator<Item = (&'a Self::Resource, &'a mut VramImage)>,
    );
}

pub struct VramTextureLoader<T: VramImageLoader> {
    id: String,
    loader: Mutex<T>,
    cache: Mutex<HashMap<T::Resource, (TextureHandle, VramImage)>>,
}

impl<T: VramImageLoader> VramTextureLoader<T> {
    pub fn new(loader: T) -> Self {
        Self {
            id: loader.id().to_string(),
            loader: Mutex::new(loader),
            cache: Mutex::new(HashMap::new()),
        }
    }
}

impl<T: VramImageLoader> TextureLoader for VramTextureLoader<T> {
    fn id(&self) -> &str {
        &self.id
    }

    fn load(
        &self,
        ctx: &Context,
        uri: &str,
        texture_options: TextureOptions,
        _size_hint: egui::SizeHint,
    ) -> Result<TexturePoll, LoadError> {
        let Some(resource) = T::Resource::from_uri(uri) else {
            return Err(LoadError::NotSupported);
        };
        if texture_options != TextureOptions::NEAREST {
            return Err(LoadError::Loading(
                "Only TextureOptions::NEAREST are supported".into(),
            ));
        }
        let loader = self.loader.lock().unwrap();
        let mut cache = self.cache.lock().unwrap();
        let resources = cache.iter_mut().map(|(k, v)| (k, &mut v.1));
        loader.update(resources);
        for (handle, image) in cache.values_mut() {
            if let Some(data) = image.take() {
                let delta = ImageDelta::full(data, TextureOptions::NEAREST);
                ctx.tex_manager().write().set(handle.id(), delta);
            }
        }
        match cache.entry(resource) {
            Entry::Occupied(entry) => {
                let texture = SizedTexture::from_handle(&entry.get().0);
                Ok(TexturePoll::Ready { texture })
            }
            Entry::Vacant(entry) => {
                let Some(mut image) = loader.add(entry.key()) else {
                    return Err(LoadError::Loading("could not load texture".into()));
                };
                let Some(data) = image.take() else {
                    return Err(LoadError::Loading("could not load texture".into()));
                };
                let handle = ctx.load_texture(uri, data, TextureOptions::NEAREST);
                let (handle, _) = entry.insert((handle, image));
                let texture = SizedTexture::from_handle(handle);
                Ok(TexturePoll::Ready { texture })
            }
        }
    }

    fn forget(&self, uri: &str) {
        if let Some(resource) = T::Resource::from_uri(uri) {
            self.cache.lock().unwrap().remove(&resource);
        }
    }

    fn forget_all(&self) {
        self.cache.lock().unwrap().clear();
    }

    fn byte_size(&self) -> usize {
        self.cache
            .lock()
            .unwrap()
            .values()
            .map(|h| h.0.byte_size())
            .sum()
    }
}