typegpu/core/function/comptime.js

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

import { WgslTypeError } from "../../errors.js";
import { setName } from "../../shared/meta.js";
import { $getNameForward, $gpuCallable, $internal } from "../../shared/symbols.js";
import { coerceToSnippet } from "../../tgsl/generationHelpers.js";
import { isKnownAtComptime, NormalState } from "../../types.js";
export function isComptimeFn(value) {
    return !!value?.[$internal]?.isComptime;
}
/**
 * Creates a version of `func` that can called safely in a TypeGPU function to
 * precompute and inject a value into the final shader code.
 *
 * Note how the function passed into `comptime` doesn't have to be marked with
 * 'use gpu'. That's because the function doesn't execute on the GPU, it gets
 * executed before the shader code gets sent to the GPU.
 *
 * @example
 * ```ts
 * const color = tgpu.comptime((int: number) => {
 *   const r = (int >> 16) & 0xff;
 *   const g = (int >> 8) & 0xff;
 *   const b = int & 0xff;
 *   return d.vec3f(r / 255, g / 255, b / 255);
 * });
 *
 * const material = (diffuse: d.v3f): d.v3f => {
 *   'use gpu';
 *   const albedo = color(0xff00ff);
 *   return albedo.mul(diffuse);
 * };
 * ```
 */
export function comptime(func) {
    const impl = ((...args) => {
        return func(...args);
    });
    impl.toString = () => 'comptime';
    impl[$getNameForward] = func;
    impl[$gpuCallable] = {
        call(ctx, args) {
            if (!args.every((s) => isKnownAtComptime(s))) {
                throw new WgslTypeError(`Called comptime function with runtime-known values: ${args
                    .filter((s) => !isKnownAtComptime(s))
                    .map((s) => `'${s.value}'`)
                    .join(', ')}`);
            }
            ctx.pushMode(new NormalState());
            try {
                return coerceToSnippet(func(...args.map((s) => s.value)));
            }
            finally {
                ctx.popMode();
            }
        },
    };
    impl.$name = (label) => {
        setName(func, label);
        return impl;
    };
    Object.defineProperty(impl, $internal, {
        value: { isComptime: true },
    });
    return impl;
}