iris/core/src/render/primitive.rs

use std::ops::{Deref, DerefMut};

use crate::{
    Color, UiRegion, WidgetId,
    render::{
        ArrBuf,
        data::{MaskIdx, PrimitiveInstance},
    },
    util::Vec2,
};
use bytemuck::Pod;
use wgpu::*;

pub struct Primitives {
    instances: Vec<PrimitiveInstance>,
    assoc: Vec<WidgetId>,
    data: PrimitiveData,
    free: Vec<usize>,
    pub updated: bool,
}

impl Default for Primitives {
    fn default() -> Self {
        Self {
            instances: Default::default(),
            assoc: Default::default(),
            data: Default::default(),
            free: Vec::new(),
            updated: true,
        }
    }
}

/// Not in `primitives!` and with no group-1 buffer: a texture instance's `idx`
/// names the texture to bind, so there is nothing to look up per-instance.
pub const TEXTURE_BINDING: u32 = 2;

pub trait Primitive: Pod {
    const BINDING: u32;
    fn vec(data: &mut PrimitiveData) -> &mut PrimitiveVec<Self>;
}

macro_rules! primitives {
    ($($name:ident: $ty:ty => $binding:expr,)*) => {
        #[derive(Default)]
        pub struct PrimitiveData {
            $(pub(crate) $name: PrimitiveVec<$ty>,)*
        }

        pub struct PrimitiveBuffers {
            $($name: ArrBuf<$ty>,)*
        }

        impl PrimitiveBuffers {
            pub fn update(&mut self, device: &Device, queue: &Queue, data: &PrimitiveData) {
                $(self.$name.update(device, queue, &data.$name);)*
            }
        }

        impl PrimitiveBuffers {
            pub const LEN: usize = primitives!(@count $($name)*);
            pub const BINDINGS: [u32; Self::LEN] = [$(<$ty>::BINDING,)*];
            pub fn buffers(&self) -> [(u32, &Buffer); Self::LEN] {
                [
                    $((<$ty>::BINDING, &self.$name.buffer),)*
                ]
            }
            pub fn new(device: &Device) -> Self {
                Self {
                    $($name: ArrBuf::new(
                        device,
                        BufferUsages::STORAGE | BufferUsages::COPY_DST,
                        stringify!($name),
                    ),)*
                }
            }
        }

        impl PrimitiveData {
            pub fn clear(&mut self) {
                $(self.$name.clear();)*
            }
            pub fn free(&mut self, binding: u32, idx: usize) {
                match binding {
                    $(<$ty>::BINDING => self.$name.free(idx),)*
                    _ => unreachable!()
                }
            }
        }

        $(
            unsafe impl bytemuck::Pod for $ty {}
            unsafe impl bytemuck::Zeroable for $ty {}
            impl Primitive for $ty {
                const BINDING: u32 = $binding;
                fn vec(data: &mut PrimitiveData) -> &mut PrimitiveVec<Self> {
                    &mut data.$name
                }
            }
        )*
    };
    (@count $t1:tt $($t:tt)+) => { 1 + primitives!(@count $($t)+) };
    (@count $t:tt) => { 1 };
}

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

impl Primitives {
    pub fn write<P: Primitive>(
        &mut self,
        layer: usize,
        PrimitiveInst {
            id,
            primitive,
            region,
            mask_idx,
        }: PrimitiveInst<P>,
    ) -> PrimitiveHandle {
        let data_idx = P::vec(&mut self.data).add(primitive);
        self.push(
            layer,
            id,
            PrimitiveInstance {
                region,
                idx: data_idx as u32,
                mask_idx,
                binding: P::BINDING,
            },
            Some(data_idx),
        )
    }

    /// Writes an instance that samples `texture` instead of a group-1 buffer.
    pub fn write_texture(
        &mut self,
        layer: usize,
        id: WidgetId,
        texture: u32,
        region: UiRegion,
        mask_idx: MaskIdx,
    ) -> PrimitiveHandle {
        self.push(
            layer,
            id,
            PrimitiveInstance {
                region,
                idx: texture,
                mask_idx,
                binding: TEXTURE_BINDING,
            },
            None,
        )
    }

    fn push(
        &mut self,
        layer: usize,
        id: WidgetId,
        inst: PrimitiveInstance,
        data_idx: Option<usize>,
    ) -> PrimitiveHandle {
        self.updated = true;
        let inst_idx = if let Some(i) = self.free.pop() {
            self.instances[i] = inst;
            self.assoc[i] = id;
            i
        } else {
            let i = self.instances.len();
            self.instances.push(inst);
            self.assoc.push(id);
            i
        };
        PrimitiveHandle {
            layer,
            inst_idx,
            data_idx,
            binding: inst.binding,
        }
    }

    /// returns (old index, new index)
    pub fn apply_free(&mut self) -> impl Iterator<Item = PrimitiveChange> {
        self.free.sort_by(|a, b| b.cmp(a));
        self.free.drain(..).filter_map(|i| {
            self.instances.swap_remove(i);
            self.assoc.swap_remove(i);
            if i == self.instances.len() {
                return None;
            }
            let id = self.assoc[i];
            let old = self.instances.len();
            Some(PrimitiveChange { id, old, new: i })
        })
    }

    pub fn free(&mut self, h: &PrimitiveHandle) -> MaskIdx {
        self.updated = true;
        if let Some(i) = h.data_idx {
            self.data.free(h.binding, i);
        }
        self.free.push(h.inst_idx);
        self.instances[h.inst_idx].mask_idx
    }

    pub fn data(&self) -> &PrimitiveData {
        &self.data
    }

    pub fn instances(&self) -> &Vec<PrimitiveInstance> {
        &self.instances
    }

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

pub struct PrimitiveChange {
    pub id: WidgetId,
    pub old: usize,
    pub new: usize,
}

#[derive(Debug)]
pub struct PrimitiveHandle {
    pub layer: usize,
    pub inst_idx: usize,
    /// `None` for a texture instance, which has no group-1 entry to free.
    pub data_idx: Option<usize>,
    pub binding: u32,
}

primitives!(
    rects: RectPrimitive => 0,
    glyphs: GlyphPrimitive => 1,
);

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

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)]
#[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,
}

pub struct PrimitiveVec<T> {
    vec: Vec<T>,
    free: Vec<usize>,
}

impl<T> PrimitiveVec<T> {
    pub fn new() -> Self {
        Self {
            vec: Vec::new(),
            free: Vec::new(),
        }
    }
    pub fn add(&mut self, t: T) -> usize {
        if let Some(i) = self.free.pop() {
            self.vec[i] = t;
            i
        } else {
            let i = self.vec.len();
            self.vec.push(t);
            i
        }
    }
    pub fn free(&mut self, i: usize) {
        self.free.push(i);
    }
    pub fn clear(&mut self) {
        self.free.clear();
        self.vec.clear();
    }
}

impl<T> Default for PrimitiveVec<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Deref for PrimitiveVec<T> {
    type Target = Vec<T>;

    fn deref(&self) -> &Self::Target {
        &self.vec
    }
}

impl<T> DerefMut for PrimitiveVec<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.vec
    }
}