iris/core/src/render/primitive.rs

use std::{any::TypeId, marker::PhantomData};

use crate::{
    Color, TextureHandle, UiData, UiRegion, WidgetId,
    render::{
        data::{MaskIdx, PrimitiveInstance},
        page::GlyphRender,
        texture::ImageRender,
    },
    util::{HashMap, Vec2},
};
use bytemuck::Pod;
use wgpu::{BindGroupLayout, Device, Queue, RenderPass};

/// One instance of a primitive, laid out as the struct its shader reads.
///
/// The type carries its own shader, so drawing one is all the wiring it needs:
/// its list, free list, buffers and pipeline follow from being registered.
pub trait Primitive: Pod + 'static {
    /// Compiled after `prelude.wgsl`, which states what it declares and what
    /// it is given.
    const WGSL: &'static str;

    /// Made once, the first time the renderer sees this primitive. It owns
    /// whatever the shader samples and records the primitive's own draws; the
    /// default owns nothing and draws every instance in one call.
    fn render(device: &Device, queue: &Queue) -> Box<dyn PrimitiveRender>
    where
        Self: Sized,
    {
        let _ = (device, queue);
        Box::new(Instanced)
    }
}

/// The renderer's half of a primitive: what it samples, what it uploads, and
/// what draws it records.
///
/// Everything a draw shares -- the pipeline, the window and masks, the list's
/// own data and instance buffer -- is set before this is called. What is left
/// is what only this primitive knows: its group 2, and how many draws its
/// instances are.
pub trait PrimitiveRender {
    /// The layout its shader reads at group 2. `None` for a primitive whose
    /// shader samples nothing, whose pipeline then has no group 2 at all.
    fn layout(&self) -> Option<&BindGroupLayout> {
        None
    }

    /// Uploads whatever this primitive owns, once a frame, before any draw.
    fn update(&mut self, ui: &mut UiData) {
        let _ = ui;
    }

    /// Keeps what the primitive needs per instance at draw time, read from
    /// the list's own data. A primitive that binds nothing per instance --
    /// most of them -- leaves this empty and draws in one call.
    fn instance_bindings(&self, list: &InstanceList, out: &mut Vec<u32>) {
        let _ = (list, out);
    }

    fn draw<'a>(&'a self, pass: &mut RenderPass<'a>, list: ListDraw<'a>);
}

/// What a `PrimitiveRender` draws: this list's instances, and whatever
/// `instance_bindings` kept for them.
pub struct ListDraw<'a> {
    pub instances: u32,
    pub bindings: &'a [u32],
}

/// The default: nothing sampled, every instance in one call.
pub struct Instanced;

impl PrimitiveRender for Instanced {
    fn draw<'a>(&'a self, pass: &mut RenderPass<'a>, list: ListDraw<'a>) {
        pass.draw(0..4, 0..list.instances);
    }
}

/// Which registered primitive an instance is.
pub struct PrimitiveKind<P> {
    id: u32,
    _p: PhantomData<fn(P)>,
}

impl<P> PrimitiveKind<P> {
    fn new(id: u32) -> Self {
        Self {
            id,
            _p: PhantomData,
        }
    }
}

impl<P> Clone for PrimitiveKind<P> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<P> Copy for PrimitiveKind<P> {}

/// Every primitive a ui can draw, in the order they were first drawn.
#[derive(Default)]
pub struct PrimitiveRegistry {
    kinds: Vec<PrimitiveSource>,
    ids: HashMap<TypeId, u32>,
}

pub struct PrimitiveSource {
    pub wgsl: &'static str,
    pub label: &'static str,
    /// Size of one instance's entry, stated as the data binding's minimum.
    pub stride: u64,
    pub render: fn(&Device, &Queue) -> Box<dyn PrimitiveRender>,
}

impl PrimitiveRegistry {
    /// Registers `P` if this is the first time it has been drawn.
    pub fn kind<P: Primitive>(&mut self) -> PrimitiveKind<P> {
        let Self { kinds, ids } = self;
        let id = *ids.entry(TypeId::of::<P>()).or_insert_with(|| {
            kinds.push(PrimitiveSource {
                wgsl: P::WGSL,
                label: std::any::type_name::<P>(),
                stride: size_of::<P>() as u64,
                render: P::render,
            });
            kinds.len() as u32 - 1
        });
        PrimitiveKind::new(id)
    }

    pub fn sources(&self) -> &[PrimitiveSource] {
        &self.kinds
    }
}

/// One registered primitive's instances in one layer. Everything per-instance
/// rides here, so it stays in step through a `swap_remove`.
pub struct InstanceList {
    instances: Vec<PrimitiveInstance>,
    /// The widget each instance belongs to, for renumbering its handles.
    assoc: Vec<WidgetId>,
    free: Vec<usize>,
    /// `stride` bytes of the primitive's own data per instance.
    data: Vec<u8>,
    /// From the type the list was made for, so a write is never checked.
    stride: usize,
}

impl InstanceList {
    fn new<P: Primitive>() -> Self {
        Self {
            instances: Vec::new(),
            assoc: Vec::new(),
            free: Vec::new(),
            data: Vec::new(),
            stride: size_of::<P>(),
        }
    }

    pub fn instances(&self) -> &[PrimitiveInstance] {
        &self.instances
    }

    pub fn data(&self) -> &[u8] {
        &self.data
    }

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

    fn push(&mut self, id: WidgetId, inst: PrimitiveInstance, data: &[u8]) -> usize {
        if let Some(i) = self.free.pop() {
            self.instances[i] = inst;
            self.assoc[i] = id;
            self.data[i * self.stride..][..self.stride].copy_from_slice(data);
            i
        } else {
            let i = self.instances.len();
            self.instances.push(inst);
            self.assoc.push(id);
            self.data.extend_from_slice(data);
            i
        }
    }

    fn free(&mut self, i: usize) -> MaskIdx {
        self.free.push(i);
        self.instances[i].mask_idx
    }

    fn apply_free(&mut self, kind: u32) -> impl Iterator<Item = PrimitiveChange> {
        self.free.sort_by(|a, b| b.cmp(a));
        let instances = &mut self.instances;
        let assoc = &mut self.assoc;
        let data = &mut self.data;
        let stride = self.stride;
        self.free.drain(..).filter_map(move |i| {
            instances.swap_remove(i);
            assoc.swap_remove(i);
            let last = instances.len();
            data.copy_within(last * stride..(last + 1) * stride, i * stride);
            data.truncate(last * stride);
            if i == last {
                return None;
            }
            let id = assoc[i];
            Some(PrimitiveChange {
                id,
                kind,
                old: last,
                new: i,
            })
        })
    }
}

/// Everything one layer draws, one list per registered primitive.
pub struct LayerDraws {
    /// `None` until this layer draws that primitive, because only the write
    /// knows the type the list is for.
    primitives: Vec<Option<InstanceList>>,
    pub updated: bool,
}

impl Default for LayerDraws {
    fn default() -> Self {
        Self {
            primitives: Vec::new(),
            updated: true,
        }
    }
}

impl LayerDraws {
    pub fn write<P: Primitive>(
        &mut self,
        layer: usize,
        PrimitiveInst {
            kind,
            id,
            primitive,
            region,
            mask_idx,
        }: PrimitiveInst<P>,
    ) -> PrimitiveHandle {
        self.updated = true;
        // Grown on first use rather than sized from the registry, which a
        // layer cannot see.
        if self.primitives.len() <= kind.id as usize {
            self.primitives.resize_with(kind.id as usize + 1, || None);
        }
        let inst_idx = self.primitives[kind.id as usize]
            .get_or_insert_with(InstanceList::new::<P>)
            .push(
                id,
                PrimitiveInstance { region, mask_idx },
                bytemuck::bytes_of(&primitive),
            );
        PrimitiveHandle {
            layer,
            kind: kind.id,
            inst_idx,
        }
    }

    pub fn primitives(&self) -> &[Option<InstanceList>] {
        &self.primitives
    }

    pub fn apply_free(&mut self) -> impl Iterator<Item = PrimitiveChange> {
        self.primitives
            .iter_mut()
            .enumerate()
            .filter_map(|(kind, list)| Some((kind as u32, list.as_mut()?)))
            .flat_map(|(kind, list)| list.apply_free(kind))
    }

    pub fn free(&mut self, h: &PrimitiveHandle) -> MaskIdx {
        self.updated = true;
        self.list(h).free(h.inst_idx)
    }

    pub fn region_mut(&mut self, h: &PrimitiveHandle) -> &mut UiRegion {
        self.updated = true;
        &mut self.list(h).instances[h.inst_idx].region
    }

    /// A handle is only ever made by `write`, which is what created the list.
    fn list(&mut self, h: &PrimitiveHandle) -> &mut InstanceList {
        self.primitives[h.kind as usize]
            .as_mut()
            .expect("handle names a primitive this layer never drew")
    }
}

pub struct PrimitiveInst<P> {
    pub kind: PrimitiveKind<P>,
    pub id: WidgetId,
    pub primitive: P,
    pub region: UiRegion,
    pub mask_idx: MaskIdx,
}

pub struct PrimitiveChange {
    pub id: WidgetId,
    /// Which registered primitive's list moved, since they index separately.
    pub kind: u32,
    pub old: usize,
    pub new: usize,
}

#[derive(Debug)]
pub struct PrimitiveHandle {
    pub layer: usize,
    pub kind: u32,
    pub inst_idx: usize,
}

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct RectPrimitive {
    pub color: Color<u8>,
    pub radius: f32,
    pub thickness: f32,
    pub inner_radius: f32,
}

impl Primitive for RectPrimitive {
    const WGSL: &'static str = include_str!("shader/rect.wgsl");
}

impl RectPrimitive {
    pub fn color(color: Color<u8>) -> Self {
        Self {
            color,
            radius: 0.0,
            thickness: 0.0,
            inner_radius: 0.0,
        }
    }
}

/// `color` is multiplied by the atlas alpha for a mask glyph; a colour glyph
/// takes the texel unchanged, which `GlyphEntry::IS_COLORED` selects.
#[repr(C, align(8))]
#[derive(Debug, Copy, Clone)]
pub struct GlyphPrimitive {
    pub uv_min: Vec2,
    pub uv_max: Vec2,
    /// Which atlas array layer this glyph is on.
    pub layer: u32,
    pub color: Color<u8>,
    pub flags: u32,
}

// Manual rather than derived: the align(8) leaves four bytes of padding, which
// is how WGSL lays the struct out.
unsafe impl bytemuck::Pod for GlyphPrimitive {}
unsafe impl bytemuck::Zeroable for GlyphPrimitive {}
impl Primitive for GlyphPrimitive {
    const WGSL: &'static str = include_str!("shader/glyph.wgsl");

    fn render(device: &Device, queue: &Queue) -> Box<dyn PrimitiveRender> {
        Box::new(GlyphRender::new(device, queue))
    }
}

/// One drawn image. Its shader reads nothing per instance; the slot names the
/// texture to bind for it.
#[repr(C)]
#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct TexturePrimitive {
    pub slot: u32,
}

impl Primitive for TexturePrimitive {
    const WGSL: &'static str = include_str!("shader/texture.wgsl");

    fn render(device: &Device, queue: &Queue) -> Box<dyn PrimitiveRender> {
        Box::new(ImageRender::new(device, queue))
    }
}

impl From<&TextureHandle> for TexturePrimitive {
    fn from(handle: &TextureHandle) -> Self {
        Self {
            slot: handle.slot(),
        }
    }
}