typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
198 lines (197 loc) • 7.24 kB
JavaScript
import { schemaCallWrapper } from "../../data/schemaCallWrapper.js";
import { Void } from "../../data/wgslTypes.js";
import { ExecutionError } from "../../errors.js";
import { provideInsideTgpuFn } from "../../execMode.js";
import { getName, setName } from "../../shared/meta.js";
import { isMarkedInternal } from "../../shared/symbols.js";
import { $getNameForward, $internal, $providing, $resolve } from "../../shared/symbols.js";
import { addArgTypesToExternals, addReturnTypeToExternals } from "../resolve/externals.js";
import { stitch } from "../resolve/stitch.js";
import { isAccessor, isMutableAccessor, } from "../slot/slotTypes.js";
import { dualImpl } from "./dualImpl.js";
import { createFnCore } from "./fnCore.js";
import { stripTemplate } from "./templateUtils.js";
import { comptime } from "./comptime.js";
export function fn(argTypesOrCallback, returnType) {
if (typeof argTypesOrCallback === 'function') {
return createGenericFn(argTypesOrCallback, []);
}
const argTypes = argTypesOrCallback;
const shell = {
[$internal]: true,
argTypes,
returnType: returnType ?? Void,
};
const call = (arg, ...values) => createFn(shell, stripTemplate(arg, ...values));
return Object.assign(call, shell);
}
export function isTgpuFn(value) {
return (isMarkedInternal(value) &&
value?.resourceType === 'function');
}
export function isGenericFn(value) {
return (isMarkedInternal(value) &&
value?.resourceType === 'generic-function');
}
// --------------
// Implementation
// --------------
function stringifyPair([slot, value]) {
return `${getName(slot) ?? '<unnamed>'}=${value}`;
}
function createFn(shell, _implementation) {
let pairs = [];
// Unwrapping generic functions
let implementation;
if (isGenericFn(_implementation)) {
pairs = _implementation[$providing]?.pairs ?? [];
implementation = _implementation[$internal].inner;
}
else {
implementation = _implementation;
}
const core = createFnCore(implementation, 'normal');
const fnBase = {
shell,
resourceType: 'function',
[$internal]: { implementation },
$uses(newExternals) {
core.setExternals('userProvided', newExternals);
return this;
},
[$getNameForward]: core,
$name(label) {
setName(this, label);
return this;
},
with: comptime((slot, value) => {
const s = isAccessor(slot) || isMutableAccessor(slot) ? slot.slot : slot;
return createBoundFunction(fn, [[s, value]]);
}),
[$resolve](ctx) {
if (typeof implementation === 'string') {
addArgTypesToExternals(implementation, shell.argTypes, core);
addReturnTypeToExternals(implementation, shell.returnType, core);
}
return core.resolve(ctx, shell.argTypes, shell.returnType);
},
};
const call = dualImpl({
name: undefined, // the name is forwarded to the core anyway
noComptime: true,
signature: { argTypes: shell.argTypes, returnType: shell.returnType },
normalImpl: (...args) => provideInsideTgpuFn(() => {
try {
if (typeof implementation === 'string') {
throw new Error('Cannot execute on the CPU functions constructed with raw WGSL');
}
const castAndCopiedArgs = args.map((arg, index) => schemaCallWrapper(shell.argTypes[index], arg));
const result = implementation(...castAndCopiedArgs);
// Casting the result to the appropriate schema
return schemaCallWrapper(shell.returnType, result);
}
catch (err) {
if (err instanceof ExecutionError) {
throw err.appendToTrace(fn);
}
throw new ExecutionError(err, [fn]);
}
}),
codegenImpl: (ctx, args) => ctx.withResetIndentLevel(() => stitch `${ctx.resolve(fn).value}(${args})`),
sideEffects: true,
});
const fn = Object.assign(call, fnBase);
Object.defineProperty(fn, 'toString', {
value() {
return `fn:${getName(core) ?? '<unnamed>'}`;
},
});
if (pairs.length > 0) {
return createBoundFunction(fn, pairs);
}
return fn;
}
function createBoundFunction(innerFn, pairs) {
const fnBase = {
resourceType: 'function',
shell: innerFn.shell,
[$internal]: { implementation: innerFn[$internal].implementation },
[$providing]: { inner: innerFn, pairs },
$uses(newExternals) {
innerFn.$uses(newExternals);
return this;
},
$name(label) {
setName(this, label);
return this;
},
with: comptime((slot, value) => {
const s = isAccessor(slot) || isMutableAccessor(slot) ? slot.slot : slot;
return createBoundFunction(innerFn, [...pairs, [s, value]]);
}),
};
const call = dualImpl({
name: undefined, // setting name here would override autonaming
noComptime: true,
signature: {
argTypes: innerFn.shell.argTypes,
returnType: innerFn.shell.returnType,
},
normalImpl: innerFn,
codegenImpl: (ctx, args) => ctx.withResetIndentLevel(() => stitch `${ctx.resolve(fn).value}(${args})`),
sideEffects: true,
});
const fn = Object.assign(call, fnBase);
Object.defineProperty(fn, 'toString', {
value() {
const fnLabel = getName(this) ?? '<unnamed>';
return `fn:${fnLabel}[${pairs.map(stringifyPair).join(', ')}]`;
},
});
const innerName = getName(innerFn);
if (innerName) {
setName(fn, innerName);
}
return fn;
}
function createGenericFn(inner, pairs) {
const fnBase = {
[$internal]: { inner },
resourceType: 'generic-function',
[$providing]: pairs.length > 0 ? { inner, pairs } : undefined,
$name(label) {
setName(this, label);
// Giving `inner` a name if it doesn't have one
if (!getName(inner)) {
setName(inner, label);
}
return this;
},
[$resolve](ctx) {
return ctx.resolve(inner);
},
with(slot, value) {
const s = isAccessor(slot) || isMutableAccessor(slot) ? slot.slot : slot;
return createGenericFn(inner, [...pairs, [s, value]]);
},
};
const call = (...args) => {
return inner(...args);
};
const genericFn = Object.assign(call, fnBase);
// Inheriting name from `inner`, if it exists
const innerName = getName(inner);
if (innerName) {
setName(genericFn, innerName);
}
Object.defineProperty(genericFn, 'toString', {
value() {
const fnLabel = getName(genericFn) ?? '<unnamed>';
if (pairs.length > 0) {
return `fn*:${fnLabel}[${pairs.map(stringifyPair).join(', ')}]`;
}
return `fn*:${fnLabel}`;
},
});
return genericFn;
}