typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
208 lines (207 loc) • 7.12 kB
JavaScript
import { $internal } from "../shared/symbols.js";
import { $gpuRepr, $inRepr, $invalidSchemaReason, $memIdent, $repr, $reprPartial, $reprPatch, $validStorageSchema, $validUniformSchema, $validVertexSchema, } from "../shared/symbols.js";
import { alignmentOf } from "./alignmentOf.js";
import { isLooseData, isLooseDecorated, } from "./dataTypes.js";
import { sizeOf } from "./sizeOf.js";
import { isAlignAttrib, isBuiltinAttrib, isDecorated, isSizeAttrib, isWgslData, } from "./wgslTypes.js";
// ----------
// Public API
// ----------
export const builtinNames = [
'vertex_index',
'instance_index',
'clip_distances',
'position',
'front_facing',
'frag_depth',
'primitive_index',
'sample_index',
'sample_mask',
'fragment',
'local_invocation_id',
'local_invocation_index',
'global_invocation_id',
'global_invocation_index',
'workgroup_id',
'workgroup_index',
'num_workgroups',
'subgroup_invocation_id',
'subgroup_size',
'subgroup_id',
'num_subgroups',
];
export function attribute(data, attrib) {
if (isDecorated(data)) {
return new DecoratedImpl(data.inner, [attrib, ...data.attribs]);
}
if (isLooseDecorated(data)) {
return new LooseDecoratedImpl(data.inner, [attrib, ...data.attribs]);
}
if (isLooseData(data)) {
return new LooseDecoratedImpl(data, [attrib]);
}
return new DecoratedImpl(data, [attrib]);
}
/**
* Gives the wrapped data-type a custom byte alignment. Useful in order to
* fulfill uniform alignment requirements.
*
* @example
* const Data = d.struct({
* a: u32, // takes up 4 bytes
* // 12 bytes of padding, because `b` is custom aligned to multiples of 16 bytes
* b: d.align(16, u32),
* });
*
* @param alignment The multiple of bytes this data should align itself to.
* @param data The data-type to align.
*/
export function align(alignment, data) {
return attribute(data, {
[$internal]: true,
type: '@align',
params: [alignment],
// oxlint-disable-next-line typescript/no-explicit-any -- tired of lying to types
});
}
/**
* Adds padding bytes after the wrapped data-type, until the whole value takes up `size` bytes.
*
* @example
* const Data = d.struct({
* a: d.size(16, u32), // takes up 16 bytes, instead of 4
* b: u32, // starts at byte 16, because `a` has a custom size
* });
*
* @param size The amount of bytes that should be reserved for this data-type.
* @param data The data-type to wrap.
*/
export function size(size, data) {
return attribute(data, {
[$internal]: true,
type: '@size',
params: [size],
// oxlint-disable-next-line typescript/no-explicit-any -- tired of lying to types
});
}
/**
* Assigns an explicit numeric location to a struct member or a parameter that has this type.
*
* @example
* const VertexOutput = {
* a: d.u32, // has implicit location 0
* b: d.location(5, d.u32),
* c: d.u32, // has implicit location 6
* };
*
* @param location The explicit numeric location.
* @param data The data-type to wrap.
*/
export function location(location, data) {
return attribute(data, {
[$internal]: true,
type: '@location',
params: [location],
// oxlint-disable-next-line typescript/no-explicit-any -- tired of lying to types
});
}
export function interpolate(interpolationType, data) {
return attribute(data, {
[$internal]: true,
type: '@interpolate',
params: [interpolationType],
// oxlint-disable-next-line typescript/no-explicit-any -- tired of lying to types
});
}
/**
* Marks a position built-in output value as invariant in vertex shaders.
* If the data and control flow match for two position outputs in different
* entry points, then the result values are guaranteed to be the same.
*
* Must only be applied to the position built-in value.
*
* @example
* const VertexOutput = {
* pos: d.invariant(d.builtin.position),
* };
*
* @param data The position built-in data-type to mark as invariant.
*/
export function invariant(data) {
// Validate that invariant is only applied to position built-in
if (!isBuiltin(data)) {
throw new Error('The @invariant attribute must only be applied to the position built-in value.');
}
// Find the builtin attribute to check if it's position
const builtinAttrib = isDecorated(data) || isLooseDecorated(data) ? data.attribs.find(isBuiltinAttrib) : undefined;
if (!builtinAttrib || builtinAttrib.params[0] !== 'position') {
throw new Error('The @invariant attribute must only be applied to the position built-in value.');
}
return attribute(data, {
[$internal]: true,
type: '@invariant',
params: [],
// oxlint-disable-next-line typescript/no-explicit-any -- tired of lying to types
});
}
export function isBuiltin(value) {
return ((isDecorated(value) || isLooseDecorated(value)) &&
value.attribs.find(isBuiltinAttrib) !== undefined);
}
export function getAttributesString(field) {
if (!isDecorated(field) && !isLooseDecorated(field)) {
return '';
}
return field.attribs
.map((attrib) => {
if (attrib.params.length === 0) {
return `${attrib.type} `;
}
return `${attrib.type}(${attrib.params.join(', ')}) `;
})
.join('');
}
// --------------
// Implementation
// --------------
class BaseDecoratedImpl {
[$internal] = {};
inner;
attribs;
// ---
constructor(inner, attribs) {
this.inner = inner;
this.attribs = attribs;
const alignAttrib = attribs.find(isAlignAttrib)?.params[0];
const sizeAttrib = attribs.find(isSizeAttrib)?.params[0];
if (alignAttrib !== undefined) {
if (alignAttrib <= 0) {
throw new Error(`Custom data alignment must be a positive number, got: ${alignAttrib}.`);
}
if (Math.log2(alignAttrib) % 1 !== 0) {
throw new Error(`Alignment has to be a power of 2, got: ${alignAttrib}.`);
}
if (isWgslData(this.inner)) {
if (alignAttrib % alignmentOf(this.inner) !== 0) {
throw new Error(`Custom alignment has to be a multiple of the standard data alignment. Got: ${alignAttrib}, expected multiple of: ${alignmentOf(this.inner)}.`);
}
}
}
if (sizeAttrib !== undefined) {
if (sizeAttrib < sizeOf(this.inner)) {
throw new Error(`Custom data size cannot be smaller then the standard data size. Got: ${sizeAttrib}, expected at least: ${sizeOf(this.inner)}.`);
}
if (sizeAttrib <= 0) {
throw new Error(`Custom data size must be a positive number. Got: ${sizeAttrib}.`);
}
}
}
}
class DecoratedImpl extends BaseDecoratedImpl {
[$internal] = {};
type = 'decorated';
}
class LooseDecoratedImpl extends BaseDecoratedImpl {
[$internal] = {};
type = 'loose-decorated';
}