matrix/js/webgpu/mirrorPass.js

import { structs } from "../../lib/gpu-buffer.js";
import {
	makeComputeTarget,
	makeUniformBuffer,
	loadShader,
	makeBindGroup,
	makePass,
} from "./utils.js";

export default ({ config, device, canvas, cache, cameraTex, cameraAspectRatio, timeBuffer }) => {
	const assets = [loadShader(device, cache, "shaders/wgsl/mirrorPass.wgsl")];

	const linearSampler = device.createSampler({
		magFilter: "linear",
		minFilter: "linear",
	});

	let start;
	const numTouches = 5;
	const touches = Array(numTouches)
		.fill()
		.map((_) => [0, 0, -Infinity, 0]);
	let aspectRatio = 1;

	let index = 0;
	let touchesChanged = true;
	canvas.onmousedown = (e) => {
		const rect = e.srcElement.getBoundingClientRect();
		touches[index][0] = 0 + (e.clientX - rect.x) / rect.width;
		touches[index][1] = 1 - (e.clientY - rect.y) / rect.height;
		touches[index][2] = (performance.now() - start) / 1000;
		index = (index + 1) % numTouches;
		touchesChanged = true;
	};

	let computePipeline;
	let configBuffer;
	let sceneUniforms;
	let sceneBuffer;
	let touchUniforms;
	let touchBuffer;
	let output;
	let screenSize;
	let computeBindGroup;

	const loaded = (async () => {
		const [mirrorShader] = await Promise.all(assets);

		computePipeline = await device.createComputePipelineAsync({
			layout: "auto",
			compute: {
				module: mirrorShader.module,
				entryPoint: "computeMain",
			},
		});

		const mirrorShaderUniforms = structs.from(mirrorShader.code);

		const configUniforms = mirrorShaderUniforms.Config;
		configBuffer = makeUniformBuffer(device, configUniforms, { unused: 0 });

		sceneUniforms = mirrorShaderUniforms.Scene;
		sceneBuffer = makeUniformBuffer(device, sceneUniforms);

		touchUniforms = mirrorShaderUniforms.Touches;
		touchBuffer = makeUniformBuffer(device, touchUniforms);
	})();

	const build = (size, inputs) => {
		output?.destroy();
		output = makeComputeTarget(device, size);
		screenSize = size;
		aspectRatio = size[0] / size[1];
		computeBindGroup = makeBindGroup(device, computePipeline, 0, [
			configBuffer,
			timeBuffer,
			sceneBuffer,
			touchBuffer,
			linearSampler,
			inputs.primary.createView(),
			inputs.bloom.createView(),
			cameraTex.createView(),
			output.createView(),
		]);

		const screenAspectRatio = size[0] / size[1];
		device.queue.writeBuffer(
			sceneBuffer,
			0,
			sceneUniforms.toBuffer({ screenAspectRatio, cameraAspectRatio }),
		);

		return { primary: output };
	};

	const run = (encoder, shouldRender) => {
		if (!shouldRender) {
			return;
		}

		if (touchesChanged) {
			touchesChanged = false;
			device.queue.writeBuffer(touchBuffer, 0, touchUniforms.toBuffer({ touches }));
		}

		const computePass = encoder.beginComputePass();
		computePass.setPipeline(computePipeline);
		computePass.setBindGroup(0, computeBindGroup);
		computePass.dispatchWorkgroups(Math.ceil(screenSize[0] / 32), screenSize[1], 1);
		computePass.end();
	};

	start = performance.now();

	return makePass("Mirror", loaded, build, run);
};