iris/core/src/primitive/texture.rs

use crate::util::{RefCounter, Vec2};
use image::{DynamicImage, GenericImageView};
use std::{
    ops::Index,
    sync::mpsc::{Receiver, Sender, channel},
};

#[derive(Debug, Clone)]
pub struct TextureHandle {
    slot: u32,
    size: Vec2,
    counter: RefCounter,
    send: Sender<u32>,
}

/// a texture manager for a ui
/// note that this is heavily oriented towards wgpu's renderer so the primitives don't need mapped
pub struct Textures {
    free: Vec<u32>,
    images: Vec<Option<DynamicImage>>,
    updates: Vec<Update>,
    send: Sender<u32>,
    recv: Receiver<u32>,
}

pub enum TextureUpdate<'a> {
    Push(&'a DynamicImage),
    Set(u32, &'a DynamicImage),
    Patch(u32, PatchRect, &'a DynamicImage),
    Free(u32),
    /// Added and freed before the renderer drained either update. It still has
    /// to push a slot to stay lined up with `images`; `Free` then empties it.
    PushFree,
    SetFree,
}

#[derive(Debug, Clone, Copy)]
pub struct PatchRect {
    pub x: u32,
    pub y: u32,
    pub width: u32,
    pub height: u32,
}

enum Update {
    Push(u32),
    Set(u32),
    Patch(u32, PatchRect),
    Free(u32),
}

impl Textures {
    pub fn new() -> Self {
        let (send, recv) = channel();
        Self {
            free: Vec::new(),
            images: Vec::new(),
            updates: Vec::new(),
            send,
            recv,
        }
    }
    pub fn add(&mut self, image: impl Into<DynamicImage>) -> TextureHandle {
        let image = image.into();
        let size = image.dimensions().into();
        TextureHandle {
            slot: self.push(image),
            size,
            counter: RefCounter::new(),
            send: self.send.clone(),
        }
    }

    fn push(&mut self, image: DynamicImage) -> u32 {
        if let Some(i) = self.free.pop() {
            self.images[i as usize] = Some(image);
            self.updates.push(Update::Set(i));
            i
        } else {
            let i = self.images.len() as u32;
            self.images.push(Some(image));
            self.updates.push(Update::Push(i));
            i
        }
    }

    pub fn image_mut(&mut self, handle: &TextureHandle) -> &mut DynamicImage {
        self.images[handle.slot as usize]
            .as_mut()
            .expect("texture was freed while still held")
    }

    /// Queue an upload of just `rect`, after writing it with `image_mut`.
    pub fn patch(&mut self, handle: &TextureHandle, rect: PatchRect) {
        self.updates.push(Update::Patch(handle.slot, rect));
    }

    /// How many textures are live, which is what a ui can ask; the renderer's
    /// copies follow from the updates it drains.
    pub fn count(&self) -> usize {
        self.images.iter().flatten().count()
    }

    pub fn free(&mut self) {
        for idx in self.recv.try_iter() {
            self.images[idx as usize] = None;
            self.updates.push(Update::Free(idx));
            self.free.push(idx);
        }
    }

    pub fn updates(&mut self) -> impl Iterator<Item = TextureUpdate<'_>> {
        self.updates.drain(..).map(|u| match u {
            Update::Push(i) => self.images[i as usize]
                .as_ref()
                .map(TextureUpdate::Push)
                .unwrap_or(TextureUpdate::PushFree),
            Update::Set(i) => self.images[i as usize]
                .as_ref()
                .map(|img| TextureUpdate::Set(i, img))
                .unwrap_or(TextureUpdate::SetFree),
            Update::Patch(i, rect) => self.images[i as usize]
                .as_ref()
                .map(|img| TextureUpdate::Patch(i, rect, img))
                .unwrap_or(TextureUpdate::SetFree),
            Update::Free(i) => TextureUpdate::Free(i),
        })
    }
}

impl TextureHandle {
    /// Index into `Textures`, and into the renderer's parallel slots.
    pub fn slot(&self) -> u32 {
        self.slot
    }
    pub fn size(&self) -> Vec2 {
        self.size
    }
}

impl Drop for TextureHandle {
    fn drop(&mut self) {
        if self.counter.drop() {
            let _ = self.send.send(self.slot);
        }
    }
}

impl Index<&TextureHandle> for Textures {
    type Output = DynamicImage;

    fn index(&self, index: &TextureHandle) -> &Self::Output {
        self.images[index.slot as usize].as_ref().unwrap()
    }
}

impl Default for Textures {
    fn default() -> Self {
        Self::new()
    }
}