typegpu/data/partialIO.js

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

import { BufferWriter, getSystemEndianness } from 'typed-binary';
import { roundUp } from "../mathUtils.js";
import { alignmentOf } from "./alignmentOf.js";
import { getCompiledWriter } from "./compiledIO.js";
import { writeData } from "./dataIO.js";
import { isDisarray, isUnstruct } from "./dataTypes.js";
import { offsetsForProps } from "./offsets.js";
import { sizeOf } from "./sizeOf.js";
import { isWgslArray, isWgslStruct } from "./wgslTypes.js";
/**
 * Converts `{idx, value}[]` sparse arrays into `Record<number, T>` format.
 */
export function convertPartialToPatch(schema, data) {
    if (data === undefined || data === null) {
        return data;
    }
    if (isWgslStruct(schema) || isUnstruct(schema)) {
        const result = {};
        const record = data;
        for (const key of Object.keys(schema.propTypes)) {
            const subSchema = schema.propTypes[key];
            const value = record[key];
            if (value !== undefined && subSchema) {
                result[key] = convertPartialToPatch(subSchema, value);
            }
        }
        return result;
    }
    if (isWgslArray(schema) || isDisarray(schema)) {
        const arrSchema = schema;
        const result = {};
        for (const { idx, value } of data) {
            result[idx] = convertPartialToPatch(arrSchema.elementType, value);
        }
        return result;
    }
    return data;
}
const isLittleEndian = getSystemEndianness() === 'little';
export function getPatchInstructions(schema, data, targetBuffer) {
    const totalSize = sizeOf(schema);
    if (totalSize === 0 || data === undefined || data === null) {
        return [];
    }
    const buf = targetBuffer ?? new ArrayBuffer(totalSize);
    const writer = new BufferWriter(buf);
    const compiledView = new DataView(buf);
    const segments = [];
    function collect(node, value, offset, padding) {
        if (value === undefined || value === null) {
            return;
        }
        if (isWgslStruct(node) || isUnstruct(node)) {
            const propOffsets = offsetsForProps(node);
            for (const [key, propOffset] of Object.entries(propOffsets)) {
                const childValue = value[key];
                const subSchema = node.propTypes[key];
                if (childValue !== undefined && subSchema) {
                    collect(subSchema, childValue, offset + propOffset.offset, propOffset.padding ?? padding);
                }
            }
            return;
        }
        if (isWgslArray(node) || isDisarray(node)) {
            const arrSchema = node;
            const elementSize = roundUp(sizeOf(arrSchema.elementType), alignmentOf(arrSchema.elementType));
            const elementPadding = elementSize - sizeOf(arrSchema.elementType);
            if (ArrayBuffer.isView(value)) {
                const copyLen = Math.min(value.byteLength, arrSchema.elementCount * elementSize);
                new Uint8Array(buf, offset, copyLen).set(new Uint8Array(value.buffer, value.byteOffset, copyLen));
                segments.push({ start: offset, end: offset + copyLen, padding });
                return;
            }
            if (Array.isArray(value)) {
                for (let i = 0; i < Math.min(arrSchema.elementCount, value.length); i++) {
                    collect(arrSchema.elementType, value[i], offset + i * elementSize, elementPadding);
                }
                return;
            }
            const sparse = value;
            for (const key of Object.keys(sparse)) {
                const idx = Number(key);
                if (!Number.isNaN(idx)) {
                    collect(arrSchema.elementType, sparse[key], offset + idx * elementSize, elementPadding);
                }
            }
            return;
        }
        const leafSize = sizeOf(node);
        const compiledWriter = getCompiledWriter(node);
        if (compiledWriter) {
            compiledWriter(compiledView, offset, value, isLittleEndian, offset + leafSize);
        }
        else {
            writer.seekTo(offset);
            writeData(writer, node, value);
        }
        segments.push({ start: offset, end: offset + leafSize, padding });
    }
    collect(schema, data, 0);
    const instructions = [];
    let run = null;
    for (const seg of segments) {
        if (run && seg.start === run.end + (run.padding ?? 0)) {
            run = { start: run.start, end: seg.end, padding: seg.padding };
        }
        else {
            if (run) {
                instructions.push({
                    gpuOffset: run.start,
                    data: new Uint8Array(buf, run.start, run.end - run.start).slice(),
                });
            }
            run = seg;
        }
    }
    if (run) {
        instructions.push({
            gpuOffset: run.start,
            data: new Uint8Array(buf, run.start, run.end - run.start).slice(),
        });
    }
    return instructions;
}
export function patchArrayBuffer(buffer, schema, data) {
    const instructions = getPatchInstructions(schema, data, buffer);
    const mappedView = new Uint8Array(buffer);
    for (const { data, gpuOffset } of instructions) {
        mappedView.set(data, gpuOffset);
    }
}