typegpu/tgsl/accessIndex.js

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

import { stitch } from "../core/resolve/stitch.js";
import { isDisarray, MatrixColumnsAccess } from "../data/dataTypes.js";
import { derefSnippet } from "../data/ref.js";
import { snip } from "../data/snippet.js";
import { vec2f, vec3f, vec4f } from "../data/vector.js";
import { isPtr, isVec, isWgslArray, isWgslStruct } from "../data/wgslTypes.js";
import { isKnownAtComptime } from "../types.js";
import { accessProp } from "./accessProp.js";
import { ArrayExpression, coerceToSnippet } from "./generationHelpers.js";
const indexableTypeToResult = {
    mat2x2f: vec2f,
    mat3x3f: vec3f,
    mat4x4f: vec4f,
};
export function accessIndex(target, indexArg) {
    const index = typeof indexArg === 'number' ? coerceToSnippet(indexArg) : indexArg;
    // array
    if (isWgslArray(target.dataType) || isDisarray(target.dataType)) {
        const elementType = target.dataType.elementType;
        let origin;
        if (target.origin === 'constant-immutable-def') {
            // Constant refs stay const unless the element/index forces runtime materialization
            origin =
                index.origin === 'constant' || index.origin === 'constant-immutable-def'
                    ? 'constant-immutable-def'
                    : 'runtime-immutable-def';
        }
        else if (target.origin === 'constant') {
            // Ephemeral constants indexed with constants stay constant, otherwise they become runtime-known
            origin =
                index.origin === 'constant' || index.origin === 'constant-immutable-def'
                    ? 'constant'
                    : 'runtime';
        }
        else {
            // Fallthrough
            origin = target.origin;
        }
        if (target.value instanceof ArrayExpression && isKnownAtComptime(index)) {
            return target.value.elements[index.value];
        }
        return snip(isKnownAtComptime(target) && isKnownAtComptime(index)
            ? // oxlint-disable-next-line typescript/no-explicit-any -- it's fine, it's there
                target.value[index.value]
            : stitch `${target}[${index}]`, elementType, 
        /* origin */ origin, target.possibleSideEffects || index.possibleSideEffects);
    }
    // vector
    if (isVec(target.dataType)) {
        return snip(isKnownAtComptime(target) && isKnownAtComptime(index)
            ? // oxlint-disable-next-line typescript/no-explicit-any -- it's fine, it's there
                target.value[index.value]
            : stitch `${target}[${index}]`, target.dataType.primitive, 
        /* origin */ target.origin, target.possibleSideEffects || index.possibleSideEffects);
    }
    if (isPtr(target.dataType)) {
        // Sometimes values that are typed as pointers aren't instances of `d.ref`, so we
        // allow indexing as if it wasn't a pointer.
        return accessIndex(derefSnippet(target), index);
    }
    // matrix.columns
    if (target.value instanceof MatrixColumnsAccess) {
        const propType = indexableTypeToResult[target.value.matrix.dataType.type];
        return snip(stitch `${target.value.matrix}[${index}]`, propType, 
        /* origin */ target.origin, target.possibleSideEffects || index.possibleSideEffects);
    }
    // matrix
    if (target.dataType.type in indexableTypeToResult) {
        throw new Error("The only way of accessing matrix elements in TypeGPU functions is through the 'columns' property.");
    }
    if (isKnownAtComptime(target) && isKnownAtComptime(index)) {
        // No idea what the type is, so we act on the snippet's value and try to guess
        return coerceToSnippet(
        // oxlint-disable-next-line typescript/no-explicit-any -- we're inspecting the value, and it could be any value
        target.value[index.value]);
    }
    if (isWgslStruct(target.dataType) &&
        isKnownAtComptime(index) &&
        typeof index.value === 'string') {
        return accessProp(target, index.value);
    }
    return undefined;
}