typegpu 0.11.1 → 0.11.3

package/tgsl/shellless.js

@@ -5,6 +5,7 @@ import { WgslTypeError } from "../errors.js";
 import { RefOperator } from "../data/ref.js";
 import { concretize } from "./generationHelpers.js";
 import { createShelllessImpl } from "../core/function/shelllessImpl.js";
+
 //#region src/tgsl/shellless.ts
 function shallowEqualSchemas(a, b) {
 	if (a.type !== b.type) return false;
@@ -42,5 +43,6 @@ var ShelllessRepository = class {
 		return shellless;
 	}
 };
+
 //#endregion
-export { ShelllessRepository };
+export { ShelllessRepository };

package/tgsl/wgslGenerator.d.ts

@@ -1,5 +1,6 @@
 import { Origin, ResolvedSnippet, Snippet } from "../data/snippet.js";
 import { GenerationCtx } from "./generationHelpers.js";
+import { FunctionDefinitionOptions } from "./shaderGenerator_members.js";
 import { ShaderGenerator } from "./shaderGenerator.js";
 import { BaseData, StorableData } from "../data/wgslTypes.js";
 import { UnknownData } from "../data/dataTypes.js";
@@ -22,7 +23,7 @@ declare class WgslGenerator implements ShaderGenerator {
    */
   _typedExpression(expression: tinyest.Expression, expectedType: BaseData | BaseData[]): Snippet;
   _expression(expression: tinyest.Expression): Snippet;
-  functionDefinition(body: tinyest.Block): string;
+  functionDefinition(options: FunctionDefinitionOptions): string;
   /**
    * Generates a WGSL type string for the given data type, and adds necessary
    * definitions to the shader preamble. This shouldn't be called directly, only
@@ -30,6 +31,7 @@ declare class WgslGenerator implements ShaderGenerator {
    */
   typeAnnotation(data: BaseData): string;
   typeInstantiation(schema: BaseData, args: readonly Snippet[]): ResolvedSnippet;
+  _return(statement: tinyest.Return): string;
   _statement(statement: tinyest.Statement): string;
 }
 //#endregion

package/tgsl/wgslGenerator.js

@@ -13,6 +13,7 @@ import { convertStructValues, convertToCommonType, tryConvertSnippet } from "./c
 import { bool, i32, u32 } from "../data/numeric.js";
 import { ArrayExpression, coerceToSnippet, concretize, numericLiteralToSnippet } from "./generationHelpers.js";
 import { vec2u, vec3u, vec4u } from "../data/vector.js";
+import { getAttributesString } from "../data/attributes.js";
 import { AutoStruct } from "../data/autoStruct.js";
 import { add, div, mul, neg, sub } from "../std/operators.js";
 import { accessProp } from "./accessProp.js";
@@ -27,6 +28,7 @@ import { mathToStd } from "./math.js";
 import { isTgpuRange } from "../std/range.js";
 import { getElementSnippet, getElementType, getLoopVarKind, getRangeSnippets } from "./forOfUtils.js";
 import * as tinyest from "tinyest";
+
 //#region src/tgsl/wgslGenerator.ts
 const { NodeTypeCatalog: NODE } = tinyest;
 const parenthesizedOps = [
@@ -215,12 +217,12 @@ ${this.ctx.pre}}`;
 			const lhsExpr = this._expression(lhs);
 			if ((op === "||" || op === "&&") && isKnownAtComptime(lhsExpr)) {
 				if (!(op === "&&" ? lhsExpr.value : !lhsExpr.value)) return snip(op === "||", bool, "constant");
-				const rhsExpr = this._expression(rhs);
-				if (rhsExpr.dataType === UnknownData) throw new WgslTypeError(`Right-hand side of '${op}' is of unknown type`);
-				if (isKnownAtComptime(rhsExpr)) return snip(!!rhsExpr.value, bool, "constant");
-				const convRhs = tryConvertSnippet(this.ctx, rhsExpr, bool, false);
-				const rhsStr = this.ctx.resolve(convRhs.value, convRhs.dataType).value;
-				return snip(rhsStr, bool, "runtime");
+				const rhsExpr$1 = this._expression(rhs);
+				if (rhsExpr$1.dataType === UnknownData) throw new WgslTypeError(`Right-hand side of '${op}' is of unknown type`);
+				if (isKnownAtComptime(rhsExpr$1)) return snip(!!rhsExpr$1.value, bool, "constant");
+				const convRhs$1 = tryConvertSnippet(this.ctx, rhsExpr$1, bool, false);
+				const rhsStr$1 = this.ctx.resolve(convRhs$1.value, convRhs$1.dataType).value;
+				return snip(rhsStr$1, bool, "runtime");
 			}
 			const rhsExpr = this._expression(rhs);
 			if (rhsExpr.value instanceof RefOperator) throw new WgslTypeError(stitch`Cannot assign a ref to an existing variable '${lhsExpr}', define a new variable instead.`);
@@ -291,7 +293,7 @@ ${this.ctx.pre}}`;
 			if (!accessed) {
 				const targetStr = this.ctx.resolve(target.value, target.dataType).value;
 				const propertyStr = this.ctx.resolve(property.value, property.dataType).value;
-				throw new Error(`Unable to index value ${targetStr} of unknown type with index ${propertyStr}. If the value is an array, to address this, consider one of the following approaches: (1) declare the array using 'tgpu.const', (2) store the array in a buffer, or (3) define the array within the GPU function scope.`);
+				throw new Error(`Index access '${targetStr}[${propertyStr}]' is invalid. If the value is an array, to address this, consider one of the following approaches: (1) declare the array using 'tgpu.const', (2) store the array in a buffer, or (3) define the array within the GPU function scope.`);
 			}
 			return accessed;
 		}
@@ -421,8 +423,13 @@ ${this.ctx.pre}}`;
 		if (expression[0] === NODE.preUpdate) throw new Error("Cannot use pre-updates in TypeGPU functions.");
 		assertExhaustive(expression);
 	}
-	functionDefinition(body) {
-		return this._block(body);
+	functionDefinition(options) {
+		const body = this._block(options.body);
+		const returnType = options.determineReturnType();
+		const argList = options.args.filter((arg) => arg.used || options.functionType === "normal").map((arg) => {
+			return `${getAttributesString(arg.decoratedType)}${arg.name}: ${this.ctx.resolve(arg.decoratedType).value}`;
+		}).join(", ");
+		return `${returnType.type !== "void" ? `(${argList}) -> ${getAttributesString(returnType)}${this.ctx.resolve(returnType).value} ` : `(${argList}) `}${body}`;
 	}
 	/**
 	* Generates a WGSL type string for the given data type, and adds necessary
@@ -436,35 +443,36 @@ ${this.ctx.pre}}`;
 		if (args.length === 1 && args[0]?.dataType === schema) return snip(stitch`${args[0]}`, schema, fallthroughCopyOrigin(args[0].origin));
 		return snip(stitch`${this.ctx.resolve(schema).value}(${args})`, schema, "runtime");
 	}
-	_statement(statement) {
-		if (typeof statement === "string") {
-			const id = this._identifier(statement);
-			const resolved = id.value && this.ctx.resolve(id.value).value;
-			return resolved ? `${this.ctx.pre}${resolved};` : "";
-		}
-		if (typeof statement === "boolean") return `${this.ctx.pre}${statement ? "true" : "false"};`;
-		if (statement[0] === NODE.return) {
-			const returnNode = statement[1];
-			if (returnNode !== void 0) {
-				const expectedReturnType = this.ctx.topFunctionReturnType;
-				let returnSnippet = expectedReturnType ? this._typedExpression(returnNode, expectedReturnType) : this._expression(returnNode);
-				if (returnSnippet.value instanceof RefOperator) throw new WgslTypeError(stitch`Cannot return references, returning '${returnSnippet.value.snippet}'`);
-				if (returnSnippet.origin === "argument" && !isNaturallyEphemeral(returnSnippet.dataType) && this.ctx.topFunctionScope?.functionType === "normal") throw new WgslTypeError(stitch`Cannot return references to arguments, returning '${returnSnippet}'. Copy the argument before returning it.`);
-				if (!expectedReturnType && !isEphemeralSnippet(returnSnippet) && returnSnippet.origin !== "this-function") {
-					const str = this.ctx.resolve(returnSnippet.value, returnSnippet.dataType).value;
-					const typeStr = this.ctx.resolve(unptr(returnSnippet.dataType)).value;
-					throw new WgslTypeError(`'return ${str};' is invalid, cannot return references.
+	_return(statement) {
+		const returnNode = statement[1];
+		if (returnNode !== void 0) {
+			const expectedReturnType = this.ctx.topFunctionReturnType;
+			let returnSnippet = expectedReturnType ? this._typedExpression(returnNode, expectedReturnType) : this._expression(returnNode);
+			if (returnSnippet.value instanceof RefOperator) throw new WgslTypeError(stitch`Cannot return references, returning '${returnSnippet.value.snippet}'`);
+			if (returnSnippet.origin === "argument" && !isNaturallyEphemeral(returnSnippet.dataType) && this.ctx.topFunctionScope?.functionType === "normal") throw new WgslTypeError(stitch`Cannot return references to arguments, returning '${returnSnippet}'. Copy the argument before returning it.`);
+			if (!expectedReturnType && !isEphemeralSnippet(returnSnippet) && returnSnippet.origin !== "this-function") {
+				const str = this.ctx.resolve(returnSnippet.value, returnSnippet.dataType).value;
+				const typeStr = this.ctx.resolve(unptr(returnSnippet.dataType)).value;
+				throw new WgslTypeError(`'return ${str};' is invalid, cannot return references.
 -----
 Try 'return ${typeStr}(${str});' instead.
 -----`);
-				}
-				returnSnippet = tryConvertSnippet(this.ctx, returnSnippet, unptr(returnSnippet.dataType), false);
-				invariant(returnSnippet.dataType !== UnknownData, "Return type should be known");
-				this.ctx.reportReturnType(returnSnippet.dataType);
-				return stitch`${this.ctx.pre}return ${returnSnippet};`;
 			}
-			return `${this.ctx.pre}return;`;
+			returnSnippet = tryConvertSnippet(this.ctx, returnSnippet, unptr(returnSnippet.dataType), false);
+			invariant(returnSnippet.dataType !== UnknownData, "Return type should be known");
+			this.ctx.reportReturnType(returnSnippet.dataType);
+			return stitch`${this.ctx.pre}return ${returnSnippet};`;
 		}
+		return `${this.ctx.pre}return;`;
+	}
+	_statement(statement) {
+		if (typeof statement === "string") {
+			const id = this._identifier(statement);
+			const resolved$1 = id.value && this.ctx.resolve(id.value).value;
+			return resolved$1 ? `${this.ctx.pre}${resolved$1};` : "";
+		}
+		if (typeof statement === "boolean") return `${this.ctx.pre}${statement ? "true" : "false"};`;
+		if (statement[0] === NODE.return) return this._return(statement);
 		if (statement[0] === NODE.if) {
 			const [_, condNode, consNode, altNode] = statement;
 			const condition = this._typedExpression(condNode, bool);
@@ -500,12 +508,12 @@ ${this.ctx.pre}else ${alternate}`;
 			}
 			if (!ephemeral) {
 				if (stmtType === NODE.let) {
-					const rhsStr = this.ctx.resolve(eq.value).value;
+					const rhsStr$1 = this.ctx.resolve(eq.value).value;
 					const rhsTypeStr = this.ctx.resolve(unptr(eq.dataType)).value;
-					throw new WgslTypeError(`'let ${rawId} = ${rhsStr}' is invalid, because references cannot be assigned to 'let' variable declarations.
+					throw new WgslTypeError(`'let ${rawId} = ${rhsStr$1}' is invalid, because references cannot be assigned to 'let' variable declarations.
 -----
-- Try 'let ${rawId} = ${rhsTypeStr}(${rhsStr})' if you need to reassign '${rawId}' later
-- Try 'const ${rawId} = ${rhsStr}' if you won't reassign '${rawId}' later.
+- Try 'let ${rawId} = ${rhsTypeStr}(${rhsStr$1})' if you need to reassign '${rawId}' later
+- Try 'const ${rawId} = ${rhsStr$1}' if you won't reassign '${rawId}' later.
 -----`);
 				}
 				if (eq.origin === "constant-tgpu-const-ref") varType = "const";
@@ -524,12 +532,12 @@ ${this.ctx.pre}else ${alternate}`;
 				else if (naturallyEphemeral) varType = eq.origin === "constant" ? "const" : "let";
 			} else if (eq.origin === "argument") {
 				if (!naturallyEphemeral) {
-					const rhsStr = this.ctx.resolve(eq.value).value;
+					const rhsStr$1 = this.ctx.resolve(eq.value).value;
 					const rhsTypeStr = this.ctx.resolve(unptr(eq.dataType)).value;
-					throw new WgslTypeError(`'let ${rawId} = ${rhsStr}' is invalid, because references to arguments cannot be assigned to 'let' variable declarations.
+					throw new WgslTypeError(`'let ${rawId} = ${rhsStr$1}' is invalid, because references to arguments cannot be assigned to 'let' variable declarations.
   -----
-  - Try 'let ${rawId} = ${rhsTypeStr}(${rhsStr})' if you need to reassign '${rawId}' later
-  - Try 'const ${rawId} = ${rhsStr}' if you won't reassign '${rawId}' later.
+  - Try 'let ${rawId} = ${rhsTypeStr}(${rhsStr$1})' if you need to reassign '${rawId}' later
+  - Try 'const ${rawId} = ${rhsStr$1}' if you won't reassign '${rawId}' later.
   -----`);
 				}
 			}
@@ -591,7 +599,7 @@ ${this.ctx.pre}else ${alternate}`;
 					const length = range.end.value;
 					if (length === 0) return "";
 					const { value } = iterableSnippet;
-					const elements = isTgpuRange(value) ? value.map((i) => coerceToSnippet(i)) : value instanceof ArrayExpression ? value.elements : Array.from({ length }, (_, i) => getElementSnippet(iterableSnippet, snip(i, u32, "constant")));
+					const elements = isTgpuRange(value) ? value.map((i) => coerceToSnippet(i)) : value instanceof ArrayExpression ? value.elements : Array.from({ length }, (_$1, i) => getElementSnippet(iterableSnippet, snip(i, u32, "constant")));
 					if (isEphemeralSnippet(elements[0]) && !isNaturallyEphemeral(elements[0]?.dataType)) throw new WgslTypeError("Cannot unroll loop. The elements of iterable are emphemeral but not naturally ephemeral.");
 					return elements.map((e, i) => `${this.ctx.pre}// unrolled iteration #${i}\n${this.ctx.pre}${this._block(blockified, { [originalLoopVarName]: e })}`).join("\n");
 				}
@@ -599,7 +607,7 @@ ${this.ctx.pre}else ${alternate}`;
 				const index = this.ctx.makeNameValid("i");
 				const forHeaderStr = stitch`${this.ctx.pre}for (var ${index} = ${range.start}; ${index} ${range.comparison} ${range.end}; ${index} += ${range.step})`;
 				let bodyStr = "";
-				if (isTgpuRange(iterableSnippet.value)) bodyStr = this._block(blockified, { [originalLoopVarName]: snip(index, u32, "runtime") });
+				if (isTgpuRange(iterableSnippet.value)) bodyStr = this._block(blockified, { [originalLoopVarName]: snip(index, range.start.dataType, "runtime") });
 				else {
 					this.ctx.indent();
 					ctxIndent = true;
@@ -644,5 +652,7 @@ function blockifySingleStatement(statement) {
 	return typeof statement !== "object" || statement[0] !== NODE.block ? [NODE.block, [statement]] : statement;
 }
 const wgslGenerator = new WgslGenerator();
+var wgslGenerator_default = wgslGenerator;
+
 //#endregion
-export { WgslGenerator, wgslGenerator as default };
+export { WgslGenerator, wgslGenerator_default as default };

package/types.d.ts

@@ -49,11 +49,17 @@ type ItemLayer = {
   type: 'item';
   usedSlots: Set<TgpuSlot<unknown>>;
 };
+type FunctionArgumentAccess = () => Snippet | undefined;
+interface FunctionArgument {
+  name: string;
+  access: FunctionArgumentAccess;
+  decoratedType: BaseData;
+  used: boolean;
+}
 type FunctionScopeLayer = {
   type: 'functionScope';
   functionType: 'normal' | 'compute' | 'vertex' | 'fragment';
-  args: Snippet[];
-  argAliases: Record<string, Snippet>;
+  argAccess: Record<string, FunctionArgumentAccess>;
   externalMap: Record<string, unknown>;
   /**
    * The return type of the function. If undefined, the type should be inferred
@@ -81,12 +87,11 @@ interface ItemStateStack {
   readonly topFunctionScope: FunctionScopeLayer | undefined;
   pushItem(): void;
   pushSlotBindings(pairs: SlotValuePair[]): void;
-  pushFunctionScope(functionType: 'normal' | TgpuShaderStage, args: Snippet[], argAliases: Record<string, Snippet>,
+  pushFunctionScope(functionType: 'normal' | TgpuShaderStage, argAccess: Record<string, FunctionArgumentAccess>,
   /**
    * The return type of the function. If undefined, the type should be inferred
    * from the implementation (relevant for shellless functions).
    */
-
   returnType: BaseData | undefined, externalMap: Record<string, unknown>): FunctionScopeLayer;
   pushBlockScope(): void;
   setBlockExternals(externals: Record<string, Snippet>): void;
@@ -218,8 +223,7 @@ interface ResolutionCtx {
    */
   resolveSnippet(snippet: Snippet): ResolvedSnippet;
   fnToWgsl(options: FnToWgslOptions): {
-    head: Wgsl;
-    body: Wgsl;
+    code: string;
     returnType: BaseData;
   };
   withVaryingLocations<T>(locations: Record<string, number>, callback: () => T): T;
@@ -262,4 +266,4 @@ type AnyFn = (...args: never[]) => unknown;
 type DualFn<T extends AnyFn> = T & GPUCallable<Parameters<T>>;
 type BindableBufferUsage = 'uniform' | 'readonly' | 'mutable';
 //#endregion
-export { BindableBufferUsage, DualFn, FunctionScopeLayer, GPUCallable, ResolutionCtx, ResolvableObject, SelfResolvable, TgpuShaderStage, Wgsl, WithCast };
+export { BindableBufferUsage, DualFn, FunctionArgument, FunctionScopeLayer, GPUCallable, ResolutionCtx, ResolvableObject, SelfResolvable, TgpuShaderStage, Wgsl, WithCast };

package/types.js

@@ -2,6 +2,7 @@ import { $cast, $gpuCallable, $ownSnippet, $resolve } from "./shared/symbols.js"
 import { isWgslData } from "./data/wgslTypes.js";
 import { UnknownData } from "./data/dataTypes.js";
 import { isLazy, isProviding, isSlot } from "./core/slot/slotTypes.js";
+
 //#region src/types.ts
 var NormalState = class {
 	type = "normal";
@@ -39,5 +40,6 @@ function isWgsl(value) {
 function isGPUBuffer(value) {
 	return !!value && typeof value === "object" && "getMappedRange" in value && "mapAsync" in value;
 }
+
 //#endregion
-export { CodegenState, NormalState, SimulationState, getOwnSnippet, hasCast, isGPUBuffer, isGPUCallable, isKnownAtComptime, isSelfResolvable, isWgsl };
+export { CodegenState, NormalState, SimulationState, getOwnSnippet, hasCast, isGPUBuffer, isGPUCallable, isKnownAtComptime, isSelfResolvable, isWgsl };

package/wgslExtensions.js

@@ -13,5 +13,6 @@ const wgslExtensionToFeatureName = {
 	subgroups: "subgroups",
 	primitive_index: "primitive-index"
 };
+
 //#endregion
-export { wgslExtensionToFeatureName, wgslExtensions };
+export { wgslExtensionToFeatureName, wgslExtensions };