typegpu/core/rawCodeSnippet/tgpuRawCodeSnippet.js

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

import { snip } from "../../data/snippet.js";
import { makeDereferenceable } from "../../tgsl/makeDereferenceable.js";
import { makeResolvable } from "../../tgsl/makeResolvable.js";
import { $gpuValueOf, $internal } from "../../shared/symbols.js";
import { replaceExternalsInWgsl } from "../resolve/externals.js";
/**
 * An advanced API that creates a typed shader expression which
 * can be injected into the final shader bundle upon use.
 *
 * @param expression The code snippet that will be injected in place of `foo.$`
 * @param type The type of the expression
 * @param [origin='runtime'] Where the value originates from.
 * @param [possibleSideEffects=true] Whether generating this snippet may produce a WGSL expression with observable side-effects (e.g. calling a barrier, discarding a fragment, or writing to memory).
 *
 * **-- Which origin to choose?**
 *
 * Usually 'runtime' (the default) is a safe bet, but if you're sure that the expression or
 * computation is constant (either a reference to a constant, a numeric literal,
 * or an operation on constants), then pass 'constant' as it might lead to better
 * optimizations.
 *
 * If what the expression is a direct reference to an existing value (e.g. a uniform, a
 * storage binding, ...), then choose from 'uniform', 'mutable', 'readonly', 'workgroup',
 * 'private' or 'handle' depending on the address space of the referred value.
 *
 * @example
 * ```ts
 * // An identifier that we know will be in the
 * // final shader bundle, but we cannot
 * // refer to it in any other way.
 * const existingGlobal = tgpu['~unstable']
 *   .rawCodeSnippet('EXISTING_GLOBAL', d.f32, 'constant', false);
 *
 * const foo = () => {
 *   'use gpu';
 *   return existingGlobal.$ * 2;
 * };
 *
 * const wgsl = tgpu.resolve([foo]);
 * // fn foo() -> f32 {
 * //   return EXISTING_GLOBAL * 2;
 * // }
 * ```
 */
export function rawCodeSnippet(expression, type, origin = 'runtime', possibleSideEffects = true) {
    return new TgpuRawCodeSnippetImpl(expression, type, origin, possibleSideEffects);
}
// --------------
// Implementation
// --------------
class TgpuRawCodeSnippetImpl {
    dataType;
    origin;
    possibleSideEffects;
    #expression;
    #externals;
    static {
        TgpuRawCodeSnippetImpl.prototype[$internal] = true;
        makeDereferenceable(makeResolvable(TgpuRawCodeSnippetImpl.prototype, {
            asString() {
                return `raw(${String(this.dataType)}): "${this.#expression}"`;
            },
            resolve(ctx) {
                const replacedExpression = replaceExternalsInWgsl(ctx, this.#externals ?? {}, this.#expression);
                return snip(replacedExpression, this.dataType, this.origin, this.possibleSideEffects);
            },
        }), {
            codegenMode: {
                getBaseSnippet(trackingProxy) {
                    return snip(trackingProxy, this.dataType, this.origin, this.possibleSideEffects);
                },
            },
            normalMode: {
                get() {
                    throw new Error('Raw code snippets can only be used on the GPU.');
                },
            },
        });
    }
    constructor(expression, type, origin, possibleSideEffects) {
        this.dataType = type;
        this.origin = origin;
        this.possibleSideEffects = possibleSideEffects;
        this.#expression = expression;
    }
    $uses(dependencyMap) {
        if (this.#externals !== undefined) {
            throw new Error("Cannot call '$uses' multiple times. If you wish to override dependencies, use slots or accessors instead.");
        }
        this.#externals = dependencyMap;
        return this;
    }
}