@multiplekex/shallot 0.2.3 → 0.2.5

package/src/standard/raster/index.ts

@@ -0,0 +1,517 @@
+import type { Plugin, State, System } from "../../core";
+import { resource } from "../../core";
+import { Compute } from "../compute";
+import type { ComputeNode, ExecutionContext } from "../compute";
+import { WorldTransform } from "../transforms";
+import { Camera, getClearColor, hasSkyComponent } from "../render/camera";
+import {
+    Mesh,
+    collectBatches,
+    updateBatches,
+    MAX_BATCH_SLOTS,
+    type Batch,
+    type MeshBuffers,
+    type BatchEntities,
+} from "../render/mesh";
+import { Surface, SurfaceType, getDefaultAllSurfaces } from "../render/surface";
+import type { SurfaceData } from "../render/surface";
+import { SCENE_STRUCT_WGSL, SKY_STRUCT_WGSL, WGSL_STRUCTS } from "../render/surface/structs";
+import {
+    compileVertexBody,
+    WGSL_LIGHTING_CALC,
+    SKY_WGSL,
+    SKY_DIR_WGSL,
+    NOISE_WGSL,
+} from "../render/surface/shaders";
+import { COLOR_FORMAT } from "../render/scene";
+import { createIndirectBuffer, INDIRECT_SIZE } from "../render/indirect";
+import { Render, RenderPlugin } from "../render";
+
+export { createRasterPipeline, createSkyPipeline, executeRasterPass };
+export type { RasterPassConfig };
+
+interface RasterState {
+    indirect: GPUBuffer;
+    batches: (Batch | null)[];
+    batchEntities: BatchEntities;
+    buffers: Map<number, MeshBuffers>;
+}
+
+export const RasterResource = resource<RasterState>("raster");
+
+const RasterUploadSystem: System = {
+    group: "draw",
+
+    update(state: State) {
+        const render = Render.from(state);
+        const raster = RasterResource.from(state);
+        const compute = Compute.from(state);
+        if (!render || !raster || !compute) return;
+
+        const { device } = compute;
+        const meshEntities = state.query([Mesh, WorldTransform]);
+
+        collectBatches(
+            meshEntities,
+            (eid) => Surface.type[eid] ?? SurfaceType.Default,
+            raster.batchEntities
+        );
+        updateBatches(device, raster.batchEntities, raster, raster.indirect);
+    },
+};
+
+function compileSurfaceVariant(id: number, data: SurfaceData): string {
+    const vertexBody = compileVertexBody(data.vertex);
+    const fragmentBody = data.fragment ?? "";
+    const lit = data.lit !== false;
+
+    return `
+fn userVertexTransform_${id}(worldPos: vec3<f32>, normal: vec3<f32>, eid: u32) -> vec3<f32> {
+    ${vertexBody}
+}
+
+fn userFragment_${id}(surface: ptr<function, SurfaceData>, position: vec4<f32>) {
+    ${fragmentBody}
+}
+
+fn applyLighting_${id}(surface: SurfaceData) -> vec3<f32> {
+    ${
+        lit
+            ? `${WGSL_LIGHTING_CALC}
+    return surface.baseColor * lighting + surface.emission;`
+            : "return surface.baseColor;"
+    }
+}
+`;
+}
+
+function compileDispatchFunctions(surfaceCount: number): string {
+    const vertexCases = Array.from(
+        { length: surfaceCount },
+        (_, i) => `        case ${i}u: { return userVertexTransform_${i}(worldPos, normal, eid); }`
+    ).join("\n");
+
+    const fragmentCases = Array.from(
+        { length: surfaceCount },
+        (_, i) => `        case ${i}u: { userFragment_${i}(surface, position); }`
+    ).join("\n");
+
+    const lightingCases = Array.from(
+        { length: surfaceCount },
+        (_, i) => `        case ${i}u: { return applyLighting_${i}(surface); }`
+    ).join("\n");
+
+    return `
+fn dispatchVertexTransform(surfaceId: u32, worldPos: vec3<f32>, normal: vec3<f32>, eid: u32) -> vec3<f32> {
+    switch surfaceId {
+${vertexCases}
+        default: { return userVertexTransform_0(worldPos, normal, eid); }
+    }
+}
+
+fn dispatchFragment(surfaceId: u32, surface: ptr<function, SurfaceData>, position: vec4<f32>) {
+    switch surfaceId {
+${fragmentCases}
+        default: { userFragment_0(surface, position); }
+    }
+}
+
+fn dispatchLighting(surfaceId: u32, surface: SurfaceData) -> vec3<f32> {
+    switch surfaceId {
+${lightingCases}
+        default: { return applyLighting_0(surface); }
+    }
+}
+`;
+}
+
+function compileRasterShader(surfaces: SurfaceData[]): string {
+    const surfaceVariants = surfaces.map((s, i) => compileSurfaceVariant(i, s)).join("\n");
+    const dispatchFunctions = compileDispatchFunctions(surfaces.length);
+
+    return /* wgsl */ `
+${WGSL_STRUCTS}
+${NOISE_WGSL}
+
+${surfaceVariants}
+${dispatchFunctions}
+
+@vertex
+fn vs(input: VertexInput) -> VertexOutput {
+    let eid = entityIds[input.instance];
+    let world = matrices[eid];
+    let scaledPos = input.position * sizes[eid].xyz;
+    let baseWorldPos = (world * vec4<f32>(scaledPos, 1.0)).xyz;
+    let worldNormal = normalize((world * vec4<f32>(input.normal, 0.0)).xyz);
+    let d = data[eid];
+    let surfaceId = d.flags & 0xFFu;
+    let finalWorldPos = dispatchVertexTransform(surfaceId, baseWorldPos, worldNormal, eid);
+
+    var output: VertexOutput;
+    output.position = scene.viewProj * vec4<f32>(finalWorldPos, 1.0);
+    output.color = d.baseColor;
+    output.worldNormal = worldNormal;
+    output.entityId = eid;
+    output.worldPos = finalWorldPos;
+    return output;
+}
+
+@fragment
+fn fs(input: VertexOutput) -> FragmentOutput {
+    let eid = input.entityId;
+    let d = data[eid];
+    let surfaceId = d.flags & 0xFFu;
+
+    var surface: SurfaceData;
+    surface.baseColor = input.color.rgb;
+    surface.roughness = d.pbr.x;
+    surface.metallic = d.pbr.y;
+    surface.emission = d.emission.rgb * d.emission.a;
+    surface.normal = normalize(input.worldNormal);
+    surface.worldPos = input.worldPos;
+
+    dispatchFragment(surfaceId, &surface, input.position);
+
+    let litColor = dispatchLighting(surfaceId, surface);
+
+    var output: FragmentOutput;
+    output.color = vec4<f32>(litColor, input.color.a);
+    output.entityId = input.entityId;
+    return output;
+}
+`;
+}
+
+async function createRasterPipeline(
+    device: GPUDevice,
+    surfaces: SurfaceData[],
+    colorFormat: GPUTextureFormat
+): Promise<GPURenderPipeline> {
+    const code = compileRasterShader(surfaces);
+    const module = device.createShaderModule({ code });
+
+    return device.createRenderPipelineAsync({
+        layout: "auto",
+        vertex: {
+            module,
+            entryPoint: "vs",
+            buffers: [
+                {
+                    arrayStride: 24,
+                    attributes: [
+                        { shaderLocation: 0, offset: 0, format: "float32x3" },
+                        { shaderLocation: 1, offset: 12, format: "float32x3" },
+                    ],
+                },
+            ],
+        },
+        fragment: {
+            module,
+            entryPoint: "fs",
+            targets: [{ format: colorFormat }, { format: "r32uint" }],
+        },
+        depthStencil: {
+            format: "depth24plus",
+            depthWriteEnabled: true,
+            depthCompare: "less",
+        },
+        primitive: {
+            topology: "triangle-list",
+            cullMode: "back",
+        },
+    });
+}
+
+function compileSkyShader(): string {
+    return /* wgsl */ `
+${SCENE_STRUCT_WGSL}
+${SKY_STRUCT_WGSL}
+
+@group(0) @binding(0) var<uniform> scene: Scene;
+@group(0) @binding(1) var<uniform> sky: Sky;
+
+${SKY_DIR_WGSL}
+${NOISE_WGSL}
+${SKY_WGSL}
+
+struct VertexOutput {
+    @builtin(position) position: vec4<f32>,
+    @location(0) uv: vec2<f32>,
+}
+
+@vertex
+fn vs(@builtin(vertex_index) vertexIndex: u32) -> VertexOutput {
+    var positions = array<vec2<f32>, 3>(
+        vec2(-1.0, -1.0),
+        vec2(3.0, -1.0),
+        vec2(-1.0, 3.0)
+    );
+    var output: VertexOutput;
+    output.position = vec4(positions[vertexIndex], 0.0, 1.0);
+    output.uv = (positions[vertexIndex] + 1.0) * 0.5;
+    output.uv.y = 1.0 - output.uv.y;
+    return output;
+}
+
+@fragment
+fn fs(input: VertexOutput) -> @location(0) vec4<f32> {
+    let dir = computeSkyDir(input.uv.x, input.uv.y);
+    let color = sampleSky(dir);
+    return vec4(color, 1.0);
+}
+`;
+}
+
+async function createSkyPipeline(
+    device: GPUDevice,
+    colorFormat: GPUTextureFormat
+): Promise<GPURenderPipeline> {
+    const code = compileSkyShader();
+    const module = device.createShaderModule({ code });
+
+    return device.createRenderPipelineAsync({
+        layout: "auto",
+        vertex: { module, entryPoint: "vs" },
+        fragment: {
+            module,
+            entryPoint: "fs",
+            targets: [{ format: colorFormat }],
+        },
+        depthStencil: {
+            format: "depth24plus",
+            depthWriteEnabled: false,
+            depthCompare: "always",
+        },
+        primitive: { topology: "triangle-list" },
+    });
+}
+
+interface RasterPassConfig {
+    scene: GPUBuffer;
+    sky: GPUBuffer;
+    data: GPUBuffer;
+    matrices: GPUBuffer;
+    sizes: GPUBuffer;
+    indirect: GPUBuffer;
+    rasterPipeline: GPURenderPipeline;
+    skyPipeline: GPURenderPipeline | null;
+    getClearColor: () => { r: number; g: number; b: number };
+    getSky: () => boolean;
+    batches: () => (Batch | null)[];
+    skyBindGroup?: GPUBindGroup | null;
+    batchBindGroups?: GPUBindGroup[];
+}
+
+function executeRasterPass(ctx: ExecutionContext, config: RasterPassConfig): void {
+    const { device, encoder } = ctx;
+
+    const colorView = ctx.getTextureView("color");
+    const eidView = ctx.getTextureView("eid");
+    const zView = ctx.getTextureView("z");
+
+    if (!colorView || !eidView || !zView) {
+        return;
+    }
+
+    const batches = config.batches();
+    const colorTexture = ctx.getTexture("color");
+    if (!colorTexture) return;
+
+    const clearColor = config.getClearColor();
+    const hasSky = config.getSky() && config.skyPipeline;
+
+    if (hasSky) {
+        const skyPass = encoder.beginRenderPass({
+            colorAttachments: [
+                {
+                    view: colorView as GPUTextureView,
+                    clearValue: { r: clearColor.r, g: clearColor.g, b: clearColor.b, a: 1 },
+                    loadOp: "clear" as const,
+                    storeOp: "store" as const,
+                },
+            ],
+            depthStencilAttachment: {
+                view: zView as GPUTextureView,
+                depthClearValue: 1.0,
+                depthLoadOp: "clear" as const,
+                depthStoreOp: "store" as const,
+            },
+        });
+
+        if (!config.skyBindGroup) {
+            config.skyBindGroup = device.createBindGroup({
+                layout: config.skyPipeline!.getBindGroupLayout(0),
+                entries: [
+                    { binding: 0, resource: { buffer: config.scene } },
+                    { binding: 1, resource: { buffer: config.sky } },
+                ],
+            });
+        }
+
+        skyPass.setPipeline(config.skyPipeline!);
+        skyPass.setBindGroup(0, config.skyBindGroup);
+        skyPass.draw(3);
+        skyPass.end();
+    }
+
+    const pass = encoder.beginRenderPass({
+        colorAttachments: [
+            {
+                view: colorView as GPUTextureView,
+                clearValue: { r: clearColor.r, g: clearColor.g, b: clearColor.b, a: 1 },
+                loadOp: hasSky ? ("load" as const) : ("clear" as const),
+                storeOp: "store" as const,
+            },
+            {
+                view: eidView as GPUTextureView,
+                clearValue: { r: 0, g: 0, b: 0, a: 0 },
+                loadOp: "clear" as const,
+                storeOp: "store" as const,
+            },
+        ],
+        depthStencilAttachment: {
+            view: zView as GPUTextureView,
+            depthClearValue: 1.0,
+            depthLoadOp: hasSky ? ("load" as const) : ("clear" as const),
+            depthStoreOp: "store" as const,
+        },
+    });
+
+    pass.setPipeline(config.rasterPipeline);
+
+    if (!config.batchBindGroups) {
+        config.batchBindGroups = [];
+    }
+
+    for (let i = 0; i < batches.length; i++) {
+        const batch = batches[i];
+        if (!batch || batch.count === 0) continue;
+
+        if (!config.batchBindGroups[i]) {
+            config.batchBindGroups[i] = device.createBindGroup({
+                layout: config.rasterPipeline.getBindGroupLayout(0),
+                entries: [
+                    { binding: 0, resource: { buffer: config.scene } },
+                    { binding: 1, resource: { buffer: batch.entityIds } },
+                    { binding: 2, resource: { buffer: config.matrices } },
+                    { binding: 3, resource: { buffer: config.sizes } },
+                    { binding: 4, resource: { buffer: config.data } },
+                ],
+            });
+        }
+
+        pass.setBindGroup(0, config.batchBindGroups[i]);
+        pass.setVertexBuffer(0, batch.buffers.vertex);
+        pass.setIndexBuffer(batch.buffers.index, "uint16");
+        pass.drawIndexedIndirect(config.indirect, i * INDIRECT_SIZE);
+    }
+
+    pass.end();
+}
+
+interface RasterForwardConfig {
+    scene: GPUBuffer;
+    sky: GPUBuffer;
+    data: GPUBuffer;
+    indirect: GPUBuffer;
+    matrices: GPUBuffer;
+    sizes: GPUBuffer;
+    getSurfaces: () => SurfaceData[];
+    getClearColor: () => { r: number; g: number; b: number };
+    getSky: () => boolean;
+    batches: () => (Batch | null)[];
+}
+
+function createRasterForwardNode(config: RasterForwardConfig): ComputeNode {
+    let rasterPipeline: GPURenderPipeline | null = null;
+    let skyPipeline: GPURenderPipeline | null = null;
+    let skyBindGroup: GPUBindGroup | null = null;
+    let batchBindGroups: GPUBindGroup[] = [];
+
+    return {
+        id: "forward",
+        inputs: [{ id: "data", access: "read" }],
+        outputs: [
+            { id: "color", access: "write" },
+            { id: "eid", access: "write" },
+        ],
+
+        async prepare(device: GPUDevice) {
+            const surfaces = config.getSurfaces();
+            rasterPipeline = await createRasterPipeline(device, surfaces, COLOR_FORMAT);
+            skyPipeline = await createSkyPipeline(device, COLOR_FORMAT);
+        },
+
+        execute(ctx: ExecutionContext) {
+            if (!rasterPipeline) return;
+            const passConfig: RasterPassConfig = {
+                scene: config.scene,
+                sky: config.sky,
+                data: config.data,
+                matrices: config.matrices,
+                sizes: config.sizes,
+                indirect: config.indirect,
+                rasterPipeline,
+                skyPipeline,
+                getClearColor: config.getClearColor,
+                getSky: config.getSky,
+                batches: config.batches,
+                skyBindGroup,
+                batchBindGroups,
+            };
+            executeRasterPass(ctx, passConfig);
+            skyBindGroup = passConfig.skyBindGroup ?? null;
+            batchBindGroups = passConfig.batchBindGroups ?? [];
+        },
+    };
+}
+
+export const RasterPlugin: Plugin = {
+    systems: [RasterUploadSystem],
+    dependencies: [RenderPlugin],
+
+    async initialize(state: State) {
+        const compute = Compute.from(state);
+        const render = Render.from(state);
+        if (!compute || !render) return;
+
+        const { device } = compute;
+
+        const rasterState: RasterState = {
+            indirect: createIndirectBuffer(device, MAX_BATCH_SLOTS),
+            batches: Array(MAX_BATCH_SLOTS).fill(null),
+            batchEntities: Array(MAX_BATCH_SLOTS).fill(null),
+            buffers: new Map(),
+        };
+
+        state.setResource(RasterResource, rasterState);
+
+        const getActiveClearColor = () => {
+            for (const eid of state.query([Camera])) {
+                if (Camera.active[eid]) return getClearColor(eid);
+            }
+            return { r: 0, g: 0, b: 0 };
+        };
+
+        const getSky = () => {
+            for (const eid of state.query([Camera])) {
+                if (Camera.active[eid]) return hasSkyComponent(state, eid);
+            }
+            return false;
+        };
+
+        const forwardNode = createRasterForwardNode({
+            scene: render.scene,
+            sky: render.sky,
+            data: render.data,
+            indirect: rasterState.indirect,
+            matrices: render.matrices,
+            sizes: render.sizes,
+            getSurfaces: getDefaultAllSurfaces,
+            getClearColor: getActiveClearColor,
+            getSky,
+            batches: () => rasterState.batches,
+        });
+        compute.graph.add(forwardNode);
+    },
+};

package/src/standard/{render → raytracing}/bvh/blas.ts

@@ -1,6 +1,6 @@
 import type { Vec3, BVHNode, MortonPair, AABB } from "./structs";
 import { LEAF_FLAG, isLeaf, leafIndex } from "./structs";
-import type { MeshData } from "../mesh";
+import type { MeshData } from "../../render/mesh";
 
 export interface BLASTriangle {
     v0: Vec3;
@@ -503,7 +503,7 @@ export function buildShapeBLAS(triangles: BLASTriangle[]): BLASData {
     };
 }
 
-export function validateBLASBounds(blas: BLASData, triangles: BLASTriangle[]): boolean {
+function validateBLASBounds(blas: BLASData, triangles: BLASTriangle[]): boolean {
     if (blas.triCount === 0) return true;
 
     const epsilon = 1e-5;
@@ -530,7 +530,7 @@ export function validateBLASBounds(blas: BLASData, triangles: BLASTriangle[]): b
     return true;
 }
 
-export function validateBLASNodes(blas: BLASData): { valid: boolean; errors: string[] } {
+function validateBLASNodes(blas: BLASData): { valid: boolean; errors: string[] } {
     const errors: string[] = [];
 
     if (blas.triCount <= 1) {

package/src/standard/{render → raytracing}/bvh/tlas.ts

@@ -721,6 +721,7 @@ export interface TLASConfig {
     instanceCount: GPUBuffer;
     tlas: TLASBuffers;
     getEntityCount: () => number;
+    getRaytracing?: () => boolean;
 }
 
 export function createTLASNode(config: TLASConfig): ComputeNode {
@@ -836,6 +837,8 @@ export function createTLASNode(config: TLASConfig): ComputeNode {
         },
 
         execute(ctx: ExecutionContext) {
+            if (config.getRaytracing && !config.getRaytracing()) return;
+
             const { device, encoder } = ctx;
 
             const workgroups = Math.ceil(MAX_ENTITIES / WORKGROUP_SIZE);

package/src/standard/{render → raytracing}/depth.ts

@@ -1,11 +1,11 @@
 import type { ComputeNode, ExecutionContext } from "../compute";
-import { SCENE_STRUCT_WGSL } from "./shaders";
+import { SCENE_STRUCT_WGSL } from "../render/surface/structs";
 
 const depthConvertShader = /* wgsl */ `
 ${SCENE_STRUCT_WGSL}
 
 @group(0) @binding(0) var<uniform> scene: Scene;
-@group(0) @binding(1) var linearDepthTex: texture_2d<f32>;
+@group(0) @binding(1) var depthTex: texture_2d<f32>;
 
 struct VertexOutput {
     @builtin(position) position: vec4<f32>,
@@ -29,7 +29,7 @@ fn vs(@builtin(vertex_index) vertexIndex: u32) -> VertexOutput {
 @fragment
 fn fs(input: VertexOutput) -> @builtin(frag_depth) f32 {
     let coords = vec2<i32>(input.position.xy);
-    let t = textureLoad(linearDepthTex, coords, 0).r;
+    let t = textureLoad(depthTex, coords, 0).r;
 
     let near = scene.near;
     let far = scene.far;
@@ -59,8 +59,8 @@ export function createDepthConvertNode(config: DepthConvertConfig): ComputeNode
 
     return {
         id: "depth-convert",
-        inputs: [{ id: "linear-depth", access: "read" }],
-        outputs: [{ id: "depth", access: "write" }],
+        inputs: [{ id: "depth", access: "read" }],
+        outputs: [{ id: "z", access: "write" }],
 
         async prepare(device: GPUDevice) {
             const module = device.createShaderModule({ code: depthConvertShader });
@@ -84,10 +84,10 @@ export function createDepthConvertNode(config: DepthConvertConfig): ComputeNode
 
             const { device, encoder } = ctx;
 
-            const linearDepthView = ctx.getTextureView("linear-depth");
             const depthView = ctx.getTextureView("depth");
+            const zView = ctx.getTextureView("z");
 
-            if (!linearDepthView || !depthView) {
+            if (!depthView || !zView) {
                 return;
             }
 
@@ -95,14 +95,14 @@ export function createDepthConvertNode(config: DepthConvertConfig): ComputeNode
                 layout: pipeline!.getBindGroupLayout(0),
                 entries: [
                     { binding: 0, resource: { buffer: config.scene } },
-                    { binding: 1, resource: linearDepthView },
+                    { binding: 1, resource: depthView },
                 ],
             });
 
             const pass = encoder.beginRenderPass({
                 colorAttachments: [],
                 depthStencilAttachment: {
-                    view: depthView,
+                    view: zView,
                     depthClearValue: 1.0,
                     depthLoadOp: "clear",
                     depthStoreOp: "store",