use std::num::NonZeroUsize;
use std::rc::Rc;
use glam::{Mat4, Vec2, Vec3Swizzles};
use hashbrown::HashMap;
use lru::LruCache;
use ordered_float::OrderedFloat;
use wgpu::util::DeviceExt;
use wgpu::{Device, Queue, RenderPass, SurfaceConfiguration};
use crate::render::DrawCommand;
use crate::renderer::traits::{SpritePipeline, SpriteRenderer};
use crate::utils::rgb::to_linear_rgba;
use super::types::{CameraUniform, TextureResource, Vertex};
const PIPELINE_CACHE_SIZE: usize = 4096;
const MAX_SPRITES_PER_BATCH: usize = 1024;
const INDEX_PER_SPRITE: usize = 6;
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
struct PipeLineCacheKey {
id: u64,
shader: &'static str,
}
impl From<&BatchDrawCacheKey> for PipeLineCacheKey {
fn from(key: &BatchDrawCacheKey) -> Self {
PipeLineCacheKey {
id: key.id,
shader: key.shader,
}
}
}
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
struct BatchDrawCacheKey {
z_index: OrderedFloat<f32>,
id: u64,
shader: &'static str,
}
impl From<&DrawCommand> for BatchDrawCacheKey {
fn from(cmd: &DrawCommand) -> Self {
Self {
id: cmd.handle.id(),
z_index: OrderedFloat(cmd.transform.pos.z),
shader: cmd.renderer.key(),
}
}
}
struct BatchDraw {
texture: TextureResource,
vertexes: Vec<Vertex>,
count: usize,
shader: Rc<dyn SpriteRenderer>,
shader_buffer: Vec<u8>,
}
enum Boo<'a> {
Borrowed(&'a dyn SpritePipeline),
Owned(Box<dyn SpritePipeline>),
}
impl<'a> Boo<'a> {
pub fn as_ref(&'a self) -> &'a dyn SpritePipeline {
match self {
Boo::Borrowed(p) => *p,
Boo::Owned(p) => p.as_ref(),
}
}
}
pub struct SpriteRender {
device: Device,
queue: Queue,
config: SurfaceConfiguration,
index_buffer: wgpu::Buffer,
camera_buffer: wgpu::Buffer,
camera_uniform: CameraUniform,
pipeline_cache: LruCache<PipeLineCacheKey, Box<dyn SpritePipeline>>,
batch_cache: HashMap<BatchDrawCacheKey, Vec<BatchDraw>>,
}
impl SpriteRender {
pub fn setup(device: Device, queue: Queue, config: SurfaceConfiguration) -> Self {
let index_buffer = Self::setup_index_buffer(&device);
let pipeline_cache = LruCache::new(NonZeroUsize::new(PIPELINE_CACHE_SIZE).unwrap());
let batch_cache = HashMap::default();
let camera_buffer = Self::setup_camera_buffer(&device);
let camera_uniform = CameraUniform::default();
Self {
device,
queue,
config,
index_buffer,
camera_buffer,
camera_uniform,
pipeline_cache,
batch_cache,
}
}
fn setup_camera_buffer(device: &Device) -> wgpu::Buffer {
device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Index Buffer"),
contents: bytemuck::cast_slice(&[CameraUniform::default()]),
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
})
}
fn setup_index_buffer(device: &Device) -> wgpu::Buffer {
let mut indexes: Vec<u16> = Vec::with_capacity(MAX_SPRITES_PER_BATCH * INDEX_PER_SPRITE);
for i in 0..(MAX_SPRITES_PER_BATCH as u16) {
indexes.extend_from_slice(&[i * 4, 1 + i * 4, 2 + i * 4, i * 4, 2 + i * 4, 3 + i * 4]);
}
device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Index Buffer"),
contents: bytemuck::cast_slice(&indexes),
usage: wgpu::BufferUsages::INDEX,
})
}
pub fn draw(&mut self, cmd: DrawCommand, texture: TextureResource) {
let key = BatchDrawCacheKey::from(&cmd);
let batches = self.batch_cache.entry(key).or_default();
let batch = match batches.last_mut() {
Some(last) if last.count < MAX_SPRITES_PER_BATCH => last,
Some(_) | None => {
batches.push(BatchDraw {
texture,
vertexes: Default::default(),
count: 0,
shader: cmd.renderer,
shader_buffer: Vec::new(),
});
batches.last_mut().unwrap()
}
};
let size = cmd.transform.scaled_size();
let color = to_linear_rgba(cmd.color);
let pos = cmd.transform.pos;
let anchor = cmd.anchor;
let min = pos.xy() - size * anchor;
let max: Vec2 = pos.xy() + size * (Vec2::ONE - anchor);
let (min_tex, max_tex) = match cmd.src.clone() {
Some(src) => {
let size = batch.texture.texture.size();
let tex_size = Vec2::new(size.width as f32, size.height as f32);
let half_pixel = Vec2::splat(0.5);
(
(src.min + half_pixel) / tex_size,
(src.max - half_pixel) / tex_size,
)
}
None => (Vec2::ZERO, Vec2::ONE),
};
let mut vertexs = [
Vertex {
pos: min.to_array(),
tex_coord: [min_tex.x, max_tex.y],
color,
},
Vertex {
pos: Vec2::new(max.x, min.y).to_array(),
tex_coord: max_tex.to_array(),
color,
},
Vertex {
pos: max.to_array(),
tex_coord: [max_tex.x, min_tex.y],
color,
},
Vertex {
pos: Vec2::new(min.x, max.y).to_array(),
tex_coord: min_tex.to_array(),
color,
},
];
let mut swap_tex = |a: usize, b: usize, t: usize| {
let v = vertexs[a].tex_coord[t];
vertexs[a].tex_coord[t] = vertexs[b].tex_coord[t];
vertexs[b].tex_coord[t] = v;
};
if cmd.flip_x {
swap_tex(0, 1, 0);
swap_tex(3, 2, 0);
}
if cmd.flip_y {
swap_tex(0, 3, 1);
swap_tex(2, 1, 1);
}
if cmd.transform.angle != 0.0 {
let angle = Vec2::from_angle(cmd.transform.angle);
for v in &mut vertexs {
v.pos = (pos.xy() + angle.rotate(Vec2::from_array(v.pos) - pos.xy())).to_array();
}
}
batch.vertexes.extend_from_slice(&vertexs);
if let Some(data) = cmd.renderer_data {
batch.shader_buffer.extend(data);
}
batch.count += 1;
}
pub fn frame_pass(&mut self, view_proj: Mat4) {
self.camera_uniform.update(view_proj);
self.queue.write_buffer(
&self.camera_buffer,
0,
bytemuck::cast_slice(&[self.camera_uniform]),
);
}
pub fn frame_end(&mut self, pass: &mut RenderPass) {
let mut batch_cache: Vec<_> = std::mem::take(&mut self.batch_cache).into_iter().collect();
batch_cache.sort_by_key(|(key, _batches)| *key);
for (key, batches) in batch_cache {
for batch in batches {
let key = PipeLineCacheKey::from(&key);
let sp = if batch.shader_buffer.is_empty() {
Boo::Borrowed(
self.pipeline_cache
.get_or_insert(key, || {
batch.shader.new_pipeline(
&batch.texture,
self.config.format,
&self.camera_buffer,
batch.shader_buffer,
)
})
.as_ref(),
)
} else {
Boo::Owned(batch.shader.new_pipeline(
&batch.texture,
self.config.format,
&self.camera_buffer,
batch.shader_buffer,
))
};
let vertex_buffer =
self.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Vertex Buffer"),
contents: bytemuck::cast_slice(&batch.vertexes),
usage: wgpu::BufferUsages::VERTEX,
});
for i in 0..batch.count {
pass.set_vertex_buffer(0, vertex_buffer.slice(..));
pass.set_index_buffer(self.index_buffer.slice(..), wgpu::IndexFormat::Uint16);
sp.as_ref().render(pass, i);
let indexes_from = (i * INDEX_PER_SPRITE) as u64;
let indexes_to = indexes_from + INDEX_PER_SPRITE as u64;
pass.draw_indexed(indexes_from as u32..indexes_to as u32, 0, 0..1);
}
}
}
}
}