typegpu/data/array.js

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

import { comptime } from "../core/function/comptime.js";
import { $internal } from "../shared/symbols.js";
import { schemaCallWrapper } from "./schemaCallWrapper.js";
import { sizeOf } from "./sizeOf.js";
import { isDecorated, isLocationAttrib } from "./wgslTypes.js";
/**
 * Creates an array schema that can be used to construct gpu buffers.
 * Describes arrays with fixed-size length, storing elements of the same type.
 *
 * The only decoration allowed on element types is `d.location`. Decorators like
 * `d.align` and `d.size` cannot be applied directly — wrap them in a struct instead,
 * e.g. `d.arrayOf(d.struct({ value: d.align(16, d.u32) }), n)`.
 *
 * @example
 * const LENGTH = 3;
 * const array = d.arrayOf(d.u32, LENGTH);
 *
 * If `elementCount` is not specified, a partially applied function is returned.
 * @example
 * const array = d.arrayOf(d.vec3f);
 * //    ^? (n: number) => WgslArray<d.Vec3f>
 *
 * @param elementType The type of elements in the array.
 * @param elementCount The number of elements in the array.
 * @throws If `elementType` is decorated with anything other than `d.location`.
 */
export const arrayOf = comptime(((elementType, elementCount) => {
    if (isDecorated(elementType) && !elementType.attribs.every(isLocationAttrib)) {
        throw new Error('Arrays cannot hold decorated types other than @location. Wrap it in a struct instead, e.g. d.arrayOf(d.struct({ value: d.align(16, d.u32) }), n).');
    }
    if (elementCount === undefined) {
        return comptime((count) => cpu_arrayOf(elementType, count));
    }
    return cpu_arrayOf(elementType, elementCount);
})).$name('arrayOf');
// --------------
// Implementation
// --------------
function cpu_arrayOf(elementType, elementCount) {
    // In the schema call, create and return a deep copy
    // by wrapping all the values in `elementType` schema calls.
    const arraySchema = (elements) => {
        if (elements && elements.length !== elementCount) {
            throw new Error(`Array schema of ${elementCount} elements of type ${elementType.type} called with ${elements.length} argument(s).`);
        }
        return Array.from({ length: elementCount }, (_, i) => schemaCallWrapper(elementType, elements?.[i]));
    };
    Object.setPrototypeOf(arraySchema, WgslArrayImpl);
    if (Number.isNaN(sizeOf(elementType))) {
        throw new Error('Cannot nest runtime sized arrays.');
    }
    arraySchema.elementType = elementType;
    if (!Number.isInteger(elementCount) || elementCount < 0) {
        throw new Error(`Cannot create array schema with invalid element count: ${elementCount}.`);
    }
    arraySchema.elementCount = elementCount;
    return arraySchema;
}
const WgslArrayImpl = {
    [$internal]: true,
    type: 'array',
    toString() {
        return `arrayOf(${String(this.elementType)}, ${this.elementCount})`;
    },
};