typegpu/core/function/createCallableSchema.js

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

import { noSideEffects, snip, } from "../../data/snippet.js";
import { isPtr } from "../../data/wgslTypes.js";
import { setName } from "../../shared/meta.js";
import { $gpuCallable } from "../../shared/symbols.js";
import { tryConvertSnippet } from "../../tgsl/conversion.js";
import { isKnownAtComptime, NormalState } from "../../types.js";
export function callableSchema(options) {
    const impl = ((...args) => {
        return options.normalImpl(...args);
    });
    setName(impl, options.name);
    impl.toString = () => options.name;
    impl[$gpuCallable] = {
        get strictSignature() {
            return undefined;
        },
        call(ctx, args) {
            const argTypes = options.argTypes(...args.map((s) => {
                // Dereference implicit pointers
                if (isPtr(s.dataType) && s.dataType.implicit) {
                    return s.dataType.inner;
                }
                return s.dataType;
            }));
            const converted = args.map((s, idx) => {
                const argType = argTypes[idx];
                if (!argType) {
                    throw new Error('Function called with invalid arguments');
                }
                return tryConvertSnippet(ctx, s, argType, false);
            });
            let result;
            if (converted.every((s) => isKnownAtComptime(s))) {
                ctx.pushMode(new NormalState());
                try {
                    result = snip(options.normalImpl(...converted.map((s) => s.value)), options.schema(), 
                    // Functions give up ownership of their return value
                    /* origin */ 'constant');
                }
                finally {
                    ctx.popMode('normal');
                }
            }
            else {
                result = options.codegenImpl(ctx, converted);
            }
            if (!args.some((a) => a.possibleSideEffects)) {
                return noSideEffects(result);
            }
            return result;
        },
    };
    return impl;
}