typegpu 0.10.0 → 0.10.2

package/core/querySet/querySet.js

@@ -0,0 +1,103 @@
+import { $internal } from "../../shared/symbols.js";
+import { setName } from "../../shared/meta.js";
+
+//#region src/core/querySet/querySet.ts
+function INTERNAL_createQuerySet(group, type, count, rawQuerySet) {
+	return new TgpuQuerySetImpl(group, type, count, rawQuerySet);
+}
+function isQuerySet(value) {
+	const maybe = value;
+	return maybe?.resourceType === "query-set" && !!maybe[$internal];
+}
+var TgpuQuerySetImpl = class {
+	resourceType = "query-set";
+	#device;
+	_querySet = null;
+	_ownQuerySet;
+	_destroyed = false;
+	_available = true;
+	_readBuffer = null;
+	_resolveBuffer = null;
+	constructor(root, type, count, rawQuerySet) {
+		this.type = type;
+		this.count = count;
+		this.rawQuerySet = rawQuerySet;
+		this.#device = root.device;
+		this._ownQuerySet = !rawQuerySet;
+		this._querySet = rawQuerySet || null;
+	}
+	get querySet() {
+		if (this._destroyed) throw new Error("This QuerySet has been destroyed.");
+		if (this.rawQuerySet) return this.rawQuerySet;
+		if (this._querySet) return this._querySet;
+		this._querySet = this.#device.createQuerySet({
+			type: this.type,
+			count: this.count
+		});
+		return this._querySet;
+	}
+	get destroyed() {
+		return this._destroyed;
+	}
+	get available() {
+		return this._available;
+	}
+	get [$internal]() {
+		const self = this;
+		return {
+			get readBuffer() {
+				if (!self._readBuffer) self._readBuffer = self.#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.#device.createBuffer({
+					size: self.count * BigUint64Array.BYTES_PER_ELEMENT,
+					usage: GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC
+				});
+				return self._resolveBuffer;
+			}
+		};
+	}
+	$name(label) {
+		setName(this, label);
+		if (this._querySet) this._querySet.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.#device.createCommandEncoder();
+		commandEncoder.resolveQuerySet(this.querySet, 0, this.count, this[$internal].resolveBuffer, 0);
+		this.#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.#device.createCommandEncoder();
+			commandEncoder.copyBufferToBuffer(this[$internal].resolveBuffer, 0, this[$internal].readBuffer, 0, this.count * BigUint64Array.BYTES_PER_ELEMENT);
+			this.#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._querySet && this._ownQuerySet) this._querySet.destroy();
+		this._readBuffer?.destroy();
+		this._resolveBuffer?.destroy();
+		this._readBuffer = this._resolveBuffer = null;
+	}
+};
+
+//#endregion
+export { INTERNAL_createQuerySet, isQuerySet };

package/core/rawCodeSnippet/tgpuRawCodeSnippet.d.ts

@@ -0,0 +1,59 @@
+import { $gpuValueOf } from "../../shared/symbols.js";
+import { Origin } from "../../data/snippet.js";
+import { BaseData } from "../../data/wgslTypes.js";
+import { InferGPU } from "../../shared/repr.js";
+import { AnyData } from "../../data/dataTypes.js";
+
+//#region src/core/rawCodeSnippet/tgpuRawCodeSnippet.d.ts
+/**
+ * Extra declaration that will be included in final WGSL code
+ * when resolving objects that use it.
+ */
+interface TgpuRawCodeSnippet<TDataType extends BaseData> {
+  $: InferGPU<TDataType>;
+  value: InferGPU<TDataType>;
+  readonly [$gpuValueOf]: InferGPU<TDataType>;
+  $uses(dependencyMap: Record<string, unknown>): this;
+}
+type RawCodeSnippetOrigin = Exclude<Origin, 'function' | 'this-function' | 'argument' | 'constant-ref'>;
+/**
+ * An advanced API that creates a typed shader expression which
+ * can be injected into the final shader bundle upon use.
+ *
+ * @param expression The code snippet that will be injected in place of `foo.$`
+ * @param type The type of the expression
+ * @param [origin='runtime'] Where the value originates from.
+ *
+ * **-- Which origin to choose?**
+ *
+ * Usually 'runtime' (the default) is a safe bet, but if you're sure that the expression or
+ * computation is constant (either a reference to a constant, a numeric literal,
+ * or an operation on constants), then pass 'constant' as it might lead to better
+ * optimizations.
+ *
+ * If what the expression is a direct reference to an existing value (e.g. a uniform, a
+ * storage binding, ...), then choose from 'uniform', 'mutable', 'readonly', 'workgroup',
+ * 'private' or 'handle' depending on the address space of the referred value.
+ *
+ * @example
+ * ```ts
+ * // An identifier that we know will be in the
+ * // final shader bundle, but we cannot
+ * // refer to it in any other way.
+ * const existingGlobal = tgpu['~unstable']
+ *   .rawCodeSnippet('EXISTING_GLOBAL', d.f32, 'constant');
+ *
+ * const foo = () => {
+ *   'use gpu';
+ *   return existingGlobal.$ * 2;
+ * };
+ *
+ * const wgsl = tgpu.resolve([foo]);
+ * // fn foo() -> f32 {
+ * //   return EXISTING_GLOBAL * 2;
+ * // }
+ * ```
+ */
+declare function rawCodeSnippet<TDataType extends AnyData>(expression: string, type: TDataType, origin?: RawCodeSnippetOrigin | undefined): TgpuRawCodeSnippet<TDataType>;
+//#endregion
+export { RawCodeSnippetOrigin, TgpuRawCodeSnippet, rawCodeSnippet };

package/core/rawCodeSnippet/tgpuRawCodeSnippet.js

@@ -0,0 +1,96 @@
+import { $gpuValueOf, $internal, $ownSnippet, $resolve } from "../../shared/symbols.js";
+import { snip } from "../../data/snippet.js";
+import { inCodegenMode } from "../../execMode.js";
+import { valueProxyHandler } from "../valueProxyUtils.js";
+import { applyExternals, replaceExternalsInWgsl } from "../resolve/externals.js";
+
+//#region src/core/rawCodeSnippet/tgpuRawCodeSnippet.ts
+/**
+* An advanced API that creates a typed shader expression which
+* can be injected into the final shader bundle upon use.
+*
+* @param expression The code snippet that will be injected in place of `foo.$`
+* @param type The type of the expression
+* @param [origin='runtime'] Where the value originates from.
+*
+* **-- Which origin to choose?**
+*
+* Usually 'runtime' (the default) is a safe bet, but if you're sure that the expression or
+* computation is constant (either a reference to a constant, a numeric literal,
+* or an operation on constants), then pass 'constant' as it might lead to better
+* optimizations.
+*
+* If what the expression is a direct reference to an existing value (e.g. a uniform, a
+* storage binding, ...), then choose from 'uniform', 'mutable', 'readonly', 'workgroup',
+* 'private' or 'handle' depending on the address space of the referred value.
+*
+* @example
+* ```ts
+* // An identifier that we know will be in the
+* // final shader bundle, but we cannot
+* // refer to it in any other way.
+* const existingGlobal = tgpu['~unstable']
+*   .rawCodeSnippet('EXISTING_GLOBAL', d.f32, 'constant');
+*
+* const foo = () => {
+*   'use gpu';
+*   return existingGlobal.$ * 2;
+* };
+*
+* const wgsl = tgpu.resolve([foo]);
+* // fn foo() -> f32 {
+* //   return EXISTING_GLOBAL * 2;
+* // }
+* ```
+*/
+function rawCodeSnippet(expression, type, origin = "runtime") {
+	return new TgpuRawCodeSnippetImpl(expression, type, origin);
+}
+var TgpuRawCodeSnippetImpl = class {
+	[$internal];
+	dataType;
+	origin;
+	#expression;
+	#externalsToApply;
+	constructor(expression, type, origin) {
+		this[$internal] = true;
+		this.dataType = type;
+		this.origin = origin;
+		this.#expression = expression;
+		this.#externalsToApply = [];
+	}
+	$uses(dependencyMap) {
+		this.#externalsToApply.push(dependencyMap);
+		return this;
+	}
+	[$resolve](ctx) {
+		const externalMap = {};
+		for (const externals of this.#externalsToApply) applyExternals(externalMap, externals);
+		return snip(replaceExternalsInWgsl(ctx, externalMap, this.#expression), this.dataType, this.origin);
+	}
+	toString() {
+		return `raw(${String(this.dataType)}): "${this.#expression}"`;
+	}
+	get [$gpuValueOf]() {
+		const dataType = this.dataType;
+		const origin = this.origin;
+		return new Proxy({
+			[$internal]: true,
+			get [$ownSnippet]() {
+				return snip(this, dataType, origin);
+			},
+			[$resolve]: (ctx) => ctx.resolve(this),
+			toString: () => `raw(${String(this.dataType)}): "${this.#expression}".$`
+		}, valueProxyHandler);
+	}
+	get $() {
+		if (!inCodegenMode()) throw new Error("Raw code snippets can only be used on the GPU.");
+		return this[$gpuValueOf];
+	}
+	get value() {
+		return this.$;
+	}
+};
+
+//#endregion
+export { rawCodeSnippet };

package/core/resolve/externals.d.ts

@@ -0,0 +1,10 @@
+import "../../types.js";
+
+//#region src/core/resolve/externals.d.ts
+/**
+ * A key-value mapping where keys represent identifiers within shader code,
+ * and values can be any type that can be resolved to a code string.
+ */
+type ExternalMap = Record<string, unknown>;
+//#endregion
+export { ExternalMap };

package/core/resolve/externals.js

@@ -0,0 +1,58 @@
+import { getName, hasTinyestMetadata, setName } from "../../shared/meta.js";
+import { isWgslStruct } from "../../data/wgslTypes.js";
+import { isLooseData } from "../../data/dataTypes.js";
+import { isWgsl } from "../../types.js";
+
+//#region src/core/resolve/externals.ts
+/**
+* Merges two external maps into one. If a key is present in both maps, the value from the new map is used.
+* If the external value is a namable object, it is given a name if it does not already have one.
+* @param existing - The existing external map.
+* @param newExternals - The new external map.
+*/
+function applyExternals(existing, newExternals) {
+	for (const [key, value] of Object.entries(newExternals)) {
+		existing[key] = value;
+		if (value && (typeof value === "object" || typeof value === "function") && getName(value) === void 0) setName(value, key);
+	}
+}
+function addArgTypesToExternals(implementation, argTypes, applyExternals) {
+	const argTypeNames = [...implementation.matchAll(/:\s*(?<arg>.*?)\s*[,)]/g)].map((found) => found ? found[1] : void 0);
+	applyExternals(Object.fromEntries(argTypes.flatMap((argType, i) => {
+		const argTypeName = argTypeNames ? argTypeNames[i] : void 0;
+		return isWgslStruct(argType) && argTypeName !== void 0 ? [[argTypeName, argType]] : [];
+	})));
+}
+function addReturnTypeToExternals(implementation, returnType, applyExternals) {
+	const matched = implementation.match(/->\s(?<output>[\w\d_]+)\s{/);
+	const outputName = matched ? matched[1]?.trim() : void 0;
+	if (isWgslStruct(returnType) && outputName && !/\s/g.test(outputName)) applyExternals({ [outputName]: returnType });
+}
+function identifierRegex(name) {
+	return new RegExp(`(?<![\\w\\$_.])${name.replaceAll(".", "\\.").replaceAll("$", "\\$")}(?![\\w\\$_])`, "g");
+}
+/**
+* Replaces all occurrences of external names in WGSL code with their resolved values.
+* It adds all necessary definitions to the resolution context.
+* @param ctx - The resolution context.
+* @param externalMap - The external map.
+* @param wgsl - The WGSL code.
+*
+* @returns The WGSL code with all external names replaced with their resolved values.
+*/
+function replaceExternalsInWgsl(ctx, externalMap, wgsl) {
+	return Object.entries(externalMap).reduce((acc, [externalName, external]) => {
+		const externalRegex = identifierRegex(externalName);
+		if (wgsl && externalName !== "Out" && externalName !== "In" && !externalRegex.test(wgsl)) console.warn(`The external '${externalName}' wasn't used in the resolved template.`);
+		if (isWgsl(external) || isLooseData(external) || hasTinyestMetadata(external)) return acc.replaceAll(externalRegex, ctx.resolve(external).value);
+		if (external !== null && typeof external === "object") {
+			const foundProperties = [...wgsl.matchAll(new RegExp(`${externalName.replaceAll(".", "\\.").replaceAll("$", "\\$")}\\.(?<prop>.*?)(?![\\w\\$_])`, "g"))].map((found) => found[1]);
+			return [...new Set(foundProperties)].reduce((innerAcc, prop) => prop && prop in external ? replaceExternalsInWgsl(ctx, { [`${externalName}.${prop}`]: external[prop] }, innerAcc) : innerAcc, acc);
+		}
+		console.warn(`During resolution, the external '${externalName}' has been omitted. Only TGPU resources, 'use gpu' functions, primitives, and plain JS objects can be used as externals.`);
+		return acc;
+	}, wgsl);
+}
+
+//#endregion
+export { addArgTypesToExternals, addReturnTypeToExternals, applyExternals, replaceExternalsInWgsl };

package/core/resolve/namespace.d.ts

@@ -0,0 +1,38 @@
+import { $internal } from "../../shared/symbols.js";
+import { ResolvedSnippet } from "../../data/snippet.js";
+import { ShelllessRepository } from "../../tgsl/shellless.js";
+import { TgpuLazy, TgpuSlot } from "../slot/slotTypes.js";
+import { NameRegistry } from "../../nameRegistry.js";
+
+//#region src/core/resolve/namespace.d.ts
+type SlotToValueMap = Map<TgpuSlot<unknown>, unknown>;
+interface NamespaceInternal {
+  readonly nameRegistry: NameRegistry;
+  readonly shelllessRepo: ShelllessRepository;
+  memoizedResolves: WeakMap<object, {
+    slotToValueMap: SlotToValueMap;
+    result: ResolvedSnippet;
+  }[]>;
+  memoizedLazy: WeakMap<TgpuLazy<unknown>, {
+    slotToValueMap: SlotToValueMap;
+    result: unknown;
+  }[]>;
+  listeners: { [K in keyof NamespaceEventMap]: Set<(event: NamespaceEventMap[K]) => void> };
+}
+type NamespaceEventMap = {
+  name: {
+    target: object;
+    name: string;
+  };
+};
+type DetachListener = () => void;
+interface Namespace {
+  readonly [$internal]: NamespaceInternal;
+  on<TEvent extends keyof NamespaceEventMap>(event: TEvent, listener: (event: NamespaceEventMap[TEvent]) => void): DetachListener;
+}
+interface NamespaceOptions {
+  names?: 'random' | 'strict' | undefined;
+}
+declare function namespace(options?: NamespaceOptions): Namespace;
+//#endregion
+export { Namespace, namespace };

package/core/resolve/namespace.js

@@ -0,0 +1,41 @@
+import { $internal } from "../../shared/symbols.js";
+import { getName } from "../../shared/meta.js";
+import { RandomNameRegistry, StrictNameRegistry } from "../../nameRegistry.js";
+import { ShelllessRepository } from "../../tgsl/shellless.js";
+
+//#region src/core/resolve/namespace.ts
+var NamespaceImpl = class {
+	[$internal];
+	constructor(nameRegistry) {
+		this[$internal] = {
+			nameRegistry,
+			shelllessRepo: new ShelllessRepository(),
+			memoizedResolves: /* @__PURE__ */ new WeakMap(),
+			memoizedLazy: /* @__PURE__ */ new WeakMap(),
+			listeners: { name: /* @__PURE__ */ new Set() }
+		};
+	}
+	on(event, listener) {
+		if (event === "name") {
+			const listeners = this[$internal].listeners.name;
+			listeners.add(listener);
+			return () => listeners.delete(listener);
+		}
+		throw new Error(`Unsupported event: ${event}`);
+	}
+};
+function getUniqueName(namespace, resource) {
+	const name = namespace.nameRegistry.makeUnique(getName(resource), true);
+	for (const listener of namespace.listeners.name) listener({
+		target: resource,
+		name
+	});
+	return name;
+}
+function namespace(options) {
+	const { names = "strict" } = options ?? {};
+	return new NamespaceImpl(names === "strict" ? new StrictNameRegistry() : new RandomNameRegistry());
+}
+
+//#endregion
+export { getUniqueName, namespace };

package/core/resolve/resolveData.js

@@ -0,0 +1,146 @@
+import { $internal } from "../../shared/symbols.js";
+import { isLooseData } from "../../data/dataTypes.js";
+import { assertExhaustive } from "../../shared/utilityTypes.js";
+import { formatToWGSLType } from "../../data/vertexFormatData.js";
+import { getAttributesString } from "../../data/attributes.js";
+import { isWgslComparisonSampler, isWgslSampler } from "../../data/sampler.js";
+import { accessModeMap, isWgslStorageTexture, isWgslTexture } from "../../data/texture.js";
+import { isAttribute } from "../vertexLayout/connectAttributesToShader.js";
+
+//#region src/core/resolve/resolveData.ts
+/**
+* Schemas for which their `type` property directly
+* translates to the resulting WGSL code.
+*/
+const identityTypes = [
+	"bool",
+	"f32",
+	"f16",
+	"i32",
+	"u32",
+	"vec2f",
+	"vec3f",
+	"vec4f",
+	"vec2h",
+	"vec3h",
+	"vec4h",
+	"vec2i",
+	"vec3i",
+	"vec4i",
+	"vec2u",
+	"vec3u",
+	"vec4u",
+	"vec2<bool>",
+	"vec3<bool>",
+	"vec4<bool>",
+	"mat2x2f",
+	"mat3x3f",
+	"mat4x4f",
+	"texture_external"
+];
+function isIdentityType(data) {
+	return identityTypes.includes(data.type);
+}
+/**
+* Resolves a single property of a struct.
+* @param ctx - The resolution context.
+* @param key - The key of the property.
+* @param property - The property itself.
+*
+* @returns The resolved property string.
+*/
+function resolveStructProperty(ctx, [key, property]) {
+	return `  ${getAttributesString(property)}${key}: ${ctx.resolve(property).value},\n`;
+}
+/**
+* Resolves a struct and adds its declaration to the resolution context.
+* @param ctx - The resolution context.
+* @param struct - The struct to resolve.
+*
+* @returns The resolved struct name.
+*/
+function resolveStruct(ctx, struct) {
+	if (struct[$internal].isAbstruct) throw new Error("Cannot resolve abstract struct types to WGSL.");
+	const id = ctx.getUniqueName(struct);
+	ctx.addDeclaration(`\
+struct ${id} {
+${Object.entries(struct.propTypes).map((prop) => resolveStructProperty(ctx, prop)).join("")}\
+}`);
+	return id;
+}
+/**
+* Resolves an unstruct (struct that does not align data by default) to its struct data counterpart.
+* @param ctx - The resolution context.
+* @param unstruct - The unstruct to resolve.
+*
+* @returns The resolved unstruct name.
+*
+* @example
+* ```ts
+* resolveUnstruct(ctx, {
+*   uv: d.float16x2, // -> d.vec2f after resolution
+*   color: d.snorm8x4, -> d.vec4f after resolution
+* });
+* ```
+*/
+function resolveUnstruct(ctx, unstruct) {
+	const id = ctx.getUniqueName(unstruct);
+	ctx.addDeclaration(`\
+struct ${id} {
+${Object.entries(unstruct.propTypes).map((prop) => isAttribute(prop[1]) ? resolveStructProperty(ctx, [prop[0], formatToWGSLType[prop[1].format]]) : resolveStructProperty(ctx, prop)).join("")}
+}`);
+	return id;
+}
+/**
+* Resolves an array.
+* @param ctx - The resolution context.
+* @param array - The array to resolve.
+*
+* @returns The resolved array name along with its element type and count (if not runtime-sized).
+*
+* @example
+* ```ts
+* resolveArray(ctx, d.arrayOf(d.u32, 0)); // 'array<u32>' (not a real pattern, a function is preferred)
+* resolveArray(ctx, d.arrayOf(d.u32, 5)); // 'array<u32, 5>'
+* ```
+*/
+function resolveArray(ctx, array) {
+	const element = ctx.resolve(array.elementType).value;
+	return array.elementCount === 0 ? `array<${element}>` : `array<${element}, ${array.elementCount}>`;
+}
+function resolveDisarray(ctx, disarray) {
+	const element = ctx.resolve(isAttribute(disarray.elementType) ? formatToWGSLType[disarray.elementType.format] : disarray.elementType).value;
+	return disarray.elementCount === 0 ? `array<${element}>` : `array<${element}, ${disarray.elementCount}>`;
+}
+/**
+* Resolves a WGSL data-type schema to a string.
+* @param ctx - The resolution context.
+* @param data - The data-type to resolve.
+*
+* @returns The resolved data-type string.
+*/
+function resolveData(ctx, data) {
+	if (isLooseData(data)) {
+		if (data.type === "unstruct") return resolveUnstruct(ctx, data);
+		if (data.type === "disarray") return resolveDisarray(ctx, data);
+		if (data.type === "loose-decorated") return ctx.resolve(isAttribute(data.inner) ? formatToWGSLType[data.inner.format] : data.inner).value;
+		return ctx.resolve(formatToWGSLType[data.type]).value;
+	}
+	if (isIdentityType(data)) return data.type;
+	if (data.type === "struct") return resolveStruct(ctx, data);
+	if (data.type === "array") return resolveArray(ctx, data);
+	if (data.type === "atomic") return `atomic<${resolveData(ctx, data.inner)}>`;
+	if (data.type === "decorated") return ctx.resolve(data.inner).value;
+	if (data.type === "ptr") {
+		if (data.addressSpace === "storage") return `ptr<storage, ${ctx.resolve(data.inner).value}, ${data.access === "read-write" ? "read_write" : data.access}>`;
+		return `ptr<${data.addressSpace}, ${ctx.resolve(data.inner).value}>`;
+	}
+	if (data.type === "abstractInt" || data.type === "abstractFloat" || data.type === "void" || data.type === "u16") throw new Error(`${data.type} has no representation in WGSL`);
+	if (isWgslStorageTexture(data)) return `${data.type}<${data.format}, ${accessModeMap[data.access]}>`;
+	if (isWgslTexture(data)) return data.type.startsWith("texture_depth") ? data.type : `${data.type}<${data.sampleType.type}>`;
+	if (isWgslComparisonSampler(data) || isWgslSampler(data)) return data.type;
+	assertExhaustive(data, "resolveData");
+}
+
+//#endregion
+export { resolveData };

package/core/resolve/stitch.js

@@ -0,0 +1,25 @@
+import { isSnippet } from "../../data/snippet.js";
+import { getResolutionCtx } from "../../execMode.js";
+
+//#region src/core/resolve/stitch.ts
+/**
+* "The reverse of snipping"
+* Injects resolved snippets into a template string.
+*/
+function stitch(strings, ...snippets) {
+	const ctx = getResolutionCtx();
+	function resolveSnippet(maybeSnippet) {
+		return isSnippet(maybeSnippet) ? ctx.resolve(maybeSnippet.value, maybeSnippet.dataType).value : maybeSnippet;
+	}
+	let result = "";
+	for (let i = 0; i < strings.length; ++i) {
+		result += strings[i];
+		const snippet = snippets[i];
+		if (Array.isArray(snippet)) result += snippet.filter((s) => s !== void 0).map(resolveSnippet).join(", ");
+		else if (snippet) result += resolveSnippet(snippet);
+	}
+	return result;
+}
+
+//#endregion
+export { stitch };