typegpu/core/function/tgpuFragmentFn.js

A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.

import { getName, isNamable, setName } from "../../shared/meta.js";
import { $getNameForward, $internal, $resolve } from "../../shared/symbols.js";
import { addReturnTypeToExternals } from "../resolve/externals.js";
import { shaderStageSlot } from "../slot/internalSlots.js";
import { createFnCore } from "./fnCore.js";
import { createIoSchema, separateBuiltins } from "./ioSchema.js";
import { stripTemplate } from "./templateUtils.js";
/**
 * Creates a shell of a typed entry function for the fragment shader stage. Any function
 * that implements this shell can run for each fragment (pixel), allowing the inner code
 * to process information received from the vertex shader stage and builtins to determine
 * the final color of the pixel (many pixels in case of multiple targets).
 *
 * @param options.in
 *  Values computed in the vertex stage and builtins to be made available to functions that implement this shell.
 * @param options.out
 *  A `vec4f`, signaling this function outputs a color for one target, or a record containing colors for multiple targets.
 */
export function fragmentFn(options) {
    const shell = {
        in: options.in,
        out: options.out,
        returnType: createIoSchema(options.out),
        entryPoint: 'fragment',
    };
    const call = (arg, ...values) => createFragmentFn(shell, stripTemplate(arg, ...values));
    return Object.assign(call, shell);
}
export function isTgpuFragmentFn(value) {
    return value?.shell?.entryPoint === 'fragment';
}
// --------------
// Implementation
// --------------
function createFragmentFn(shell, implementation) {
    const core = createFnCore(implementation, 'fragment');
    const outputType = shell.returnType;
    if (typeof implementation === 'string') {
        addReturnTypeToExternals(implementation, outputType, core);
    }
    const result = {
        shell,
        outputType,
        $uses(newExternals) {
            core.setExternals('userProvided', newExternals);
            return this;
        },
        [$internal]: true,
        [$getNameForward]: core,
        $name(newLabel) {
            setName(this, newLabel);
            if (isNamable(outputType)) {
                outputType.$name(`${newLabel}_Output`);
            }
            return this;
        },
        [$resolve](ctx) {
            const entryInput = separateBuiltins(shell.in ?? {}, ctx.varyingLocations ?? {});
            if (entryInput.dataSchema && isNamable(entryInput.dataSchema)) {
                entryInput.dataSchema.$name(`${getName(this) ?? ''}_Input`);
            }
            if (typeof implementation === 'string') {
                core.setExternals('out', { Out: outputType });
            }
            return ctx.withSlots([[shaderStageSlot, 'fragment']], () => core.resolve(ctx, [], shell.returnType, entryInput));
        },
        toString() {
            return `fragmentFn:${getName(core) ?? '<unnamed>'}`;
        },
    };
    return result;
}