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, TextureFormat};
use crate::commands::Draw;
use crate::renderer::traits::{Shader, SpritePipeline};
use crate::types::{DrawContent, DrawContentCacheKey};
use crate::utils::bytes::to_wgsl_bytes;
use crate::utils::rgb::to_linear_rgba;
use crate::platform::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, Hash, PartialEq, Eq, PartialOrd, Ord)]
struct PipeLineCacheKey {
shader: &'static str,
transparent: bool,
}
impl From<BatchDrawCacheKey> for PipeLineCacheKey {
fn from(key: BatchDrawCacheKey) -> Self {
PipeLineCacheKey {
shader: key.shader,
transparent: key.transparent,
}
}
}
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct BatchDrawCacheKey {
z_index: OrderedFloat<f32>,
content: DrawContentCacheKey,
shader: &'static str,
transparent: bool,
}
impl BatchDrawCacheKey {
fn from_draw(cmd: &Draw) -> Self {
Self {
content: cmd.content.clone().into(),
z_index: OrderedFloat(cmd.transform.pos.z),
shader: cmd.shader.key(),
transparent: cmd.transparent,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PrimitiveKind {
IndexedQuads,
TriangleList,
}
struct BatchDraw {
texture: Option<TextureResource>,
vertices: Vec<Vertex>,
count: usize,
shader: Rc<dyn Shader>,
shader_buffer: Vec<u8>,
primitive: PrimitiveKind,
}
pub struct Render {
device: Device,
queue: Queue,
color_format: TextureFormat,
depth_format: TextureFormat,
index_buffer: wgpu::Buffer,
camera_buffer: wgpu::Buffer,
camera_uniform: CameraUniform,
pipeline_cache: LruCache<PipeLineCacheKey, Box<dyn SpritePipeline>>,
batch_cache: HashMap<BatchDrawCacheKey, Vec<BatchDraw>>,
transparent_batch_cache: HashMap<BatchDrawCacheKey, Vec<BatchDraw>>,
}
impl Render {
pub fn setup(
device: Device,
queue: Queue,
color_format: TextureFormat,
depth_format: TextureFormat,
) -> 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 transparent_batch_cache = HashMap::default();
let camera_buffer = Self::setup_camera_buffer(&device);
let camera_uniform = CameraUniform::default();
Self {
device,
queue,
color_format,
depth_format,
index_buffer,
camera_buffer,
camera_uniform,
pipeline_cache,
batch_cache,
transparent_batch_cache,
}
}
fn setup_camera_buffer(device: &Device) -> wgpu::Buffer {
device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Index Buffer"),
contents: to_wgsl_bytes(&CameraUniform::default()).as_ref(),
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
})
}
fn setup_index_buffer(device: &Device) -> wgpu::Buffer {
let mut indexes = [0u16; MAX_SPRITES_PER_BATCH * INDEX_PER_SPRITE];
for i in 0..(MAX_SPRITES_PER_BATCH as u16) {
let start = i as usize * INDEX_PER_SPRITE;
let end = start + INDEX_PER_SPRITE;
indexes[start..end].copy_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("Sprite Index Buffer"),
contents: bytemuck::cast_slice(&indexes),
usage: wgpu::BufferUsages::INDEX,
})
}
pub fn draw(&mut self, cmd: Draw, texture: Option<TextureResource>) {
match &cmd.content {
DrawContent::Texture { .. } => self.draw_sprite(cmd, texture),
DrawContent::Lines(_) => self.draw_lines_batch(cmd, texture),
DrawContent::Circle { .. } => self.draw_circle(cmd, texture),
DrawContent::Polygon { .. } | DrawContent::PolygonTextured { .. } => {
self.draw_polygon(cmd, texture)
}
}
}
fn get_or_create_batch(
&mut self,
cmd: &Draw,
texture: Option<TextureResource>,
) -> &mut BatchDraw {
let key = BatchDrawCacheKey::from_draw(cmd);
let cache = if cmd.transparent {
&mut self.transparent_batch_cache
} else {
&mut self.batch_cache
};
let batches = cache.entry(key).or_default();
if batches.is_empty()
|| batches
.last()
.is_some_and(|last| last.count >= MAX_SPRITES_PER_BATCH)
|| cmd.shader_data.is_some()
{
batches.push(BatchDraw {
texture,
vertices: Default::default(),
count: 0,
shader: cmd.shader.clone(),
shader_buffer: Vec::new(),
primitive: PrimitiveKind::IndexedQuads,
});
}
batches.last_mut().unwrap()
}
fn draw_sprite(&mut self, cmd: Draw, texture: Option<TextureResource>) {
let batch = self.get_or_create_batch(&cmd, texture);
let size = cmd.size * cmd.transform.scale;
let color = match &cmd.content {
DrawContent::Texture { color, .. } => to_linear_rgba(*color),
DrawContent::Circle { color, .. } => to_linear_rgba(*color),
_ => to_linear_rgba(crate::color::WHITE),
};
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(), &batch.texture) {
(Some(src), Some(texture)) => {
let size = 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,
)
}
_ => (Vec2::ZERO, Vec2::ONE),
};
let mut vertices = [
Vertex {
pos: [min.x, min.y, pos.z],
tex_coord: [min_tex.x, max_tex.y],
color,
},
Vertex {
pos: [max.x, min.y, pos.z],
tex_coord: max_tex.to_array(),
color,
},
Vertex {
pos: [max.x, max.y, pos.z],
tex_coord: [max_tex.x, min_tex.y],
color,
},
Vertex {
pos: [min.x, max.y, pos.z],
tex_coord: min_tex.to_array(),
color,
},
];
let mut swap_tex = |a: usize, b: usize, t: usize| {
let v = vertices[a].tex_coord[t];
vertices[a].tex_coord[t] = vertices[b].tex_coord[t];
vertices[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);
}
let angle = cmd.transform.angle();
if angle != 0.0 {
let angle = Vec2::from_angle(angle);
for v in &mut vertices {
let p = Vec2::from_slice(&v.pos[0..2]);
let rotated = pos.xy() + angle.rotate(p - pos.xy());
v.pos[0] = rotated.x;
v.pos[1] = rotated.y;
}
}
batch.vertices.extend_from_slice(&vertices);
if let Some(data) = cmd.shader_data {
batch.shader_buffer.extend(data);
}
batch.count += 1;
}
fn draw_lines_batch(&mut self, cmd: Draw, texture: Option<TextureResource>) {
if let DrawContent::Lines(lines) = &cmd.content {
let batch = self.get_or_create_batch(&cmd, texture.clone());
for line in lines {
let color = to_linear_rgba(line.color);
let line_dir = (line.end - line.start).normalize_or_zero();
let perp = Vec2::new(-line_dir.y, line_dir.x) * (line.thickness * 0.5);
let start_pos = cmd.transform.pos;
let vertices = [
Vertex {
pos: (line.start - perp).extend(start_pos.z).to_array(),
tex_coord: [0.0, 0.0],
color,
},
Vertex {
pos: (line.end - perp).extend(start_pos.z).to_array(),
tex_coord: [1.0, 0.0],
color,
},
Vertex {
pos: (line.end + perp).extend(start_pos.z).to_array(),
tex_coord: [1.0, 1.0],
color,
},
Vertex {
pos: (line.start + perp).extend(start_pos.z).to_array(),
tex_coord: [0.0, 1.0],
color,
},
];
batch.vertices.extend_from_slice(&vertices);
batch.count += 1;
}
}
}
fn draw_circle(&mut self, cmd: Draw, texture: Option<TextureResource>) {
self.draw_sprite(cmd, texture)
}
fn draw_polygon(&mut self, cmd: Draw, texture: Option<TextureResource>) {
let key = BatchDrawCacheKey::from_draw(&cmd);
let cache = if cmd.transparent {
&mut self.transparent_batch_cache
} else {
&mut self.batch_cache
};
let batches = cache.entry(key).or_default();
let Draw {
shader,
mut shader_data,
content,
transform,
..
} = cmd;
match content {
DrawContent::Polygon { vertices, color } => {
if vertices.len() < 3 {
return;
}
let mut batch = BatchDraw {
texture: None,
vertices: Vec::with_capacity((vertices.len() - 2) * 3),
count: 0,
shader,
shader_buffer: shader_data.take().unwrap_or_default(),
primitive: PrimitiveKind::TriangleList,
};
let color = to_linear_rgba(color);
let pivot = vertices[0];
for i in 1..vertices.len() - 1 {
for point in [pivot, vertices[i], vertices[i + 1]] {
batch.vertices.push(Vertex {
pos: point.extend(transform.pos.z).to_array(),
tex_coord: [0.0, 0.0],
color,
});
}
}
batches.push(batch);
}
DrawContent::PolygonTextured {
vertices,
uvs,
color,
..
} => {
if vertices.len() < 3 || uvs.len() != vertices.len() {
return;
}
let Some(texture) = texture else {
return;
};
let mut batch = BatchDraw {
texture: Some(texture),
vertices: Vec::with_capacity((vertices.len() - 2) * 3),
count: 0,
shader,
shader_buffer: shader_data.take().unwrap_or_default(),
primitive: PrimitiveKind::TriangleList,
};
let color = to_linear_rgba(color);
let pivot = vertices[0];
let pivot_uv = uvs[0];
for i in 1..vertices.len() - 1 {
for (point, uv) in [
(pivot, pivot_uv),
(vertices[i], uvs[i]),
(vertices[i + 1], uvs[i + 1]),
] {
batch.vertices.push(Vertex {
pos: point.extend(transform.pos.z).to_array(),
tex_coord: uv.to_array(),
color,
});
}
}
batches.push(batch);
}
_ => {}
}
}
pub fn frame_pass(&mut self, view_proj: Mat4) {
self.camera_uniform.update(view_proj);
self.queue
.write_buffer(&self.camera_buffer, 0, &to_wgsl_bytes(&self.camera_uniform));
}
pub fn frame_end(&mut self, pass: &mut RenderPass) {
let batch_cache: Vec<_> = std::mem::take(&mut self.batch_cache).into_iter().collect();
for (key, batches) in batch_cache {
let pipeline_key = PipeLineCacheKey::from(key.clone());
for batch in batches {
let sp = self
.pipeline_cache
.get_or_insert(pipeline_key.clone(), || {
batch.shader.new_pipeline(
self.color_format,
Some(self.depth_format),
&self.camera_buffer,
false,
)
})
.as_ref();
write_batch(&self.device, &self.index_buffer, pass, &batch, sp);
}
}
let mut transparent_batch_cache: Vec<_> = std::mem::take(&mut self.transparent_batch_cache)
.into_iter()
.collect();
transparent_batch_cache.sort_by(|(a, _), (b, _)| b.z_index.cmp(&a.z_index));
for (key, batches) in transparent_batch_cache {
let pipeline_key = PipeLineCacheKey::from(key.clone());
for batch in batches {
let sp = self
.pipeline_cache
.get_or_insert(pipeline_key.clone(), || {
batch.shader.new_pipeline(
self.color_format,
Some(self.depth_format),
&self.camera_buffer,
true,
)
})
.as_ref();
write_batch(&self.device, &self.index_buffer, pass, &batch, sp);
}
}
}
}
fn write_batch(
device: &Device,
index_buffer: &wgpu::Buffer,
pass: &mut RenderPass,
batch: &BatchDraw,
sp: &dyn SpritePipeline,
) {
let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Vertex Buffer"),
contents: bytemuck::cast_slice(&batch.vertices),
usage: wgpu::BufferUsages::VERTEX,
});
pass.set_vertex_buffer(0, vertex_buffer.slice(..));
let shader_data_opt = if batch.shader_buffer.is_empty() {
None
} else {
Some(batch.shader_buffer.as_slice())
};
sp.render(pass, batch.texture.as_ref(), shader_data_opt);
match batch.primitive {
PrimitiveKind::IndexedQuads => {
pass.set_index_buffer(index_buffer.slice(..), wgpu::IndexFormat::Uint16);
let index_count = (batch.count * INDEX_PER_SPRITE) as u32;
pass.draw_indexed(0..index_count, 0, 0..1);
}
PrimitiveKind::TriangleList => {
pass.draw(0..batch.vertices.len() as u32, 0..1);
}
}
}