typegpu/core/querySet/querySet.js

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

import { setName } from "../../shared/meta.js";
import { $internal, $soul } from "../../shared/symbols.js";
export function INTERNAL_createQuerySet(group, type, count, rawQuerySet) {
    return new TgpuQuerySetImpl(group, type, count, rawQuerySet);
}
export function isQuerySet(value) {
    const maybe = value;
    return maybe?.resourceType === 'query-set' && !!maybe[$internal];
}
class TgpuQuerySetImpl {
    resourceType = 'query-set';
    #ownQuerySet;
    #destroyed = false;
    #available = true;
    #readBuffer = undefined;
    #resolveBuffer = undefined;
    [$soul];
    [$internal];
    constructor(root, type, count, rawQuerySet) {
        this.#ownQuerySet = !rawQuerySet;
        this[$soul] = {
            type: 'query-set',
            device: root.device,
            queryType: type,
            count,
            raw: rawQuerySet,
            label: undefined,
        };
        // oxlint-disable-next-line typescript/no-this-alias
        const self = this;
        this[$internal] = {
            get readBuffer() {
                if (!self.#readBuffer) {
                    self.#readBuffer = self[$soul].device.createBuffer({
                        size: self.count * BigUint64Array.BYTES_PER_ELEMENT,
                        usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
                    });
                }
                return self.#readBuffer;
            },
            get resolveBuffer() {
                if (!self.#resolveBuffer) {
                    self.#resolveBuffer = self[$soul].device.createBuffer({
                        size: self.count * BigUint64Array.BYTES_PER_ELEMENT,
                        usage: GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC,
                    });
                }
                return self.#resolveBuffer;
            },
            materialize: () => {
                if (this.#destroyed) {
                    throw new Error('This QuerySet has been destroyed.');
                }
                const soul = this[$soul];
                if (!soul.raw) {
                    soul.raw = soul.device.createQuerySet({
                        type: soul.queryType,
                        count: soul.count,
                    });
                }
                return soul.raw;
            },
        };
    }
    get type() {
        return this[$soul].queryType;
    }
    get count() {
        return this[$soul].count;
    }
    get querySet() {
        return this[$internal].materialize();
    }
    get destroyed() {
        return this.#destroyed;
    }
    get available() {
        return this.#available;
    }
    $name(label) {
        setName(this, label);
        const raw = this[$soul].raw;
        if (raw) {
            raw.label = label;
        }
        return this;
    }
    resolve() {
        if (this.#destroyed) {
            throw new Error('This QuerySet has been destroyed.');
        }
        if (!this.#available) {
            throw new Error('This QuerySet is busy resolving or reading.');
        }
        const commandEncoder = this[$soul].device.createCommandEncoder();
        commandEncoder.resolveQuerySet(this.querySet, 0, this.count, this[$internal].resolveBuffer, 0);
        this[$soul].device.queue.submit([commandEncoder.finish()]);
    }
    async read() {
        if (!this.#resolveBuffer) {
            throw new Error('QuerySet must be resolved before reading.');
        }
        this.#available = false;
        try {
            const commandEncoder = this[$soul].device.createCommandEncoder();
            commandEncoder.copyBufferToBuffer(this[$internal].resolveBuffer, 0, this[$internal].readBuffer, 0, this.count * BigUint64Array.BYTES_PER_ELEMENT);
            this[$soul].device.queue.submit([commandEncoder.finish()]);
            const readBuffer = this[$internal].readBuffer;
            await readBuffer.mapAsync(GPUMapMode.READ);
            const data = new BigUint64Array(readBuffer.getMappedRange().slice());
            readBuffer.unmap();
            return Array.from(data);
        }
        finally {
            this.#available = true;
        }
    }
    destroy() {
        if (this.#destroyed) {
            return;
        }
        this.#destroyed = true;
        if (this[$soul].raw && this.#ownQuerySet) {
            this[$soul].raw.destroy();
        }
        this.#readBuffer?.destroy();
        this.#resolveBuffer?.destroy();
        this.#readBuffer = this.#resolveBuffer = undefined;
    }
}