@multiplekex/shallot 0.2.5 → 0.3.0

package/src/standard/raster/index.ts

@@ -1,473 +1,52 @@
-import type { Plugin, State, System } from "../../core";
-import { resource } from "../../core";
+import type { Plugin, State } from "../../core";
 import { Compute } from "../compute";
-import type { ComputeNode, ExecutionContext } from "../compute";
+import type { ComputeNode } 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 { Camera, Shadows, getClearColor, hasSkyComponent } from "../render/camera";
+import { DirectionalLight, normalizeDirection } from "../render/light";
+import { Mesh, MeshColors } 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 { Render, RenderPlugin, registerShadowSettings } from "../render";
+import { setTraits } from "../../core/component";
 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);
-    },
+    createShadowAtlas,
+    createShadowBuffer,
+    createShadowUploadNode,
+    createShadowForwardNode,
+} from "./shadow";
+import { createBatch, uploadBatch, type Batch } from "./batch";
+import { createRasterForwardNode } from "./forward";
+
+export { compileRasterShader, compileCompositeShader } from "./forward";
+
+export const RasterShadowSettings = {
+    softness: [] as number[],
+    samples: [] as number[],
+    distance: [] as number[],
 };
 
-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;
+setTraits(RasterShadowSettings, {
+    defaults: () => ({
+        softness: 0,
+        samples: 1,
+        distance: 100,
+    }),
+});
 
-    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);
-        },
+import { resource } from "../../core";
 
-        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 RasterResource = resource<{
+    batching: Batch;
+    shadowAtlas: GPUTexture;
+    shadowBuffer: GPUBuffer;
+    rtRendered: boolean;
+}>("raster");
 
 export const RasterPlugin: Plugin = {
-    systems: [RasterUploadSystem],
+    systems: [],
+    components: {
+        RasterShadowSettings,
+    },
     dependencies: [RenderPlugin],
 
     async initialize(state: State) {
@@ -475,42 +54,137 @@ export const RasterPlugin: Plugin = {
         const render = Render.from(state);
         if (!compute || !render) return;
 
+        registerShadowSettings({
+            softness: RasterShadowSettings.softness,
+            samples: RasterShadowSettings.samples,
+            isActive: (s, eid) => s.hasComponent(eid, RasterShadowSettings),
+        });
+
         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(),
+        const shadowAtlas = createShadowAtlas(device);
+        const shadowBuffer = createShadowBuffer(device);
+        const batching = createBatch(device);
+
+        const rasterState = {
+            batching,
+            shadowAtlas,
+            shadowBuffer,
+            rtRendered: false,
         };
 
         state.setResource(RasterResource, rasterState);
 
-        const getActiveClearColor = () => {
+        const getActiveCamera = () => {
             for (const eid of state.query([Camera])) {
-                if (Camera.active[eid]) return getClearColor(eid);
+                if (Camera.active[eid]) return eid;
             }
+            return -1;
+        };
+
+        const getActiveClearColor = () => {
+            const eid = getActiveCamera();
+            if (eid >= 0) 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);
+            const eid = getActiveCamera();
+            return eid >= 0 && hasSkyComponent(state, eid);
+        };
+
+        const getShadowsEnabled = () => {
+            const eid = getActiveCamera();
+            return eid >= 0 && state.hasComponent(eid, Shadows);
+        };
+
+        const getShadowDistance = () => {
+            const eid = getActiveCamera();
+            if (eid >= 0 && state.hasComponent(eid, RasterShadowSettings)) {
+                return RasterShadowSettings.distance[eid];
+            }
+            return 100;
+        };
+
+        const getCameraData = () => {
+            const eid = getActiveCamera();
+            if (eid >= 0) {
+                return {
+                    world: WorldTransform.data.subarray(eid * 16, eid * 16 + 16),
+                    fov: Camera.fov[eid],
+                    near: Camera.near[eid],
+                    far: Camera.far[eid],
+                    width: render.width,
+                    height: render.height,
+                };
             }
-            return false;
+            return null;
         };
 
+        const getLightDir = (): [number, number, number] => {
+            for (const eid of state.query([DirectionalLight])) {
+                const [dx, dy, dz] = normalizeDirection(
+                    DirectionalLight.directionX[eid],
+                    DirectionalLight.directionY[eid],
+                    DirectionalLight.directionZ[eid]
+                );
+                return [dx, dy, dz];
+            }
+            return [-0.5, -1.0, -0.5];
+        };
+
+        const getOpacity = (eid: number) => MeshColors.data[eid * 4 + 3];
+
+        const batchUploadNode: ComputeNode = {
+            id: "batch-upload",
+            inputs: [{ id: "data", access: "read" }],
+            outputs: [{ id: "batched", access: "write" }],
+            execute() {
+                const meshEntities = state.query([Mesh, WorldTransform]);
+                uploadBatch(
+                    device,
+                    meshEntities,
+                    (eid) => Surface.type[eid] ?? SurfaceType.Default,
+                    getOpacity,
+                    rasterState.batching
+                );
+            },
+        };
+        compute.graph.add(batchUploadNode);
+
+        const shadowUploadNode = createShadowUploadNode({
+            shadow: shadowBuffer,
+            getCameraData,
+            getLightDir,
+            shadowsEnabled: getShadowsEnabled,
+            getShadowDistance,
+        });
+        compute.graph.add(shadowUploadNode);
+
+        const shadowNodes = createShadowForwardNode({
+            shadow: shadowBuffer,
+            matrices: render.matrices,
+            sizes: render.sizes,
+            batching: rasterState.batching,
+            atlas: shadowAtlas,
+            shadowsEnabled: getShadowsEnabled,
+        });
+        for (const node of shadowNodes) compute.graph.add(node);
+
         const forwardNode = createRasterForwardNode({
             scene: render.scene,
             sky: render.sky,
             data: render.data,
-            indirect: rasterState.indirect,
             matrices: render.matrices,
             sizes: render.sizes,
+            batching: rasterState.batching,
+            shadowBuffer,
+            shadowAtlas,
             getSurfaces: getDefaultAllSurfaces,
             getClearColor: getActiveClearColor,
             getSky,
-            batches: () => rasterState.batches,
+            getShadowsEnabled,
+            getRTRendered: () => rasterState.rtRendered,
         });
         compute.graph.add(forwardNode);
     },