iris/core/src/render/texture.rs

use image::{DynamicImage, EncodableLayout, GenericImageView};
use wgpu::{util::DeviceExt, *};

use crate::{PatchRect, TextureUpdate, Textures};

pub struct GpuTextures {
    device: Device,
    queue: Queue,
    /// Parallel to `views`; patches require textures rather than views.
    textures: Vec<Option<Texture>>,
    views: Vec<TextureView>,
    view_count: usize,
    samplers: Vec<Sampler>,
    null_view: TextureView,
    no_views: Vec<TextureView>,
}

impl GpuTextures {
    pub fn update(&mut self, textures: &mut Textures) -> bool {
        let mut bindings_changed = false;
        for update in textures.updates() {
            bindings_changed |= match update {
                TextureUpdate::Push(image) => {
                    self.push(image);
                    true
                }
                TextureUpdate::Set(i, image) => {
                    self.set(i, image);
                    true
                }
                TextureUpdate::Patch(i, rect, image) => {
                    self.patch(i, rect, image);
                    false
                }
                TextureUpdate::SetFree => {
                    self.view_count += 1;
                    true
                }
                TextureUpdate::Free(i) => {
                    self.free(i);
                    true
                }
                TextureUpdate::PushFree => {
                    self.push_free();
                    true
                }
            };
        }
        bindings_changed
    }
    fn set(&mut self, i: u32, image: &DynamicImage) {
        self.view_count += 1;
        let (texture, view) = self.create(image);
        self.textures[i as usize] = Some(texture);
        self.views[i as usize] = view;
    }
    fn free(&mut self, i: u32) {
        self.view_count -= 1;
        self.textures[i as usize] = None;
        self.views[i as usize] = self.null_view.clone();
    }
    fn push(&mut self, image: &DynamicImage) {
        self.view_count += 1;
        let (texture, view) = self.create(image);
        self.textures.push(Some(texture));
        self.views.push(view);
    }
    fn push_free(&mut self) {
        self.view_count += 1;
        self.textures.push(None);
        self.views.push(self.null_view.clone());
    }

    fn patch(&mut self, i: u32, rect: PatchRect, image: &DynamicImage) {
        let Some(texture) = &self.textures[i as usize] else {
            return;
        };
        if rect.width == 0 || rect.height == 0 {
            return;
        }
        // `write_texture` requires tightly packed rows, unlike the atlas image.
        let sub = image
            .view(rect.x, rect.y, rect.width, rect.height)
            .to_image();
        self.queue.write_texture(
            TexelCopyTextureInfo {
                texture,
                mip_level: 0,
                origin: Origin3d {
                    x: rect.x,
                    y: rect.y,
                    z: 0,
                },
                aspect: TextureAspect::All,
            },
            sub.as_bytes(),
            TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(rect.width * 4),
                rows_per_image: Some(rect.height),
            },
            Extent3d {
                width: rect.width,
                height: rect.height,
                depth_or_array_layers: 1,
            },
        );
    }

    fn create(&self, image: &DynamicImage) -> (Texture, TextureView) {
        let image = image.to_rgba8();
        let (width, height) = image.dimensions();
        let texture = self.device.create_texture_with_data(
            &self.queue,
            &TextureDescriptor {
                label: None,
                size: Extent3d {
                    width,
                    height,
                    depth_or_array_layers: 1,
                },
                mip_level_count: 1,
                sample_count: 1,
                dimension: TextureDimension::D2,
                format: TextureFormat::Rgba8Unorm,
                usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
                view_formats: &[],
            },
            wgt::TextureDataOrder::MipMajor,
            image.as_bytes(),
        );
        let view = texture.create_view(&TextureViewDescriptor::default());
        (texture, view)
    }

    pub fn new(device: &Device, queue: &Queue) -> Self {
        let null_view = null_texture_view(device);
        Self {
            device: device.clone(),
            queue: queue.clone(),
            textures: Vec::new(),
            views: Vec::new(),
            samplers: vec![default_sampler(device)],
            no_views: vec![null_view.clone()],
            null_view,
            view_count: 0,
        }
    }

    pub fn views(&self) -> Vec<&TextureView> {
        if self.views.is_empty() {
            &self.no_views
        } else {
            &self.views
        }
        .iter()
        .by_ref()
        .collect()
    }

    pub fn samplers(&self) -> Vec<&Sampler> {
        self.samplers.iter().by_ref().collect()
    }

    pub fn view_count(&self) -> usize {
        self.view_count
    }
}

pub fn null_texture_view(device: &Device) -> TextureView {
    device
        .create_texture(&TextureDescriptor {
            label: Some("null"),
            size: Extent3d {
                width: 1,
                height: 1,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: TextureDimension::D2,
            format: TextureFormat::Rgba8Unorm,
            usage: TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        })
        .create_view(&TextureViewDescriptor::default())
}

pub fn default_sampler(device: &Device) -> Sampler {
    device.create_sampler(&SamplerDescriptor::default())
}