typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
228 lines (227 loc) • 8.61 kB
JavaScript
import { roundUp } from "../mathUtils.js";
import { alignmentOf } from "./alignmentOf.js";
import { offsetsForProps } from "./offsets.js";
import { sizeOf } from "./sizeOf.js";
import { isContiguous } from "./isContiguous.js";
import { getLongestContiguousPrefix } from "./getLongestContiguousPrefix.js";
import { isVec, isWgslArray, isWgslStruct } from "./wgslTypes.js";
import { undecorate } from "./dataTypes.js";
const OFFSET_MARKER = Symbol('indirectOffset');
const CONTIGUOUS_MARKER = Symbol('indirectContiguous');
function isOffsetProxy(value) {
return (typeof value === 'object' &&
value !== null &&
OFFSET_MARKER in value &&
CONTIGUOUS_MARKER in value);
}
function scalarNode(offset, contiguous) {
return { [OFFSET_MARKER]: offset, [CONTIGUOUS_MARKER]: contiguous };
}
function getMarker(target, prop) {
if (prop === OFFSET_MARKER) {
return target[OFFSET_MARKER];
}
if (prop === CONTIGUOUS_MARKER) {
return target[CONTIGUOUS_MARKER];
}
return undefined;
}
function makeProxy(schema, baseOffset, contiguous = sizeOf(schema)) {
const unwrapped = undecorate(schema);
const vecComponentCount = isVec(unwrapped) ? unwrapped.componentCount : undefined;
if (vecComponentCount !== undefined) {
const componentSize = sizeOf(unwrapped.primitive);
return makeVecProxy(scalarNode(baseOffset, contiguous), componentSize, vecComponentCount);
}
if (isWgslStruct(unwrapped)) {
return makeStructProxy(unwrapped, scalarNode(baseOffset, contiguous));
}
if (isWgslArray(unwrapped)) {
return makeArrayProxy(unwrapped, scalarNode(baseOffset, contiguous));
}
return scalarNode(baseOffset, contiguous);
}
export function createOffsetProxy(schema, baseOffset = 0) {
return makeProxy(schema, baseOffset, sizeOf(schema));
}
function makeVecProxy(target, componentSize, componentCount) {
const baseOffset = target[OFFSET_MARKER];
return new Proxy(target, {
get(t, prop) {
const marker = getMarker(t, prop);
if (marker !== undefined) {
return marker;
}
const idx = prop === 'x' || prop === '0'
? 0
: prop === 'y' || prop === '1'
? 1
: prop === 'z' || prop === '2'
? 2
: prop === 'w' || prop === '3'
? 3
: -1;
if (idx < 0 || idx >= componentCount) {
return undefined;
}
const byteOffset = idx * componentSize;
const contiguous = Math.max(0, t[CONTIGUOUS_MARKER] - byteOffset);
return scalarNode(baseOffset + byteOffset, contiguous);
},
});
}
function makeArrayProxy(array, target) {
const elementType = array.elementType;
const elementSize = sizeOf(elementType);
const stride = roundUp(elementSize, alignmentOf(elementType));
const hasPadding = stride > elementSize;
return new Proxy(target, {
get(t, prop) {
const marker = getMarker(t, prop);
if (marker !== undefined) {
return marker;
}
if (prop === 'length') {
return array.elementCount;
}
if (typeof prop !== 'string') {
return undefined;
}
const index = Number(prop);
if (!Number.isInteger(index) || index < 0 || index >= array.elementCount) {
return undefined;
}
const elementOffset = index * stride;
const remainingFromHere = !isContiguous(elementType)
? elementSize + getLongestContiguousPrefix(elementType) // it is too much, but we correct it later
: Math.max(0, t[CONTIGUOUS_MARKER] - elementOffset);
const childContiguous = hasPadding
? Math.min(remainingFromHere, elementSize)
: remainingFromHere;
return makeProxy(elementType, t[OFFSET_MARKER] + elementOffset, childContiguous);
},
});
}
function makeStructProxy(struct, target) {
const offsets = offsetsForProps(struct);
const propTypes = struct.propTypes;
const propNames = Object.keys(propTypes);
const meta = new Map();
let runStart = 0;
for (let i = 0; i < propNames.length; i++) {
const name = propNames[i];
if (!name) {
continue;
}
const type = propTypes[name];
if (!type) {
continue;
}
const info = offsets[name];
const padding = info.padding ?? 0;
const typeContiguous = isContiguous(type);
const isRunEnd = i === propNames.length - 1 || padding > 0 || !typeContiguous;
if (!isRunEnd) {
continue;
}
const runEnd = info.offset + (typeContiguous ? info.size : getLongestContiguousPrefix(type));
for (let j = runStart; j <= i; j++) {
const runName = propNames[j];
if (!runName) {
continue;
}
const runInfo = offsets[runName];
meta.set(runName, { offset: runInfo.offset, runEnd });
}
runStart = i + 1;
}
return new Proxy(target, {
get(t, prop) {
const marker = getMarker(t, prop);
if (marker !== undefined) {
return marker;
}
if (typeof prop !== 'string') {
return undefined;
}
const m = meta.get(prop);
if (!m) {
return undefined;
}
const remainingFromHere = Math.max(0, t[CONTIGUOUS_MARKER] - m.offset);
const localLimit = Math.max(0, m.runEnd - m.offset);
const propSchema = propTypes[prop];
if (!propSchema) {
return undefined;
}
return makeProxy(propSchema, t[OFFSET_MARKER] + m.offset, sizeOf(struct) === m.runEnd ? remainingFromHere : localLimit);
},
});
}
function getRootContiguous(schema) {
const unwrapped = undecorate(schema);
if (isWgslStruct(unwrapped)) {
const offsets = offsetsForProps(unwrapped);
const propTypes = unwrapped.propTypes;
const propNames = Object.keys(propTypes);
for (let i = 0; i < propNames.length; i++) {
const name = propNames[i];
if (!name) {
continue;
}
const info = offsets[name];
const padding = info.padding ?? 0;
const runEnd = info.offset + info.size;
const isRunEnd = i === propNames.length - 1 || padding > 0;
if (isRunEnd) {
return runEnd;
}
}
return 0;
}
if (isWgslArray(unwrapped)) {
const elementType = unwrapped.elementType;
const elementSize = sizeOf(elementType);
const stride = roundUp(elementSize, alignmentOf(elementType));
const totalSize = sizeOf(schema);
if (!Number.isFinite(totalSize)) {
return elementSize;
}
return stride > elementSize ? elementSize : totalSize;
}
return sizeOf(schema);
}
/**
* A function that retrieves offset and information for a specific primitive within a data schema.
* Example usage:
* ```ts
* const Boid = d.struct({
* position: d.vec3f,
* velocity: d.vec3f,
* });
* const memLayout = d.memoryLayoutOf(Boid, (b) => b.velocity.y);
* console.log(memLayout.offset); // Byte offset of velocity.y within Boid (here 20 bytes)
* console.log(memLayout.contiguous); // Contiguous bytes available from that offset (here 8 bytes)
* ```
*
* @param schema - The data schema to analyze.
* @param accessor - Optional function that accesses a specific element within the schema. If omitted, uses the root offset (0).
* @returns An object containing the offset and contiguous byte information.
*/
export function memoryLayoutOf(schema, accessor) {
if (!accessor) {
return {
offset: 0,
contiguous: getRootContiguous(schema),
};
}
const proxy = createOffsetProxy(schema);
const result = accessor(proxy);
if (isOffsetProxy(result)) {
return {
offset: result[OFFSET_MARKER],
contiguous: result[CONTIGUOUS_MARKER],
};
}
throw new Error('memoryLayoutOf: accessor did not return a schema element. Make sure the accessor navigates to a field or element of the schema (e.g. `(s) => s.position.x`).');
}