typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
210 lines (209 loc) • 7.71 kB
JavaScript
import * as d from "../data/index.js";
import { isWgslStorageTexture, isWgslTexture } from "../data/texture.js";
import { assertExhaustive } from "../shared/utilityTypes.js";
// Maps schema singletons (e.g. `d.f32`) back to their `d` export names
let leafKeys;
function getDataSchemaKey(schema) {
if (!leafKeys) {
leafKeys = new Map();
for (const [key, value] of Object.entries(d)) {
if ((typeof value === 'object' || typeof value === 'function') && value !== null) {
leafKeys.set(value, key);
}
}
}
return leafKeys.get(schema);
}
let builtinsByName;
function getBuiltinByName(value) {
if (!builtinsByName) {
builtinsByName = new Map();
for (const candidate of Object.values(d.builtin)) {
if (!d.isDecorated(candidate) && !d.isLooseDecorated(candidate)) {
continue;
}
const builtin = candidate.attribs.find(d.isBuiltinAttrib);
if (builtin) {
builtinsByName.set(builtin.params[0], candidate);
}
}
}
const builtin = builtinsByName.get(value);
if (!builtin) {
throw new Error(`TypeGPU builtin '${value}' could not be reconstructed.`);
}
return builtin;
}
function serializeAttrib(attrib) {
if (d.isAlignAttrib(attrib)) {
return { type: 'align', value: attrib.params[0] };
}
if (d.isSizeAttrib(attrib)) {
return { type: 'size', value: attrib.params[0] };
}
if (d.isLocationAttrib(attrib)) {
return { type: 'location', value: attrib.params[0] };
}
if (d.isInterpolateAttrib(attrib)) {
return { type: 'interpolate', value: attrib.params[0] };
}
if (d.isBuiltinAttrib(attrib)) {
return { type: 'builtin', value: attrib.params[0] };
}
if (d.isInvariantAttrib(attrib)) {
return { type: 'invariant' };
}
throw new Error('This TypeGPU schema decorator cannot be serialized yet.');
}
function applyAttrib(schema, attrib) {
if (attrib.type === 'align') {
return d.align(attrib.value, schema);
}
if (attrib.type === 'size') {
return d.size(attrib.value, schema);
}
if (attrib.type === 'location') {
return d.location(attrib.value, schema);
}
if (attrib.type === 'interpolate') {
return d.interpolate(attrib.value, schema);
}
if (attrib.type === 'builtin') {
return getBuiltinByName(attrib.value);
}
if (attrib.type === 'invariant') {
return d.invariant(schema);
}
assertExhaustive(attrib, 'schema.ts#applyAttrib');
}
function serializeProps(propTypes) {
return Object.entries(propTypes).map(([prop, propType]) => [prop, serializeDataSchema(propType)]);
}
function deserializeProps(props) {
return Object.fromEntries(props.map(([prop, schema]) => [prop, deserializeDataSchema(schema)]));
}
const sampledTextureConstructors = {
texture_1d: (sampleType) => d.texture1d(sampleType),
texture_2d: (sampleType) => d.texture2d(sampleType),
texture_2d_array: (sampleType) => d.texture2dArray(sampleType),
texture_3d: (sampleType) => d.texture3d(sampleType),
texture_cube: (sampleType) => d.textureCube(sampleType),
texture_cube_array: (sampleType) => d.textureCubeArray(sampleType),
texture_multisampled_2d: (sampleType) => d.textureMultisampled2d(sampleType),
texture_depth_2d: () => d.textureDepth2d(),
texture_depth_2d_array: () => d.textureDepth2dArray(),
texture_depth_cube: () => d.textureDepthCube(),
texture_depth_cube_array: () => d.textureDepthCubeArray(),
texture_depth_multisampled_2d: () => d.textureDepthMultisampled2d(),
};
const storageTextureConstructors = {
texture_storage_1d: (format, access) => d.textureStorage1d(format, access),
texture_storage_2d: (format, access) => d.textureStorage2d(format, access),
texture_storage_2d_array: (format, access) => d.textureStorage2dArray(format, access),
texture_storage_3d: (format, access) => d.textureStorage3d(format, access),
};
export function serializeDataSchema(schema) {
const key = getDataSchemaKey(schema);
if (key) {
return { type: 'd', key };
}
if (isWgslTexture(schema)) {
return {
type: 'sampled-texture',
kind: schema.type,
sampleType: serializeDataSchema(schema.sampleType),
};
}
if (isWgslStorageTexture(schema)) {
return {
type: 'storage-texture',
kind: schema.type,
format: schema.format,
access: schema.access,
};
}
if (schema.type === 'texture_external') {
return { type: 'external-texture' };
}
if (d.isDecorated(schema) || d.isLooseDecorated(schema)) {
return {
type: 'decorated',
inner: serializeDataSchema(schema.inner),
attribs: schema.attribs.map(serializeAttrib),
};
}
if (d.isAtomic(schema)) {
return { type: 'atomic', inner: serializeDataSchema(schema.inner) };
}
if (d.isWgslArray(schema)) {
return {
type: 'array',
element: serializeDataSchema(schema.elementType),
count: schema.elementCount,
};
}
if (d.isDisarray(schema)) {
return {
type: 'disarray',
element: serializeDataSchema(schema.elementType),
count: schema.elementCount,
};
}
if (d.isWgslStruct(schema)) {
return { type: 'struct', props: serializeProps(schema.propTypes) };
}
if (d.isUnstruct(schema)) {
return { type: 'unstruct', props: serializeProps(schema.propTypes) };
}
throw new Error(`TypeGPU schema '${schema.type}' cannot be serialized yet.`);
}
export function deserializeDataSchema(schema) {
if (schema.type === 'd') {
const leaf = d[schema.key];
if (!leaf) {
throw new Error(`TypeGPU schema 'd.${schema.key}' could not be reconstructed.`);
}
return leaf;
}
if (schema.type === 'array') {
return d.arrayOf(deserializeDataSchema(schema.element), schema.count);
}
if (schema.type === 'disarray') {
return d.disarrayOf(deserializeDataSchema(schema.element), schema.count);
}
if (schema.type === 'struct') {
return d.struct(deserializeProps(schema.props));
}
if (schema.type === 'unstruct') {
return d.unstruct(deserializeProps(schema.props));
}
if (schema.type === 'atomic') {
return d.atomic(deserializeDataSchema(schema.inner));
}
if (schema.type === 'decorated') {
let result = deserializeDataSchema(schema.inner);
for (let i = schema.attribs.length - 1; i >= 0; i--) {
const attrib = schema.attribs[i];
result = applyAttrib(result, attrib);
}
return result;
}
if (schema.type === 'sampled-texture') {
const constructor = sampledTextureConstructors[schema.kind];
if (!constructor) {
throw new Error(`TypeGPU texture schema '${schema.kind}' could not be reconstructed.`);
}
return constructor(deserializeDataSchema(schema.sampleType));
}
if (schema.type === 'external-texture') {
return d.textureExternal();
}
if (schema.type === 'storage-texture') {
const constructor = storageTextureConstructors[schema.kind];
if (!constructor) {
throw new Error(`TypeGPU storage texture schema '${schema.kind}' could not be reconstructed.`);
}
return constructor(schema.format, schema.access);
}
throw new Error('TypeGPU schema payload could not be reconstructed.');
}