iris/src/render/primitive.rs

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

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

pub struct Primitives {
    instances: Vec<PrimitiveInstance>,
    assoc: Vec<Id>,
    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,
        }
    }
}

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 {
            $($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 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: Id,
    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 vec = P::vec(&mut self.data);
        let i = vec.add(primitive);
        let inst = PrimitiveInstance {
            region,
            idx: i as u32,
            mask_idx,
            binding: P::BINDING,
        };
        let inst_i = 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::new::<P>(layer, inst_i, i)
    }

    /// 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.data.free(h.binding, h.data_idx);
        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: Id,
    pub old: usize,
    pub new: usize,
}

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

impl PrimitiveHandle {
    fn new<P: Primitive>(layer: usize, inst_idx: usize, data_idx: usize) -> Self {
        Self {
            layer,
            inst_idx,
            data_idx,
            binding: P::BINDING,
        }
    }
}

primitives!(
    rects: RectPrimitive => 0,
    textures: TexturePrimitive => 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,
        }
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct TexturePrimitive {
    pub view_idx: u32,
    pub sampler_idx: 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
    }
}