typegpu/core/function/ioSchema.js

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

import { location, } from "../../data/attributes.js";
import { isBuiltin } from "../../data/attributes.js";
import { getCustomLocation, isData } from "../../data/dataTypes.js";
import { INTERNAL_createStruct } from "../../data/struct.js";
import { isVoid } from "../../data/wgslTypes.js";
export function withLocations(members, locations = {}) {
    let nextLocation = 0;
    const usedCustomLocations = new Set();
    return Object.fromEntries(Object.entries(members ?? {})
        .map(([key, member]) => {
        const customLocation = getCustomLocation(member);
        if (customLocation !== undefined) {
            if (usedCustomLocations.has(customLocation)) {
                throw new Error('Duplicate custom location attributes found');
            }
            usedCustomLocations.add(customLocation);
        }
        return [key, member];
    })
        .map(([key, member]) => {
        if (isBuiltin(member)) {
            // skipping builtins
            return [key, member];
        }
        if (getCustomLocation(member) !== undefined) {
            // this member is already marked
            return [key, member];
        }
        if (locations[key]) {
            // location has been determined by a previous procedure
            return [key, location(locations[key], member)];
        }
        while (usedCustomLocations.has(nextLocation)) {
            nextLocation++;
        }
        return [key, location(nextLocation++, member)];
    }));
}
export function separateBuiltins(schema, locations = {}) {
    const positionalArgs = [];
    const dataFields = {};
    for (const [key, type] of Object.entries(schema)) {
        if (isBuiltin(type)) {
            positionalArgs.push({ schemaKey: key, type });
        }
        else {
            dataFields[key] = type;
        }
    }
    const dataSchema = Object.keys(dataFields).length > 0
        ? INTERNAL_createStruct(withLocations(dataFields, locations), /* isAbstruct */ false)
        : undefined;
    return { dataSchema, positionalArgs };
}
export function separateAllAsPositional(schema) {
    const withLocs = withLocations(schema);
    const positionalArgs = Object.entries(withLocs).map(([key, type]) => ({ schemaKey: key, type }));
    return { dataSchema: undefined, positionalArgs };
}
export function createIoSchema(layout, locations = {}) {
    return (isData(layout)
        ? isVoid(layout)
            ? layout
            : isBuiltin(layout)
                ? layout
                : getCustomLocation(layout) !== undefined
                    ? layout
                    : location(0, layout)
        : INTERNAL_createStruct(withLocations(layout, locations), 
        /* isAbstruct */ false));
}