iris/core/src/render/texture.rs

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

use crate::{
    PatchRect, TextureUpdate, Textures, UiData,
    render::{
        TexturePrimitive,
        primitive::{ListDraw, PrimitiveRender},
    },
};

/// Draws standalone images, which it owns. Each is its own texture, so each
/// instance binds its own and is a draw of its own.
pub struct ImageRender {
    textures: GpuTextures,
    layout: BindGroupLayout,
    sampler: Sampler,
}

impl ImageRender {
    pub fn new(device: &Device, queue: &Queue) -> Self {
        Self {
            textures: GpuTextures::new(device, queue),
            layout: sampled_layout(device, TextureViewDimension::D2, "ui image"),
            sampler: default_sampler(device),
        }
    }
}

impl PrimitiveRender for ImageRender {
    fn layout(&self) -> Option<&BindGroupLayout> {
        Some(&self.layout)
    }

    fn update(&mut self, ui: &mut UiData) {
        self.textures
            .update(&mut ui.textures, &self.layout, &self.sampler);
    }

    fn instance_bindings(&self, list: &super::InstanceList, out: &mut Vec<u32>) {
        let slots = list
            .data()
            .chunks_exact(list.stride())
            .map(|data| bytemuck::pod_read_unaligned::<TexturePrimitive>(data).slot);
        out.extend(slots);
    }

    fn draw<'a>(&'a self, pass: &mut RenderPass<'a>, list: ListDraw<'a>) {
        for (i, &slot) in list.bindings.iter().enumerate() {
            let Some(image) = self.textures.group(slot) else {
                continue;
            };
            pass.set_bind_group(2, image, &[]);
            pass.draw(0..4, i as u32..i as u32 + 1);
        }
    }
}

/// The standalone images a ui draws, each its own texture and bind group --
/// unlike the glyph atlas in `super::page`, which is one array they share.
pub struct GpuTextures {
    device: Device,
    queue: Queue,
    slots: Vec<Option<ImageGpu>>,
}

struct ImageGpu {
    /// Kept for `patch`, which needs the texture rather than the view.
    texture: Texture,
    group: BindGroup,
}

impl GpuTextures {
    pub fn new(device: &Device, queue: &Queue) -> Self {
        Self {
            device: device.clone(),
            queue: queue.clone(),
            slots: Vec::new(),
        }
    }

    pub fn update(&mut self, textures: &mut Textures, layout: &BindGroupLayout, sampler: &Sampler) {
        for update in textures.updates() {
            match update {
                TextureUpdate::Push(image) => {
                    let image = self.create(image, layout, sampler);
                    self.slots.push(Some(image));
                }
                TextureUpdate::Set(i, image) => {
                    let image = self.create(image, layout, sampler);
                    self.slots[i as usize] = Some(image);
                }
                TextureUpdate::Patch(i, rect, image) => self.patch(i, rect, image),
                TextureUpdate::PushFree => self.slots.push(None),
                TextureUpdate::SetFree => {}
                TextureUpdate::Free(i) => self.slots[i as usize] = None,
            }
        }
    }

    pub fn group(&self, slot: u32) -> Option<&BindGroup> {
        self.slots.get(slot as usize)?.as_ref().map(|i| &i.group)
    }

    fn create(
        &self,
        image: &DynamicImage,
        layout: &BindGroupLayout,
        sampler: &Sampler,
    ) -> ImageGpu {
        let rgba = image.to_rgba8();
        let (width, height) = rgba.dimensions();
        let texture = self.device.create_texture_with_data(
            &self.queue,
            &TextureDescriptor {
                label: Some("image"),
                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,
            rgba.as_bytes(),
        );
        let view = texture.create_view(&TextureViewDescriptor::default());
        let group = sampled_group(&self.device, layout, &view, sampler, "ui image");
        ImageGpu { texture, group }
    }

    fn patch(&mut self, i: u32, rect: PatchRect, image: &DynamicImage) {
        let Some(Some(slot)) = self.slots.get(i as usize) else {
            return;
        };
        let dst = TexelCopyTextureInfo {
            texture: &slot.texture,
            mip_level: 0,
            origin: Origin3d {
                x: rect.x,
                y: rect.y,
                z: 0,
            },
            aspect: TextureAspect::All,
        };
        match image.as_rgba8() {
            Some(rgba) => write_region(&self.queue, dst, rgba, rect),
            // The texture is rgba8, so any other layout has to be converted --
            // and converting the rectangle is cheaper than the whole image.
            None => {
                let sub = image
                    .view(rect.x, rect.y, rect.width, rect.height)
                    .to_image();
                write_region(&self.queue, dst, &sub, PatchRect { x: 0, y: 0, ..rect });
            }
        }
    }
}

pub fn write_region(queue: &Queue, dst: TexelCopyTextureInfo, src: &RgbaImage, rect: PatchRect) {
    if rect.width == 0 || rect.height == 0 {
        return;
    }
    let stride = src.width() * 4;
    queue.write_texture(
        dst,
        src.as_bytes(),
        TexelCopyBufferLayout {
            offset: (rect.y * stride + rect.x * 4) as u64,
            bytes_per_row: Some(stride),
            rows_per_image: Some(rect.height),
        },
        Extent3d {
            width: rect.width,
            height: rect.height,
            depth_or_array_layers: 1,
        },
    );
}

/// What a primitive that samples binds: a texture, and the sampler that reads
/// it.
pub fn sampled_group(
    device: &Device,
    layout: &BindGroupLayout,
    view: &TextureView,
    sampler: &Sampler,
    label: &'static str,
) -> BindGroup {
    device.create_bind_group(&BindGroupDescriptor {
        layout,
        entries: &[
            BindGroupEntry {
                binding: 0,
                resource: BindingResource::TextureView(view),
            },
            BindGroupEntry {
                binding: 1,
                resource: BindingResource::Sampler(sampler),
            },
        ],
        label: Some(label),
    })
}

/// The layout for one of those. The dimension differs -- the atlas is an
/// array of pages and an image is not -- and nothing else does.
pub fn sampled_layout(
    device: &Device,
    dimension: TextureViewDimension,
    label: &'static str,
) -> BindGroupLayout {
    device.create_bind_group_layout(&BindGroupLayoutDescriptor {
        entries: &[
            BindGroupLayoutEntry {
                binding: 0,
                visibility: ShaderStages::FRAGMENT,
                ty: BindingType::Texture {
                    sample_type: TextureSampleType::Float { filterable: false },
                    view_dimension: dimension,
                    multisampled: false,
                },
                count: None,
            },
            BindGroupLayoutEntry {
                binding: 1,
                visibility: ShaderStages::FRAGMENT,
                ty: BindingType::Sampler(SamplerBindingType::NonFiltering),
                count: None,
            },
        ],
        label: Some(label),
    })
}

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