bworld/client/renderer/core.ts

import { Camera } from "../components/camera.ts";
import { mat4 } from "gl-matrix";

export let device: GPUDevice;
export let canvas: HTMLCanvasElement;

const MAX_SPRITES = 100000;
const VERTS_PER_SPRITE = 6;
const FLOATS_PER_VERT = 9;
const VERTEX_STRIDE = FLOATS_PER_VERT * 4;

const SAMPLE_COUNT = 4;
const DEPTH_FORMAT: GPUTextureFormat = "depth24plus";

// dynamic uniform offsets have to be 256 aligned
const UNIFORM_SLOT_SIZE = 256;

let context: GPUCanvasContext;
let canvas_format: GPUTextureFormat;

let pipeline_2d: GPURenderPipeline;
let pipeline_3d: GPURenderPipeline;
let uniform_layout: GPUBindGroupLayout;
let texture_layout: GPUBindGroupLayout;
let sampler: GPUSampler;

const vertex_data = new Float32Array(MAX_SPRITES * VERTS_PER_SPRITE * FLOATS_PER_VERT);
let vert_index = 0;

// every flush in a frame appends to this, so earlier draws dont get overwritten before the submit
let stream_buffer: GPUBuffer;
let stream_offset = 0;

let uniform_buffer: GPUBuffer;
let uniform_bind_group: GPUBindGroup;
let uniform_slot = -1;

const texture_bind_groups = new WeakMap<GPUTexture, GPUBindGroup>();

let current_texture: GPUTexture | null = null;
export let white_tex: GPUTexture | null = null;

let mode3d = false;

let camera: Camera | undefined;

const ortho = mat4.create();
const proj = mat4.create();
const view = mat4.create();
const mvp = mat4.create();

// per frame state
let encoder: GPUCommandEncoder | undefined;
let pass: GPURenderPassEncoder | undefined;
let frame_view: GPUTextureView;
let msaa_texture: GPUTexture | undefined;
let depth_texture: GPUTexture | undefined;
let current_pipeline: GPURenderPipeline | undefined;
let clear_color: GPUColor = { r: 0, g: 0, b: 0, a: 1 };
let pending_color_clear = false;
let pending_depth_clear = false;
let scissor: { x: number; y: number; width: number; height: number } | undefined;
let pending_destroy: GPUBuffer[] = [];

const shader_src = /* wgsl */ `
struct Uniforms {
	mvp: mat4x4<f32>,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(1) @binding(0) var texture0: texture_2d<f32>;
@group(1) @binding(1) var sampler0: sampler;

struct VertexOut {
	@builtin(position) position: vec4<f32>,
	@location(0) tex_coord: vec2<f32>,
	@location(1) color: vec4<f32>,
}

@vertex
fn vs_main(
	@location(0) position: vec3<f32>,
	@location(1) tex_coord: vec2<f32>,
	@location(2) color: vec4<f32>,
) -> VertexOut {
	var out: VertexOut;
	out.position = uniforms.mvp * vec4<f32>(position, 1.0);
	out.tex_coord = tex_coord;
	out.color = color;
	return out;
}

@fragment
fn fs_main(in: VertexOut) -> @location(0) vec4<f32> {
	return textureSample(texture0, sampler0, in.tex_coord) * in.color;
}
`;

export async function init_window(canvas_element: HTMLCanvasElement) {
	canvas = canvas_element;
	canvas.width = 1800;
	canvas.height = 900;

	if (!navigator.gpu) {
		throw new Error("WebGPU not supported");
	}
	const adapter = await navigator.gpu.requestAdapter();
	if (!adapter) {
		throw new Error("No WebGPU adapter found");
	}
	device = await adapter.requestDevice();
	device.addEventListener("uncapturederror", (event) => {
		console.error("WebGPU error:", (event as GPUUncapturedErrorEvent).error.message);
	});
	device.lost.then((info) => console.error(`WebGPU device lost: ${info.message}`));

	// the dom typings dont know about the webgpu overload yet
	const ctx = canvas.getContext("webgpu") as GPUCanvasContext | null;
	if (!ctx) {
		throw new Error("WebGPU not supported");
	}
	context = ctx;
	canvas_format = navigator.gpu.getPreferredCanvasFormat();
	context.configure({ device, format: canvas_format, alphaMode: "opaque" });

	create_pipelines();

	sampler = device.createSampler({
		magFilter: "nearest",
		minFilter: "nearest",
		addressModeU: "clamp-to-edge",
		addressModeV: "clamp-to-edge",
	});

	stream_buffer = create_stream_buffer(8 * 1024 * 1024);
	create_uniform_buffer(64);

	create_white_texture();
}

export function begin_drawing() {
	encoder = device.createCommandEncoder();
	frame_view = context.getCurrentTexture().createView();
	ensure_render_targets();

	update_2d_mvp();
	stream_offset = 0;
	uniform_slot = -1;
	write_mvp();

	pending_color_clear = true;
	pending_depth_clear = true;
	vert_index = 0;
}

export function end_drawing() {
	flush_batch();
	// make sure the frame gets cleared even if nothing was drawn
	ensure_pass();

	pass!.end();
	pass = undefined;
	device.queue.submit([encoder!.finish()]);
	encoder = undefined;

	for (const buffer of pending_destroy) {
		buffer.destroy();
	}
	pending_destroy = [];
}

function update_2d_mvp() {
	const w = canvas.width;
	const h = canvas.height;

	const l = 0;
	const r = w;
	const t = 0;
	const b = h;
	const n = -1;
	const f = 1;

	mat4.orthoZO(ortho, l, r, b, t, n, f);
}

export function begin_clip(x: number, y: number, width: number, height: number) {
	flush_batch();
	scissor = { x, y, width, height };
	if (pass) {
		apply_scissor(pass);
	}
}

export function end_clip() {
	flush_batch();
	scissor = undefined;
	if (pass) {
		apply_scissor(pass);
	}
}

export function begin_mode_3d(new_camera: Camera) {
	flush_batch();

	mode3d = true;
	camera = new_camera;

	update_camera();
	write_mvp();
}

export function end_mode_3d() {
	if (!mode3d) {
		return;
	}

	flush_batch();

	mode3d = false;
	write_mvp();
}

export function clear_background(r: number, g: number, b: number, a = 1) {
	flush_batch();
	clear_color = { r, g, b, a };
	// the clear happens when a render pass starts, so start a new one
	if (pass) {
		pass.end();
		pass = undefined;
	}
	pending_color_clear = true;
}

export function flush_batch() {
	if (vert_index === 0 || !current_texture) {
		return;
	}

	const byte_length = vert_index * 4;
	if (stream_offset + byte_length > stream_buffer.size) {
		// the old buffer might still be used by draws in this frame, destroy it after the submit
		pending_destroy.push(stream_buffer);
		stream_buffer = create_stream_buffer(Math.max(stream_buffer.size * 2, byte_length));
		stream_offset = 0;
	}

	device.queue.writeBuffer(stream_buffer, stream_offset, vertex_data, 0, vert_index);
	draw(stream_buffer, stream_offset, vert_index / FLOATS_PER_VERT);
	stream_offset += byte_length;

	vert_index = 0;
}

export function flush_buffer(buffer: GPUBuffer, draw_count: number) {
	draw(buffer, 0, draw_count);
}

export function create_vertex_buffer(vertices: Float32Array): GPUBuffer {
	const buffer = device.createBuffer({
		size: Math.max(4, vertices.byteLength),
		usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
	});
	device.queue.writeBuffer(buffer, 0, vertices);
	return buffer;
}

export function destroy_vertex_buffer(buffer: GPUBuffer) {
	// it might be used by a draw thats not submitted yet
	if (encoder) {
		pending_destroy.push(buffer);
	} else {
		buffer.destroy();
	}
}

export function resize_canvas() {
	const width = self.innerWidth;
	const height = self.innerHeight;

	if (canvas.width !== width || canvas.height !== height) {
		canvas.width = width;
		canvas.height = height;
		canvas.style.width = canvas.width + "px";
		canvas.style.height = canvas.height + "px";
		// render targets get recreated at the start of the next frame
	}
}

export function get_current_texture(): GPUTexture | null {
	return current_texture;
}

export function set_current_texture(texture: GPUTexture) {
	current_texture = texture;
}

export function create_texture(width: number, height: number): GPUTexture {
	return device.createTexture({
		size: [width, height],
		format: "rgba8unorm",
		usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST | GPUTextureUsage.RENDER_ATTACHMENT,
	});
}

export function push_vertex(px: number, py: number, pz: number, u: number, vv: number, r = 1, g = 1, b = 1, a = 1) {
	let i = vert_index;
	vertex_data[i++] = px;
	vertex_data[i++] = py;
	vertex_data[i++] = pz;
	vertex_data[i++] = u;
	vertex_data[i++] = vv;
	vertex_data[i++] = r;
	vertex_data[i++] = g;
	vertex_data[i++] = b;
	vertex_data[i++] = a;
	vert_index = i;
}

export function push_quad(
	dx: number,
	dy: number,
	dw: number,
	dh: number,
	u0: number,
	v0: number,
	u1: number,
	v1: number,
	r = 1,
	g = 1,
	b = 1,
	a = 1,
) {
	const x2 = dx + dw;
	const y2 = dy + dh;

	// triangle 1
	push_vertex(dx, dy, 0, u0, v0, r, g, b, a);
	push_vertex(x2, dy, 0, u1, v0, r, g, b, a);
	push_vertex(dx, y2, 0, u0, v1, r, g, b, a);

	// triangle 2
	push_vertex(x2, dy, 0, u1, v0, r, g, b, a);
	push_vertex(x2, y2, 0, u1, v1, r, g, b, a);
	push_vertex(dx, y2, 0, u0, v1, r, g, b, a);
}

export function push_quad_vertices(
	x0: number,
	y0: number,
	x1: number,
	y1: number,
	x2: number,
	y2: number,
	x3: number,
	y3: number,
	u0: number,
	v0: number,
	u1: number,
	v1: number,
	r = 1,
	g = 1,
	b = 1,
	a = 1,
) {
	// triangle 1
	push_vertex(x0, y0, 0, u0, v0, r, g, b, a);
	push_vertex(x1, y1, 0, u1, v0, r, g, b, a);
	push_vertex(x2, y2, 0, u1, v1, r, g, b, a);

	// triangle 2
	push_vertex(x0, y0, 0, u0, v0, r, g, b, a);
	push_vertex(x2, y2, 0, u1, v1, r, g, b, a);
	push_vertex(x3, y3, 0, u0, v1, r, g, b, a);
}

// internal

function draw(buffer: GPUBuffer, offset: number, vertex_count: number) {
	if (!current_texture || vertex_count === 0) {
		return;
	}

	const render_pass = ensure_pass();

	const pipeline = mode3d ? pipeline_3d : pipeline_2d;
	if (current_pipeline !== pipeline) {
		render_pass.setPipeline(pipeline);
		current_pipeline = pipeline;
	}

	render_pass.setBindGroup(0, uniform_bind_group, [uniform_slot * UNIFORM_SLOT_SIZE]);
	render_pass.setBindGroup(1, get_texture_bind_group(current_texture));
	render_pass.setVertexBuffer(0, buffer, offset);
	render_pass.draw(vertex_count);
}

function ensure_pass(): GPURenderPassEncoder {
	if (pass) {
		return pass;
	}

	pass = encoder!.beginRenderPass({
		colorAttachments: [{
			view: msaa_texture!.createView(),
			resolveTarget: frame_view,
			clearValue: clear_color,
			loadOp: pending_color_clear ? "clear" : "load",
			storeOp: "store",
		}],
		depthStencilAttachment: {
			view: depth_texture!.createView(),
			depthClearValue: 1,
			depthLoadOp: pending_depth_clear ? "clear" : "load",
			depthStoreOp: "store",
		},
	});
	pending_color_clear = false;
	pending_depth_clear = false;
	current_pipeline = undefined;
	apply_scissor(pass);

	return pass;
}

function apply_scissor(render_pass: GPURenderPassEncoder) {
	const target_width = msaa_texture!.width;
	const target_height = msaa_texture!.height;

	if (!scissor) {
		render_pass.setScissorRect(0, 0, target_width, target_height);
		return;
	}

	// webgpu errors on rects outside the target instead of clipping them
	const x0 = Math.min(target_width, Math.max(0, Math.floor(scissor.x)));
	const y0 = Math.min(target_height, Math.max(0, Math.floor(scissor.y)));
	const x1 = Math.min(target_width, Math.max(x0, Math.ceil(scissor.x + scissor.width)));
	const y1 = Math.min(target_height, Math.max(y0, Math.ceil(scissor.y + scissor.height)));

	render_pass.setScissorRect(x0, y0, x1 - x0, y1 - y0);
}

function ensure_render_targets() {
	const width = canvas.width;
	const height = canvas.height;

	if (msaa_texture && msaa_texture.width === width && msaa_texture.height === height) {
		return;
	}

	msaa_texture?.destroy();
	depth_texture?.destroy();

	msaa_texture = device.createTexture({
		size: [width, height],
		format: canvas_format,
		sampleCount: SAMPLE_COUNT,
		usage: GPUTextureUsage.RENDER_ATTACHMENT,
	});
	depth_texture = device.createTexture({
		size: [width, height],
		format: DEPTH_FORMAT,
		sampleCount: SAMPLE_COUNT,
		usage: GPUTextureUsage.RENDER_ATTACHMENT,
	});
}

// puts the current matrix in a new uniform slot, earlier draws in the frame keep using theirs
function write_mvp() {
	uniform_slot += 1;
	if (uniform_slot * UNIFORM_SLOT_SIZE >= uniform_buffer.size) {
		pending_destroy.push(uniform_buffer);
		create_uniform_buffer((uniform_buffer.size / UNIFORM_SLOT_SIZE) * 2);
		uniform_slot = 0;
	}

	const matrix = mode3d ? mvp : ortho;
	device.queue.writeBuffer(uniform_buffer, uniform_slot * UNIFORM_SLOT_SIZE, matrix as Float32Array);
}

function update_camera() {
	if (!camera) {
		return;
	}

	mat4.perspectiveZO(
		proj,
		camera.fov,
		canvas.width / canvas.height,
		camera.near,
		camera.far,
	);

	mat4.identity(view);

	mat4.rotateZ(view, view, -camera.roll);
	mat4.rotateX(view, view, -camera.pitch);
	mat4.rotateY(view, view, -camera.yaw);

	mat4.translate(view, view, [-camera.x, -camera.y, -camera.z]);

	mat4.multiply(mvp, proj, view);
}

function create_pipelines() {
	const module = device.createShaderModule({ code: shader_src });

	uniform_layout = device.createBindGroupLayout({
		entries: [{
			binding: 0,
			visibility: GPUShaderStage.VERTEX,
			buffer: { type: "uniform", hasDynamicOffset: true, minBindingSize: 64 },
		}],
	});
	texture_layout = device.createBindGroupLayout({
		entries: [
			{ binding: 0, visibility: GPUShaderStage.FRAGMENT, texture: {} },
			{ binding: 1, visibility: GPUShaderStage.FRAGMENT, sampler: {} },
		],
	});
	const layout = device.createPipelineLayout({ bindGroupLayouts: [uniform_layout, texture_layout] });

	const vertex: GPUVertexState = {
		module,
		entryPoint: "vs_main",
		buffers: [{
			arrayStride: VERTEX_STRIDE,
			attributes: [
				{ shaderLocation: 0, offset: 0, format: "float32x3" },
				{ shaderLocation: 1, offset: 12, format: "float32x2" },
				{ shaderLocation: 2, offset: 20, format: "float32x4" },
			],
		}],
	};
	const fragment: GPUFragmentState = {
		module,
		entryPoint: "fs_main",
		targets: [{
			format: canvas_format,
			blend: {
				color: { srcFactor: "src-alpha", dstFactor: "one-minus-src-alpha", operation: "add" },
				alpha: { srcFactor: "src-alpha", dstFactor: "one-minus-src-alpha", operation: "add" },
			},
		}],
	};
	const multisample: GPUMultisampleState = { count: SAMPLE_COUNT };

	pipeline_2d = device.createRenderPipeline({
		layout,
		vertex,
		fragment,
		multisample,
		primitive: { topology: "triangle-list", cullMode: "none" },
		depthStencil: { format: DEPTH_FORMAT, depthWriteEnabled: false, depthCompare: "always" },
	});

	pipeline_3d = device.createRenderPipeline({
		layout,
		vertex,
		fragment,
		multisample,
		primitive: { topology: "triangle-list", cullMode: "back", frontFace: "ccw" },
		depthStencil: {
			format: DEPTH_FORMAT,
			depthWriteEnabled: true,
			depthCompare: "less-equal",
			// same as the old polygonOffset(1, 1)
			depthBias: 1,
			depthBiasSlopeScale: 1,
		},
	});
}

function create_stream_buffer(size: number) {
	return device.createBuffer({ size, usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST });
}

function create_uniform_buffer(slots: number) {
	uniform_buffer = device.createBuffer({
		size: slots * UNIFORM_SLOT_SIZE,
		usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
	});
	uniform_bind_group = device.createBindGroup({
		layout: uniform_layout,
		entries: [{ binding: 0, resource: { buffer: uniform_buffer, size: 64 } }],
	});
}

function get_texture_bind_group(texture: GPUTexture) {
	let bind_group = texture_bind_groups.get(texture);
	if (!bind_group) {
		bind_group = device.createBindGroup({
			layout: texture_layout,
			entries: [
				{ binding: 0, resource: texture.createView() },
				{ binding: 1, resource: sampler },
			],
		});
		texture_bind_groups.set(texture, bind_group);
	}
	return bind_group;
}

function create_white_texture() {
	const tex = create_texture(1, 1);
	device.queue.writeTexture({ texture: tex }, new Uint8Array([255, 255, 255, 255]), { bytesPerRow: 4 }, [1, 1]);
	white_tex = tex;
}