typegpu/data/ref.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 { WgslTypeError } from "../errors.js";
import { setName } from "../shared/meta.js";
import { $gpuCallable, $internal, $ownSnippet, $resolve } from "../shared/symbols.js";
import { UnknownData } from "./dataTypes.js";
import { createPtrFromOrigin, explicitFrom } from "./ptr.js";
import { isAlias, snip, withDataType } from "./snippet.js";
import { isNaturallyEphemeral, isPtr } from "./wgslTypes.js";
export const _ref = (() => {
    const impl = ((value) => INTERNAL_createRef(value));
    setName(impl, 'ref');
    impl.toString = () => 'ref';
    impl[$internal] = true;
    impl[$gpuCallable] = {
        call(ctx, [value]) {
            if (value.origin === 'argument') {
                throw new WgslTypeError(stitch `d.ref(${value}) is illegal, cannot take a reference of an argument. Copy the value first, and take a reference of the copy.`);
            }
            if (value.origin === 'constant-immutable-def' || value.origin === 'runtime-immutable-def') {
                const typeStr = ctx.resolve(value.dataType).value;
                throw new WgslTypeError(stitch `d.ref(${value}) is illegal, cannot take a reference to a constant.
-----
- Try 'd.ref(${typeStr}(${value}));' instead to create a new referencable value.
-----`);
            }
            if (isAlias(value) && isNaturallyEphemeral(value.dataType)) {
                const typeStr = ctx.resolve(value.dataType).value;
                throw new WgslTypeError(stitch `d.ref(${value}) is illegal, cannot take a reference to a scalar value.
-----
- Try 'd.ref(${typeStr}(${value}));' instead to create a new referencable scalar.
-----`);
            }
            if (isPtr(value.dataType)) {
                // This can happen if we take a reference of an *implicit* pointer, one
                // made by assigning a reference to a `const`.
                return withDataType(explicitFrom(value.dataType), value);
            }
            /**
             * Pointer type only exists if the ref was created from a reference (buttery-butter).
             *
             * @example
             * ```ts
             * const life = ref(42); // created from a value
             * const boid = ref(layout.$.boids[0]); // created from a reference
             * ```
             */
            const ptrType = createPtrFromOrigin(value.origin, value.dataType);
            return snip(new RefOperator(value, ptrType), ptrType ?? UnknownData, 
            /* origin */ 'runtime', value.possibleSideEffects);
        },
    };
    return impl;
})();
export function isRef(value) {
    return value?.[$internal]?.type === 'ref';
}
// --------------
// Implementation
// --------------
export function INTERNAL_createRef(value) {
    const target = {
        [$internal]: { type: 'ref' },
        get $() {
            return value;
        },
        set $(newValue) {
            if (newValue && typeof newValue === 'object') {
                // Setting an object means updating the properties of the original object.
                // e.g.: foo.$ = Boid();
                for (const key of Object.keys(newValue)) {
                    value[key] = newValue[key];
                }
            }
            else {
                value = newValue;
            }
        },
    };
    if (value === undefined || value === null) {
        throw new Error('Cannot create a ref from undefined or null');
    }
    if (typeof value === 'object') {
        return new Proxy(target, {
            get(target, prop) {
                if (prop in target) {
                    return target[prop];
                }
                return value[prop];
            },
            set(_target, prop, propValue) {
                if (prop === $internal) {
                    return false;
                }
                if (prop === '$') {
                    return Reflect.set(target, prop, propValue);
                }
                return Reflect.set(value, prop, propValue);
            },
        });
    }
    return target;
}
/**
 * The result of calling `d.ref(...)`. The code responsible for
 * generating shader code can check if the value of a snippet is
 * an instance of `RefOperator`, and act accordingly.
 */
export class RefOperator {
    [$internal];
    snippet;
    #ptrType;
    constructor(snippet, ptrType) {
        this[$internal] = true;
        this.snippet = snippet;
        this.#ptrType = ptrType;
    }
    get [$ownSnippet]() {
        if (!this.#ptrType) {
            throw new Error(stitch `Cannot take a reference of ${this.snippet}`);
        }
        return snip(this, this.#ptrType, this.snippet.origin, this.snippet.possibleSideEffects);
    }
    [$resolve]() {
        if (!this.#ptrType) {
            throw new Error(stitch `Cannot take a reference of ${this.snippet}`);
        }
        return snip(stitch `(&${this.snippet})`, this.#ptrType, this.snippet.origin, this.snippet.possibleSideEffects);
    }
}
export function derefSnippet(snippet) {
    if (!isPtr(snippet.dataType)) {
        return snippet;
    }
    const innerType = snippet.dataType.inner;
    if (snippet.value instanceof RefOperator) {
        return snip(stitch `${snippet.value.snippet}`, innerType, snippet.origin, snippet.possibleSideEffects);
    }
    return snip(stitch `(*${snippet})`, innerType, snippet.origin, snippet.possibleSideEffects);
}