typegpu 0.10.2 → 0.11.0-rc.1

package/tgsl/generationHelpers.js

@@ -7,7 +7,6 @@ import { getOwnSnippet } from "../types.js";
 import { stitch } from "../core/resolve/stitch.js";
 import { isRef } from "../data/ref.js";
 import { abstractFloat, abstractInt, bool, f32, i32 } from "../data/numeric.js";
-
 //#region src/tgsl/generationHelpers.ts
 function numericLiteralToSnippet(value) {
 	if (value >= 2 ** 63 || value < -(2 ** 63)) return snip(value, abstractFloat, "constant");
@@ -46,6 +45,8 @@ function coerceToSnippet(value) {
 */
 var ArrayExpression = class {
 	[$internal] = true;
+	type;
+	elements;
 	constructor(type, elements) {
 		this.type = type;
 		this.elements = elements;
@@ -62,6 +63,5 @@ var ArrayExpression = class {
 		return snip(stitch`${`array<${ctx.resolve(this.type.elementType).value}, ${this.elements.length}>`}(${this.elements})`, this.type, "runtime");
 	}
 };
-
 //#endregion
-export { ArrayExpression, coerceToSnippet, concretize, concretizeSnippets, numericLiteralToSnippet };
+export { ArrayExpression, coerceToSnippet, concretize, concretizeSnippets, numericLiteralToSnippet };

package/tgsl/math.js

@@ -1,6 +1,5 @@
 import { f32 } from "../data/numeric.js";
 import { abs, acos, acosh, asin, asinh, atan, atan2, atanh, ceil, cos, cosh, countLeadingZeros, exp, floor, log, log2, max, min, pow, sign, sin, sinh, sqrt, tan, tanh, trunc } from "../std/numeric.js";
-
 //#region src/tgsl/math.ts
 const mathToStd = {
 	abs,
@@ -40,6 +39,5 @@ const mathToStd = {
 	imul: void 0,
 	round: void 0
 };
-
 //#endregion
-export { mathToStd };
+export { mathToStd };

package/tgsl/shaderGenerator.d.ts

@@ -1,18 +1,20 @@
-import { Snippet } from "../data/snippet.js";
+import { ResolvedSnippet, Snippet } from "../data/snippet.js";
 import { GenerationCtx } from "./generationHelpers.js";
-import { ExternalMap } from "../core/resolve/externals.js";
 import { BaseData } from "../data/wgslTypes.js";
-import { Block, Expression, Statement } from "tinyest";
+import { Block } from "tinyest";
 
 //#region src/tgsl/shaderGenerator.d.ts
+/**
+ * **NOTE: This is an unstable API and may change in the future.**
+ *
+ * An interface meant to be used by other systems to generate snippets of
+ * shader code in the target language (WGSL, GLSL, etc.).
+ */
 interface ShaderGenerator {
   initGenerator(ctx: GenerationCtx): void;
-  block(body: Block, externalMap?: ExternalMap): string;
-  identifier(id: string): Snippet;
-  typedExpression(expression: Expression, expectedType: BaseData): Snippet;
-  expression(expression: Expression): Snippet;
-  statement(statement: Statement): string;
   functionDefinition(body: Block): string;
+  typeInstantiation(schema: BaseData, args: readonly Snippet[]): ResolvedSnippet;
+  typeAnnotation(schema: BaseData): string;
 }
 //#endregion
 export { ShaderGenerator };

package/tgsl/shaderGenerator_members.d.ts

@@ -0,0 +1,2 @@
+import { ResolutionCtx } from "../types.js";
+import { UnknownData } from "../data/dataTypes.js";

package/tgsl/shaderGenerator_members.js

@@ -0,0 +1,6 @@
+import { __exportAll } from "../_virtual/_rolldown/runtime.js";
+import { UnknownData } from "../data/dataTypes.js";
+//#region src/tgsl/shaderGenerator_members.ts
+var shaderGenerator_members_exports = /* @__PURE__ */ __exportAll({ UnknownData: () => UnknownData });
+//#endregion
+export { shaderGenerator_members_exports };

package/tgsl/shellless.js

@@ -2,11 +2,9 @@ import { getMetaData, getName } from "../shared/meta.js";
 import { isPtr, isWgslArray, isWgslStruct } from "../data/wgslTypes.js";
 import { UnknownData } from "../data/dataTypes.js";
 import { WgslTypeError } from "../errors.js";
-import { getResolutionCtx } from "../execMode.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;
@@ -28,10 +26,6 @@ var ShelllessRepository = class {
 			}
 			if (s.dataType === UnknownData) throw new Error(`Passed illegal value ${s.value} as the #${index} argument to ${meta.name}(...)\nShellless functions can only accept arguments representing WGSL resources: constructible WGSL types, d.refs, samplers or texture views.\nRemember, that arguments such as samplers, texture views, accessors, slots etc. should be dereferenced via '.$' first.`);
 			let type = concretize(s.dataType);
-			if (s.origin === "constant-tgpu-const-ref" || s.origin === "runtime-tgpu-const-ref") {
-				const ctx = getResolutionCtx();
-				throw new Error(`Cannot pass constant references as function arguments. Explicitly copy them by wrapping them in a schema: '${ctx.resolve(type).value}(...)'`);
-			}
 			if (isPtr(type) && type.implicit) type = type.inner;
 			return type;
 		});
@@ -48,6 +42,5 @@ var ShelllessRepository = class {
 		return shellless;
 	}
 };
-
 //#endregion
-export { ShelllessRepository };
+export { ShelllessRepository };

package/tgsl/wgslGenerator.d.ts

@@ -0,0 +1,36 @@
+import { Origin, ResolvedSnippet, Snippet } from "../data/snippet.js";
+import { GenerationCtx } from "./generationHelpers.js";
+import { ShaderGenerator } from "./shaderGenerator.js";
+import { BaseData, StorableData } from "../data/wgslTypes.js";
+import { UnknownData } from "../data/dataTypes.js";
+import { ExternalMap } from "../core/resolve/externals.js";
+import * as tinyest from "tinyest";
+
+//#region src/tgsl/wgslGenerator.d.ts
+declare class WgslGenerator implements ShaderGenerator {
+  #private;
+  initGenerator(ctx: GenerationCtx): void;
+  protected get ctx(): GenerationCtx;
+  _block([_, statements]: tinyest.Block, externalMap?: ExternalMap): string;
+  refVariable(id: string, dataType: StorableData): string;
+  blockVariable(varType: 'var' | 'let' | 'const', id: string, dataType: BaseData | UnknownData, origin: Origin): Snippet;
+  protected emitVarDecl(pre: string, keyword: 'var' | 'let' | 'const', name: string, _dataType: BaseData | UnknownData, rhsStr: string): string;
+  _identifier(id: string): Snippet;
+  /**
+   * A wrapper for `generateExpression` that updates `ctx.expectedType`
+   * and tries to convert the result when it does not match the expected type.
+   */
+  _typedExpression(expression: tinyest.Expression, expectedType: BaseData | BaseData[]): Snippet;
+  _expression(expression: tinyest.Expression): Snippet;
+  functionDefinition(body: tinyest.Block): 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
+   * through `ctx.resolve` to properly cache the result.
+   */
+  typeAnnotation(data: BaseData): string;
+  typeInstantiation(schema: BaseData, args: readonly Snippet[]): ResolvedSnippet;
+  _statement(statement: tinyest.Statement): string;
+}
+//#endregion
+export { WgslGenerator };

package/tgsl/wgslGenerator.js

@@ -1,8 +1,8 @@
 import { $gpuCallable, $internal, $providing, isMarkedInternal } from "../shared/symbols.js";
 import { getName } from "../shared/meta.js";
-import { Void, isNaturallyEphemeral, isPtr, isWgslArray, isWgslStruct } from "../data/wgslTypes.js";
+import { Void, isBool, isNaturallyEphemeral, isNumericSchema, isPtr, isVec, isWgslArray, isWgslStruct } from "../data/wgslTypes.js";
 import { ConsoleLog, InfixDispatch, UnknownData, isLooseData, unptr } from "../data/dataTypes.js";
-import { isEphemeralOrigin, isEphemeralSnippet, snip } from "../data/snippet.js";
+import { fallthroughCopyOrigin, isEphemeralOrigin, isEphemeralSnippet, snip } from "../data/snippet.js";
 import { ResolutionError, WgslTypeError, invariant } from "../errors.js";
 import { isGPUCallable, isKnownAtComptime } from "../types.js";
 import { stitch } from "../core/resolve/stitch.js";
@@ -11,8 +11,9 @@ import { RefOperator } from "../data/ref.js";
 import { safeStringify } from "../shared/stringify.js";
 import { convertStructValues, convertToCommonType, tryConvertSnippet } from "./conversion.js";
 import { bool, i32, u32 } from "../data/numeric.js";
-import { AutoStruct } from "../data/autoStruct.js";
 import { ArrayExpression, coerceToSnippet, concretize, numericLiteralToSnippet } from "./generationHelpers.js";
+import { vec2u, vec3u, vec4u } from "../data/vector.js";
+import { AutoStruct } from "../data/autoStruct.js";
 import { add, div, mul, neg, sub } from "../std/operators.js";
 import { accessProp } from "./accessProp.js";
 import { accessIndex } from "./accessIndex.js";
@@ -20,11 +21,12 @@ import { constant } from "../core/constant/tgpuConstant.js";
 import { isGenericFn } from "../core/function/tgpuFn.js";
 import { arrayOf } from "../data/array.js";
 import { pow } from "../std/numeric.js";
+import { resolveData } from "../core/resolve/resolveData.js";
 import { UnrollableIterable } from "../core/unroll/tgpuUnroll.js";
 import { mathToStd } from "./math.js";
-import { getElementCountSnippet, getElementSnippet, getElementType, getLoopVarKind } from "./forOfUtils.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 = [
@@ -61,6 +63,12 @@ const binaryLogicalOps = [
 	">",
 	">="
 ];
+const bitShiftOps = [
+	"<<",
+	">>",
+	"<<=",
+	">>="
+];
 const OP_MAP = {
 	"===": "==",
 	"!==": "!=",
@@ -106,7 +114,19 @@ function operatorToType(lhs, op, rhs) {
 }
 const unaryOpCodeToCodegen = {
 	"-": neg[$gpuCallable].call.bind(neg),
-	void: () => snip(void 0, Void, "constant")
+	void: () => snip(void 0, Void, "constant"),
+	"!": (ctx, [argExpr]) => {
+		if (argExpr === void 0) throw new Error("The unary operator `!` expects 1 argument, but 0 were provided.");
+		if (isKnownAtComptime(argExpr)) return snip(!argExpr.value, bool, "constant");
+		const { value, dataType } = argExpr;
+		const argStr = ctx.resolve(value, dataType).value;
+		if (isBool(dataType)) return snip(`!${argStr}`, bool, "runtime");
+		if (isNumericSchema(dataType)) {
+			const resultStr = `!bool(${argStr})`;
+			return snip(dataType.type === "f32" ? `(((bitcast<u32>(${argStr}) & 0x7fffffff) > 0x7f800000) || ${resultStr})` : dataType.type === "f16" ? `(((bitcast<u32>(${argStr}) & 0x7fff) > 0x7c00) || ${resultStr})` : resultStr, bool, "runtime");
+		}
+		return snip(false, bool, "constant");
+	}
 };
 const binaryOpCodeToCodegen = {
 	"+": add[$gpuCallable].call.bind(add),
@@ -125,7 +145,7 @@ var WgslGenerator = class {
 		if (!this.#ctx) throw new Error("WGSL Generator has not yet been initialized. Please call initialize(ctx) before using the generator.");
 		return this.#ctx;
 	}
-	block([_, statements], externalMap) {
+	_block([_, statements], externalMap) {
 		this.ctx.pushBlockScope();
 		if (externalMap) {
 			const externals = Object.fromEntries(Object.entries(externalMap).map(([id, value]) => [id, coerceToSnippet(value)]));
@@ -133,7 +153,7 @@ var WgslGenerator = class {
 		}
 		try {
 			this.ctx.indent();
-			const body = statements.map((statement) => this.statement(statement)).filter((statement) => statement.length > 0).join("\n");
+			const body = statements.map((statement) => this._statement(statement)).filter((statement) => statement.length > 0).join("\n");
 			this.ctx.dedent();
 			return `{
 ${body}
@@ -162,7 +182,10 @@ ${this.ctx.pre}}`;
 		this.ctx.defineVariable(id, snippet);
 		return snippet;
 	}
-	identifier(id) {
+	emitVarDecl(pre, keyword, name, _dataType, rhsStr) {
+		return `${pre}${keyword} ${name} = ${rhsStr};`;
+	}
+	_identifier(id) {
 		if (!id) throw new Error("Cannot resolve an empty identifier");
 		if (id === "undefined") return snip(void 0, Void, "constant");
 		const res = this.ctx.getById(id);
@@ -173,24 +196,24 @@ ${this.ctx.pre}}`;
 	* A wrapper for `generateExpression` that updates `ctx.expectedType`
 	* and tries to convert the result when it does not match the expected type.
 	*/
-	typedExpression(expression, expectedType) {
+	_typedExpression(expression, expectedType) {
 		const prevExpectedType = this.ctx.expectedType;
 		this.ctx.expectedType = expectedType;
 		try {
-			const result = this.expression(expression);
+			const result = this._expression(expression);
 			if (expectedType instanceof AutoStruct) return result;
 			return tryConvertSnippet(this.ctx, result, expectedType);
 		} finally {
 			this.ctx.expectedType = prevExpectedType;
 		}
 	}
-	expression(expression) {
-		if (typeof expression === "string") return this.identifier(expression);
+	_expression(expression) {
+		if (typeof expression === "string") return this._identifier(expression);
 		if (typeof expression === "boolean") return snip(expression, bool, "constant");
 		if (expression[0] === NODE.logicalExpr || expression[0] === NODE.binaryExpr || expression[0] === NODE.assignmentExpr) {
 			const [exprType, lhs, op, rhs] = expression;
-			const lhsExpr = this.expression(lhs);
-			const rhsExpr = this.expression(rhs);
+			const lhsExpr = this._expression(lhs);
+			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.`);
 			if (op === "==") throw new Error("Please use the === operator instead of ==");
 			if (op === "===" && isKnownAtComptime(lhsExpr) && isKnownAtComptime(rhsExpr)) return snip(lhsExpr.value === rhsExpr.value, bool, "constant");
@@ -198,8 +221,20 @@ ${this.ctx.pre}}`;
 			if (rhsExpr.dataType === UnknownData) throw new WgslTypeError(`Right-hand side of '${op}' is of unknown type`);
 			const codegen = binaryOpCodeToCodegen[op];
 			if (codegen) return codegen(this.ctx, [lhsExpr, rhsExpr]);
-			const forcedType = exprType === NODE.assignmentExpr ? [lhsExpr.dataType] : void 0;
-			const [convLhs, convRhs] = convertToCommonType(this.ctx, [lhsExpr, rhsExpr], forcedType) ?? [lhsExpr, rhsExpr];
+			let convLhs;
+			let convRhs;
+			if (bitShiftOps.includes(op)) {
+				let rhsTarget;
+				if (isVec(lhsExpr.dataType)) {
+					const cc = lhsExpr.dataType.componentCount;
+					rhsTarget = cc === 2 ? vec2u : cc === 3 ? vec3u : vec4u;
+				} else rhsTarget = u32;
+				convRhs = tryConvertSnippet(this.ctx, rhsExpr, rhsTarget, false);
+				convLhs = lhsExpr;
+			} else {
+				const forcedType = exprType === NODE.assignmentExpr ? [lhsExpr.dataType] : void 0;
+				[convLhs, convRhs] = convertToCommonType(this.ctx, [lhsExpr, rhsExpr], forcedType) ?? [lhsExpr, rhsExpr];
+			}
 			const lhsStr = this.ctx.resolve(convLhs.value, convLhs.dataType).value;
 			const rhsStr = this.ctx.resolve(convRhs.value, convRhs.dataType).value;
 			const type = operatorToType(convLhs.dataType, op, convRhs.dataType);
@@ -213,13 +248,13 @@ ${this.ctx.pre}}`;
 		}
 		if (expression[0] === NODE.postUpdate) {
 			const [_, op, arg] = expression;
-			const argExpr = this.expression(arg);
+			const argExpr = this._expression(arg);
 			const argStr = this.ctx.resolve(argExpr.value, argExpr.dataType).value;
 			return snip(`${argStr}${op}`, argExpr.dataType, "runtime");
 		}
 		if (expression[0] === NODE.unaryExpr) {
 			const [_, op, arg] = expression;
-			const argExpr = this.expression(arg);
+			const argExpr = this._expression(arg);
 			const codegen = unaryOpCodeToCodegen[op];
 			if (codegen) return codegen(this.ctx, [argExpr]);
 			const argStr = this.ctx.resolve(argExpr.value, argExpr.dataType).value;
@@ -228,7 +263,7 @@ ${this.ctx.pre}}`;
 		}
 		if (expression[0] === NODE.memberAccess) {
 			const [_, targetNode, property] = expression;
-			const target = this.expression(targetNode);
+			const target = this._expression(targetNode);
 			if (target.value === console) return snip(new ConsoleLog(property), UnknownData, "runtime");
 			if (target.value === Math) {
 				if (property in mathToStd && mathToStd[property]) return snip(mathToStd[property], UnknownData, "runtime");
@@ -240,8 +275,8 @@ ${this.ctx.pre}}`;
 		}
 		if (expression[0] === NODE.indexAccess) {
 			const [_, targetNode, propertyNode] = expression;
-			const target = this.expression(targetNode);
-			const inProperty = this.expression(propertyNode);
+			const target = this._expression(targetNode);
+			const inProperty = this._expression(propertyNode);
 			const property = convertToCommonType(this.ctx, [inProperty], [u32, i32], false)?.[0] ?? inProperty;
 			const accessed = accessIndex(target, property);
 			if (!accessed) {
@@ -258,27 +293,27 @@ ${this.ctx.pre}}`;
 		}
 		if (expression[0] === NODE.call) {
 			const [_, calleeNode, argNodes] = expression;
-			const callee = this.expression(calleeNode);
+			const callee = this._expression(calleeNode);
 			if (isWgslStruct(callee.value)) {
 				if (argNodes.length > 1) throw new WgslTypeError("Struct schemas should always be called with at most 1 argument");
 				if (!argNodes[0]) return snip(`${this.ctx.resolve(callee.value).value}()`, callee.value, "runtime");
-				const arg = this.typedExpression(argNodes[0], callee.value);
+				const arg = this._typedExpression(argNodes[0], callee.value);
 				return snip(this.ctx.resolve(arg.value, callee.value).value, callee.value, "runtime");
 			}
 			if (isWgslArray(callee.value)) {
 				if (argNodes.length > 1) throw new WgslTypeError("Array schemas should always be called with at most 1 argument");
 				if (!argNodes[0]) return snip(`${this.ctx.resolve(callee.value).value}()`, callee.value, "runtime");
-				const arg = this.typedExpression(argNodes[0], callee.value);
+				const arg = this._typedExpression(argNodes[0], callee.value);
 				if (arg.value instanceof ArrayExpression) return snip(stitch`${this.ctx.resolve(callee.value).value}(${arg.value.elements})`, arg.dataType, "runtime");
 				return snip(this.ctx.resolve(arg.value, callee.value).value, callee.value, "runtime");
 			}
 			if (callee.value === constant) throw new Error("Constants cannot be defined within TypeGPU function scope. To address this, move the constant definition outside the function scope.");
 			if (callee.value instanceof InfixDispatch) {
 				if (!argNodes[0]) throw new WgslTypeError(`An infix operator '${callee.value.name}' was called without any arguments`);
-				const rhs = this.expression(argNodes[0]);
+				const rhs = this._expression(argNodes[0]);
 				return callee.value.operator(this.ctx, [callee.value.lhs, rhs]);
 			}
-			if (callee.value instanceof ConsoleLog) return this.ctx.generateLog(callee.value.op, argNodes.map((arg) => this.expression(arg)));
+			if (callee.value instanceof ConsoleLog) return this.ctx.generateLog(callee.value.op, argNodes.map((arg) => this._expression(arg)));
 			if (isGPUCallable(callee.value)) {
 				const callable = callee.value[$gpuCallable];
 				const strictSignature = callable.strictSignature;
@@ -286,9 +321,9 @@ ${this.ctx.pre}}`;
 				if (strictSignature) convertedArguments = argNodes.map((arg, i) => {
 					const argType = strictSignature.argTypes[i];
 					if (!argType) throw new WgslTypeError(`Function '${getName(callee.value)}' was called with too many arguments`);
-					return this.typedExpression(arg, argType);
+					return this._typedExpression(arg, argType);
 				});
-				else convertedArguments = argNodes.map((arg) => this.expression(arg));
+				else convertedArguments = argNodes.map((arg) => this._expression(arg));
 				try {
 					return callable.call(this.ctx, convertedArguments);
 				} catch (err) {
@@ -301,7 +336,7 @@ ${this.ctx.pre}}`;
 				const slotPairs = isGeneric ? callee.value[$providing]?.pairs ?? [] : [];
 				const callback = isGeneric ? callee.value[$internal].inner : callee.value;
 				const shelllessCall = this.ctx.withRenamed(callback, getName(callee.value), () => this.ctx.withSlots(slotPairs, () => {
-					const args = argNodes.map((arg) => this.expression(arg));
+					const args = argNodes.map((arg) => this._expression(arg));
 					const shellless = this.ctx.shelllessRepo.get(callback, args);
 					if (!shellless) return;
 					const converted = args.map((s, idx) => {
@@ -324,11 +359,11 @@ ${this.ctx.pre}}`;
 				const entries = Object.fromEntries(Object.entries(obj).map(([key, value]) => {
 					let accessed = structType.accessProp(key);
 					let expr;
-					if (accessed) expr = this.typedExpression(value, accessed.type);
+					if (accessed) expr = this._typedExpression(value, accessed.type);
 					else {
-						expr = this.expression(value);
+						expr = this._expression(value);
 						if (expr.dataType === UnknownData) throw new WgslTypeError(stitch`Property ${key} in object literal has a value of unknown type: '${expr}'`);
-						accessed = structType.provideProp(key, concretize(expr.dataType));
+						accessed = structType.provideProp(key, unptr(concretize(expr.dataType)));
 					}
 					return [accessed.prop, expr];
 				}));
@@ -340,7 +375,7 @@ ${this.ctx.pre}}`;
 				const entries = Object.fromEntries(Object.entries(structType.propTypes).map(([key, value]) => {
 					const val = obj[key];
 					if (val === void 0) throw new WgslTypeError(`Missing property ${key} in object literal for struct ${structType}`);
-					return [key, this.typedExpression(val, value)];
+					return [key, this._typedExpression(val, value)];
 				}));
 				const convertedSnippets = convertStructValues(this.ctx, structType, entries);
 				return snip(stitch`${this.ctx.resolve(structType).value}(${convertedSnippets})`, structType, "runtime");
@@ -354,10 +389,10 @@ ${this.ctx.pre}}`;
 			let values;
 			if (isWgslArray(arrType)) {
 				elemType = arrType.elementType;
-				values = valueNodes.map((value) => this.typedExpression(value, elemType));
+				values = valueNodes.map((value) => this._typedExpression(value, elemType));
 				if (values.length !== arrType.elementCount) throw new WgslTypeError(`Cannot create value of type '${arrType}' from an array of length: ${values.length}`);
 			} else {
-				const valuesSnippets = valueNodes.map((value) => this.expression(value));
+				const valuesSnippets = valueNodes.map((value) => this._expression(value));
 				if (valuesSnippets.length === 0) throw new WgslTypeError("Cannot infer the type of an empty array literal.");
 				const converted = convertToCommonType(this.ctx, valuesSnippets);
 				if (!converted) throw new WgslTypeError("The given values cannot be automatically converted to a common type. Consider wrapping the array in an appropriate schema");
@@ -369,8 +404,8 @@ ${this.ctx.pre}}`;
 		}
 		if (expression[0] === NODE.conditionalExpr) {
 			const [_, test, consequent, alternative] = expression;
-			const testExpression = this.expression(test);
-			if (isKnownAtComptime(testExpression)) return testExpression.value ? this.expression(consequent) : this.expression(alternative);
+			const testExpression = this._expression(test);
+			if (isKnownAtComptime(testExpression)) return testExpression.value ? this._expression(consequent) : this._expression(alternative);
 			else throw new Error(`Ternary operator is only supported for comptime-known checks (used with '${testExpression.value}'). For runtime checks, please use 'std.select' or if/else statements.`);
 		}
 		if (expression[0] === NODE.stringLiteral) return snip(expression[1], UnknownData, "constant");
@@ -378,11 +413,23 @@ ${this.ctx.pre}}`;
 		assertExhaustive(expression);
 	}
 	functionDefinition(body) {
-		return this.block(body);
+		return this._block(body);
 	}
-	statement(statement) {
+	/**
+	* 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
+	* through `ctx.resolve` to properly cache the result.
+	*/
+	typeAnnotation(data) {
+		return resolveData(this.ctx, data);
+	}
+	typeInstantiation(schema, args) {
+		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 id = this._identifier(statement);
 			const resolved = id.value && this.ctx.resolve(id.value).value;
 			return resolved ? `${this.ctx.pre}${resolved};` : "";
 		}
@@ -391,7 +438,7 @@ ${this.ctx.pre}}`;
 			const returnNode = statement[1];
 			if (returnNode !== void 0) {
 				const expectedReturnType = this.ctx.topFunctionReturnType;
-				let returnSnippet = expectedReturnType ? this.typedExpression(returnNode, expectedReturnType) : this.expression(returnNode);
+				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") {
@@ -411,11 +458,17 @@ Try 'return ${typeStr}(${str});' instead.
 		}
 		if (statement[0] === NODE.if) {
 			const [_, condNode, consNode, altNode] = statement;
-			const condition = this.typedExpression(condNode, bool);
-			const consequent = condition.value === false ? void 0 : this.block(blockifySingleStatement(consNode));
-			const alternate = condition.value === true || !altNode ? void 0 : this.block(blockifySingleStatement(altNode));
-			if (condition.value === true) return `${this.ctx.pre}${consequent}`;
-			if (condition.value === false) return alternate ? `${this.ctx.pre}${alternate}` : "";
+			const condition = this._typedExpression(condNode, bool);
+			if (typeof condition.value === "boolean") {
+				let node = condition.value ? consNode : altNode;
+				if (node === void 0) return "";
+				if (!Array.isArray(node)) node = blockifySingleStatement(node);
+				if (node[0] === NODE.block && node[1].length === 1 && node[1][0][0] === NODE.if) return this._statement(node[1][0]);
+				if (node[0] === NODE.if) return this._statement(node);
+				return `${this.ctx.pre}${this._block(blockifySingleStatement(node))}`;
+			}
+			const consequent = this._block(blockifySingleStatement(consNode));
+			const alternate = !altNode ? void 0 : this._block(blockifySingleStatement(altNode));
 			if (!alternate) return stitch`${this.ctx.pre}if (${condition}) ${consequent}`;
 			return stitch`\
 ${this.ctx.pre}if (${condition}) ${consequent}
@@ -424,7 +477,7 @@ ${this.ctx.pre}else ${alternate}`;
 		if (statement[0] === NODE.let || statement[0] === NODE.const) {
 			let varType = "var";
 			const [stmtType, rawId, rawValue] = statement;
-			const eq = rawValue !== void 0 ? this.expression(rawValue) : void 0;
+			const eq = rawValue !== void 0 ? this._expression(rawValue) : void 0;
 			if (!eq) throw new Error(`Cannot create variable '${rawId}' without an initial value.`);
 			const ephemeral = isEphemeralSnippet(eq);
 			let dataType = eq.dataType;
@@ -472,9 +525,11 @@ ${this.ctx.pre}else ${alternate}`;
 				}
 			}
 			const snippet = this.blockVariable(varType, rawId, concretize(dataType), eq.origin);
-			return stitch`${this.ctx.pre}${varType} ${snippet.value} = ${tryConvertSnippet(this.ctx, eq, dataType, false)};`;
+			const rhsSnippet = tryConvertSnippet(this.ctx, eq, dataType, false);
+			const rhsStr = this.ctx.resolve(rhsSnippet.value, rhsSnippet.dataType).value;
+			return this.emitVarDecl(this.ctx.pre, varType, snippet.value, concretize(dataType), rhsStr);
 		}
-		if (statement[0] === NODE.block) return `${this.ctx.pre}${this.block(statement)}`;
+		if (statement[0] === NODE.block) return `${this.ctx.pre}${this._block(statement)}`;
 		if (statement[0] === NODE.for) {
 			const [_, init, condition, update, body] = statement;
 			const prevUnrollingFlag = this.#unrolling;
@@ -482,13 +537,13 @@ ${this.ctx.pre}else ${alternate}`;
 			try {
 				this.ctx.pushBlockScope();
 				const [initStatement, conditionExpr, updateStatement] = this.ctx.withResetIndentLevel(() => [
-					init ? this.statement(init) : void 0,
-					condition ? this.typedExpression(condition, bool) : void 0,
-					update ? this.statement(update) : void 0
+					init ? this._statement(init) : void 0,
+					condition ? this._typedExpression(condition, bool) : void 0,
+					update ? this._statement(update) : void 0
 				]);
 				const initStr = initStatement ? initStatement.slice(0, -1) : "";
 				const updateStr = updateStatement ? updateStatement.slice(0, -1) : "";
-				const bodyStr = this.block(blockifySingleStatement(body));
+				const bodyStr = this._block(blockifySingleStatement(body));
 				return stitch`${this.ctx.pre}for (${initStr}; ${conditionExpr}; ${updateStr}) ${bodyStr}`;
 			} finally {
 				this.#unrolling = prevUnrollingFlag;
@@ -500,9 +555,9 @@ ${this.ctx.pre}else ${alternate}`;
 			this.#unrolling = false;
 			try {
 				const [_, condition, body] = statement;
-				const condSnippet = this.typedExpression(condition, bool);
+				const condSnippet = this._typedExpression(condition, bool);
 				const conditionStr = this.ctx.resolve(condSnippet.value).value;
-				const bodyStr = this.block(blockifySingleStatement(body));
+				const bodyStr = this._block(blockifySingleStatement(body));
 				return `${this.ctx.pre}while (${conditionStr}) ${bodyStr}`;
 			} finally {
 				this.#unrolling = prevUnrollingFlag;
@@ -515,40 +570,40 @@ ${this.ctx.pre}else ${alternate}`;
 			const prevUnrollingFlag = this.#unrolling;
 			try {
 				this.ctx.pushBlockScope();
-				const iterableExpr = this.expression(iterable);
+				const iterableExpr = this._expression(iterable);
 				const shouldUnroll = iterableExpr.value instanceof UnrollableIterable;
 				const iterableSnippet = shouldUnroll ? iterableExpr.value.snippet : iterableExpr;
-				const elementCountSnippet = getElementCountSnippet(this.ctx, iterableSnippet, shouldUnroll);
+				const range = getRangeSnippets(this.ctx, iterableSnippet, shouldUnroll);
 				const originalLoopVarName = loopVar[1];
 				const blockified = blockifySingleStatement(body);
 				if (shouldUnroll) {
-					if (!isKnownAtComptime(elementCountSnippet)) throw new Error("Cannot unroll loop. Length of iterable is unknown at comptime.");
+					if (!isKnownAtComptime(range.end)) throw new Error("Cannot unroll loop. Length of iterable is unknown at comptime.");
 					this.#unrolling = true;
-					const length = elementCountSnippet.value;
+					const length = range.end.value;
 					if (length === 0) return "";
 					const { value } = iterableSnippet;
-					const elements = 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 }, (_, 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");
+					return elements.map((e, i) => `${this.ctx.pre}// unrolled iteration #${i}\n${this.ctx.pre}${this._block(blockified, { [originalLoopVarName]: e })}`).join("\n");
 				}
-				if (isEphemeralSnippet(iterableSnippet)) throw new Error(`\`for ... of ...\` loops only support iterables stored in variables.
-  -----
-  You can wrap iterable with \`tgpu.unroll(...)\`. If iterable is known at comptime, the loop will be unrolled.
-  -----`);
 				this.#unrolling = false;
 				const index = this.ctx.makeNameValid("i");
-				const elementSnippet = getElementSnippet(iterableSnippet, snip(index, u32, "runtime"));
-				const loopVarName = this.ctx.makeNameValid(originalLoopVarName);
-				const loopVarKind = getLoopVarKind(elementSnippet);
-				const elementType = getElementType(elementSnippet, iterableSnippet);
-				const forHeaderStr = stitch`${this.ctx.pre}for (var ${index} = 0u; ${index} < ${elementCountSnippet}; ${index}++) {`;
-				this.ctx.indent();
-				ctxIndent = true;
-				const loopVarDeclStr = stitch`${this.ctx.pre}${loopVarKind} ${loopVarName} = ${tryConvertSnippet(this.ctx, elementSnippet, elementType, false)};`;
-				const bodyStr = `${this.ctx.pre}${this.block(blockified, { [originalLoopVarName]: snip(loopVarName, elementType, elementSnippet.origin) })}`;
-				this.ctx.dedent();
-				ctxIndent = false;
-				return stitch`${forHeaderStr}\n${loopVarDeclStr}\n${bodyStr}\n${this.ctx.pre}}`;
+				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") });
+				else {
+					this.ctx.indent();
+					ctxIndent = true;
+					const loopVarName = this.ctx.makeNameValid(originalLoopVarName);
+					const elementSnippet = getElementSnippet(iterableSnippet, snip(index, u32, "runtime"));
+					const loopVarKind = getLoopVarKind(elementSnippet);
+					const elementType = getElementType(elementSnippet, iterableSnippet);
+					bodyStr = `{\n${stitch`${this.ctx.pre}${loopVarKind} ${loopVarName} = ${tryConvertSnippet(this.ctx, elementSnippet, elementType, false)};`}\n${this.ctx.pre}${this._block(blockified, { [originalLoopVarName]: snip(loopVarName, elementType, elementSnippet.origin) })}\n`;
+					this.ctx.dedent();
+					bodyStr += `${this.ctx.pre}}`;
+					ctxIndent = false;
+				}
+				return stitch`${forHeaderStr} ${bodyStr.trim()}`;
 			} finally {
 				if (ctxIndent) this.ctx.dedent();
 				this.#unrolling = prevUnrollingFlag;
@@ -563,7 +618,7 @@ ${this.ctx.pre}else ${alternate}`;
 			if (this.#unrolling) throw new WgslTypeError("Cannot unroll loop containing `break`");
 			return `${this.ctx.pre}break;`;
 		}
-		const expr = this.expression(statement);
+		const expr = this._expression(statement);
 		const resolved = expr.value && this.ctx.resolve(expr.value).value;
 		return resolved ? `${this.ctx.pre}${resolved};` : "";
 	}
@@ -580,6 +635,5 @@ function blockifySingleStatement(statement) {
 	return typeof statement !== "object" || statement[0] !== NODE.block ? [NODE.block, [statement]] : statement;
 }
 const wgslGenerator = new WgslGenerator();
-
 //#endregion
-export { wgslGenerator as default };
+export { WgslGenerator, wgslGenerator as default };