From fc3fa83179f70c6a31c47ec507bd23433c6e6cec Mon Sep 17 00:00:00 2001 From: William Candillon Date: Tue, 18 Aug 2026 21:57:49 +0200 Subject: [PATCH 1/3] feat: Support default parameter values in TGSL functions Allows 'use gpu' functions to declare defaulted identifier parameters, e.g. (t: number, cyclic = false) => ... - tinyest: FuncParameter identifier variant carries an optional default Expression. - tinyest-for-wgsl: AssignmentPattern params (with an Identifier target) transpile their default expression; identifiers used only in defaults are gathered as externals, other params are not. - typegpu: when a call site provides fewer arguments than declared params, the trailing defaults are prepended to the function body as const statements, so each arity resolves to its own concrete overload. - eslint-plugin: no-unsupported-syntax no longer flags defaulted identifier params, still flags destructuring defaults. Co-Authored-By: Claude Fable 5 --- .../src/rules/noUnsupportedSyntax.ts | 11 +- .../tests/rules/noUnsupportedSyntax.test.ts | 21 ++-- .../tinyest-for-wgsl/tests/parsers.test.ts | 35 ++++++ packages/tinyest/src/nodes.ts | 5 + packages/typegpu/src/resolutionCtx.ts | 20 +++- packages/typegpu/tests/tgsl/shellless.test.ts | 109 ++++++++++++++++++ 6 files changed, 188 insertions(+), 13 deletions(-) diff --git a/packages/eslint-plugin/src/rules/noUnsupportedSyntax.ts b/packages/eslint-plugin/src/rules/noUnsupportedSyntax.ts index be092ff776..bb6510226e 100644 --- a/packages/eslint-plugin/src/rules/noUnsupportedSyntax.ts +++ b/packages/eslint-plugin/src/rules/noUnsupportedSyntax.ts @@ -49,7 +49,16 @@ export const noUnsupportedSyntax = createRule({ if (!directives.getEnclosingTypegpuFunction()) { return; } - report(node, 'assignment pattern (default parameter)'); + // Default values for plain identifier parameters are supported; + // defaults inside destructuring patterns are not. + const isDefaultedIdentifierParam = + (node.parent.type === 'ArrowFunctionExpression' || + node.parent.type === 'FunctionExpression' || + node.parent.type === 'FunctionDeclaration') && + node.left.type === 'Identifier'; + if (!isDefaultedIdentifierParam) { + report(node, 'assignment pattern (destructuring default)'); + } }, AwaitExpression(node) { diff --git a/packages/eslint-plugin/tests/rules/noUnsupportedSyntax.test.ts b/packages/eslint-plugin/tests/rules/noUnsupportedSyntax.test.ts index 1e18b8139e..e53dea908c 100644 --- a/packages/eslint-plugin/tests/rules/noUnsupportedSyntax.test.ts +++ b/packages/eslint-plugin/tests/rules/noUnsupportedSyntax.test.ts @@ -11,6 +11,9 @@ describe('noUnsupportedSyntax', () => { "const cls = new (class { #priv = 1; fn = () => { 'use gpu'; const a = this.#priv; } } )()", "const fn = () => { 'use gpu'; switch (x) { case 1: break; } }", "const fn = () => { 'use gpu'; const obj = { [key]: 1 }; }", + "const fn = (arg = 1) => { 'use gpu'; }", + "function fn(arg = 1) { 'use gpu'; }", + "const fn = function(arg = 1) { 'use gpu'; }", ], invalid: [ { @@ -44,29 +47,27 @@ describe('noUnsupportedSyntax', () => { ], }, { - code: "const fn = (arg = 1) => { 'use gpu'; }", + code: "const fn = ({ arg = 1 }) => { 'use gpu'; }", errors: [ { messageId: 'unexpected', - data: { snippet: 'arg = 1', syntax: 'assignment pattern (default parameter)' }, + data: { snippet: 'arg = 1', syntax: 'assignment pattern (destructuring default)' }, }, ], }, { - code: "function fn(arg = 1) { 'use gpu'; }", + code: "const fn = () => { 'use gpu'; const { a = 1 } = foo; }", errors: [ { messageId: 'unexpected', - data: { snippet: 'arg = 1', syntax: 'assignment pattern (default parameter)' }, + data: { + snippet: '{ a = 1 } = foo', + syntax: 'variable declaration using destructuring', + }, }, - ], - }, - { - code: "const fn = function(arg = 1) { 'use gpu'; }", - errors: [ { messageId: 'unexpected', - data: { snippet: 'arg = 1', syntax: 'assignment pattern (default parameter)' }, + data: { snippet: 'a = 1', syntax: 'assignment pattern (destructuring default)' }, }, ], }, diff --git a/packages/tinyest-for-wgsl/tests/parsers.test.ts b/packages/tinyest-for-wgsl/tests/parsers.test.ts index 8eca3814d1..d80855fe61 100644 --- a/packages/tinyest-for-wgsl/tests/parsers.test.ts +++ b/packages/tinyest-for-wgsl/tests/parsers.test.ts @@ -95,6 +95,41 @@ describe('transpileBabelFn and transpileAcornFn', () => { }), ); + it( + 'supports default parameter values', + dualTest((p) => { + const { params, body, externalNames } = transpileFn(p('(a, b = 2, c = e) => a + b + c')); + + expect(params).toStrictEqual([ + { type: 'i', name: 'a' }, + { type: 'i', name: 'b', default: [5, '2'] }, + { type: 'i', name: 'c', default: 'e' }, + ]); + expect(JSON.stringify(body)).toMatchInlineSnapshot( + `"[0,[[10,[1,[1,"a","+","b"],"+","c"]]]]"`, + ); + // 'e' is external, even though it only appears in a default value. + expect(externalNames).toMatchInlineSnapshot(` + Map { + "e" => "e", + } + `); + }), + ); + + it( + 'does not treat params referenced in default values as external names', + dualTest((p) => { + const { params, externalNames } = transpileFn(p('(a, b = a) => a + b')); + + expect(params).toStrictEqual([ + { type: 'i', name: 'a' }, + { type: 'i', name: 'b', default: 'a' }, + ]); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }), + ); + it( 'respects local declarations when gathering external names', dualTest((p, transpileFn) => { diff --git a/packages/tinyest/src/nodes.ts b/packages/tinyest/src/nodes.ts index 3d57bb0658..89121876fa 100644 --- a/packages/tinyest/src/nodes.ts +++ b/packages/tinyest/src/nodes.ts @@ -290,6 +290,11 @@ export type FuncParameter = | { type: typeof FuncParameterType.identifier; name: string; + /** + * The parameter's default value expression (`(a, b = 2) => ...`), + * used when a call site omits the argument. + */ + default?: Expression | undefined; } | { type: typeof FuncParameterType.destructuredObject; diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index c183d04cdd..173cf8ab8b 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -58,7 +58,7 @@ import type { import { CodegenState, isSelfResolvable, NormalState, type FunctionArgument } from './types.ts'; import type { WgslEnableExtension } from './wgslExtensions.ts'; import { getName, hasTinyestMetadata, isNamable, setName } from './shared/meta.ts'; -import { FuncParameterType } from 'tinyest'; +import { FuncParameterType, NodeTypeCatalog, type Block, type Statement } from 'tinyest'; import { accessProp } from './tgsl/accessProp.ts'; import { createIoSchema } from './core/function/ioSchema.ts'; import { isShelllessImpl } from './core/function/shelllessImpl.ts'; @@ -734,6 +734,22 @@ export class ResolutionCtxImpl implements ResolutionCtx { } } + // Parameters beyond the call site's arguments fall back to their + // default value expressions (`(a, b = 2) => ...`), prepended to the + // body as const statements so they generate in the function's scope. + let body = options.body; + if (!options.entryInput) { + const defaultStatements: Statement[] = []; + for (const param of options.params.slice(options.argTypes.length)) { + if (param.type === FuncParameterType.identifier && param.default !== undefined) { + defaultStatements.push([NodeTypeCatalog.const, param.name, param.default]); + } + } + if (defaultStatements.length > 0) { + body = [NodeTypeCatalog.block, [...defaultStatements, ...body[1]]] as Block; + } + } + let returnType: BaseData | undefined; const code = this.gen.functionDefinition({ @@ -741,7 +757,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { name: options.name, workgroupSize: options.workgroupSize, args, - body: options.body, + body, determineReturnType: () => { if (returnType) { // Already determined diff --git a/packages/typegpu/tests/tgsl/shellless.test.ts b/packages/typegpu/tests/tgsl/shellless.test.ts index e9edc0fb7f..84098f3ba3 100644 --- a/packages/typegpu/tests/tgsl/shellless.test.ts +++ b/packages/typegpu/tests/tgsl/shellless.test.ts @@ -511,4 +511,113 @@ describe('shellless', () => { Remember, that arguments such as samplers, texture views, accessors, slots etc. should be dereferenced via '.$' first.] `); }); + + it('falls back to default parameter values when arguments are omitted', () => { + const scale = (v: d.v2f, factor: number = 2) => { + 'use gpu'; + return std.mul(v, d.f32(factor)); + }; + + const foo = () => { + 'use gpu'; + return std.add(scale(d.vec2f(1, 2)), scale(d.vec2f(3, 4), 0.5)); + }; + + expect(tgpu.resolve([foo])).toMatchInlineSnapshot(` + "fn scale(v: vec2f) -> vec2f { + const factor = 2; + return (v * f32(factor)); + } + + fn scale_1(v: vec2f, factor: f32) -> vec2f { + return (v * factor); + } + + fn foo() -> vec2f { + return (scale(vec2f(1, 2)) + scale_1(vec2f(3, 4), 0.5f)); + }" + `); + }); + + it('supports comptime-style mode flags with default values', () => { + const interpolate = (t: number, cyclic: boolean = false) => { + 'use gpu'; + if (cyclic) { + return std.fract(t); + } + return std.clamp(t, 0, 1); + }; + + const foo = () => { + 'use gpu'; + return interpolate(1.5) + interpolate(2.5, true); + }; + + expect(tgpu.resolve([foo])).toMatchInlineSnapshot(` + "fn interpolate(t: f32) -> f32 { + const cyclic = false; + if (cyclic) { + return fract(t); + } + return clamp(t, 0f, 1f); + } + + fn interpolate_1(t: f32, cyclic: bool) -> f32 { + if (cyclic) { + return fract(t); + } + return clamp(t, 0f, 1f); + } + + fn foo() -> f32 { + return (interpolate(1.5f) + interpolate_1(2.5f, true)); + }" + `); + }); + + it('evaluates default values referencing earlier parameters', () => { + const madd = (a: d.v2f, b: d.v2f = a) => { + 'use gpu'; + return std.add(a, b); + }; + + const foo = () => { + 'use gpu'; + return madd(d.vec2f(1, 2)); + }; + + expect(tgpu.resolve([foo])).toMatchInlineSnapshot(` + "fn madd(a: vec2f) -> vec2f { + let b = a; + return (a + b); + } + + fn foo() -> vec2f { + return madd(vec2f(1, 2)); + }" + `); + }); + + it('falls back to default parameter values in shelled functions', () => { + const addBias = tgpu.fn( + [d.f32], + d.f32, + )((x, bias: number = 10) => { + 'use gpu'; + return x + d.f32(bias); + }); + + const main = tgpu.fn([], d.f32)(() => addBias(1)); + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "fn addBias(x: f32) -> f32 { + const bias = 10; + return (x + f32(bias)); + } + + fn main() -> f32 { + return addBias(1f); + }" + `); + }); }); From 66fa72348545d409496851aa6920af2b65cc9493 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Wed, 21 Oct 2026 18:19:29 +0200 Subject: [PATCH 2/3] refactor: Adapt default parameter values to refactored tinyest-for-wgsl Moves default value extraction into functionParts.ts and transpiles defaults in createParser's transpileFn, so the acorn, babel and legacy parsers all support them. --- .../tinyest-for-wgsl/src/functionParts.ts | 43 ++++++++++++++----- packages/tinyest-for-wgsl/src/parsers.ts | 11 ++++- .../tinyest-for-wgsl/tests/parsers.test.ts | 4 +- 3 files changed, 45 insertions(+), 13 deletions(-) diff --git a/packages/tinyest-for-wgsl/src/functionParts.ts b/packages/tinyest-for-wgsl/src/functionParts.ts index 4fad4b9330..c9852125ce 100644 --- a/packages/tinyest-for-wgsl/src/functionParts.ts +++ b/packages/tinyest-for-wgsl/src/functionParts.ts @@ -72,7 +72,11 @@ function validateFunction(functionNode: FunctionNode): void { const unsupportedTypes = new Set( functionNode.params.flatMap((param) => - param.type === 'ObjectPattern' || param.type === 'Identifier' ? [] : [param.type], + param.type === 'ObjectPattern' || + param.type === 'Identifier' || + (param.type === 'AssignmentPattern' && param.left.type === 'Identifier') + ? [] + : [param.type], ), ); if (unsupportedTypes.size > 0) { @@ -80,15 +84,16 @@ function validateFunction(functionNode: FunctionNode): void { } } +type SupportedParam = + | babel.Identifier + | acorn.Identifier + | babel.ObjectPattern + | acorn.ObjectPattern + | babel.AssignmentPattern + | acorn.AssignmentPattern; + function parseParams(functionNode: FunctionNode): tinyest.FuncParameter[] { - return ( - functionNode.params as ( - | babel.Identifier - | acorn.Identifier - | babel.ObjectPattern - | acorn.ObjectPattern - )[] - ).map((param) => + return (functionNode.params as SupportedParam[]).map((param) => param.type === 'ObjectPattern' ? { type: tinyest.FuncParameterType.destructuredObject, @@ -102,13 +107,30 @@ function parseParams(functionNode: FunctionNode): tinyest.FuncParameter[] { } : { type: tinyest.FuncParameterType.identifier, - name: param.name, + name: + param.type === 'AssignmentPattern' + ? (param.left as babel.Identifier | acorn.Identifier).name + : param.name, }, ); } +function parseParamDefaults( + functionNode: FunctionNode, +): (babel.Expression | acorn.Expression | undefined)[] { + return (functionNode.params as SupportedParam[]).map((param) => + param.type === 'AssignmentPattern' ? param.right : undefined, + ); +} + export function extractFunctionParts(rootNode: JsNode): { params: tinyest.FuncParameter[]; + /** + * Raw default value expressions (`(a, b = 2) => ...`), aligned by index + * with `params`. Transpiled and attached to the params by `transpileFn`, + * which owns the transpilation context. + */ + paramDefaults: (babel.Expression | acorn.Expression | undefined)[]; body: acorn.BlockStatement | acorn.Expression | babel.BlockStatement | babel.Expression; } { const functionNode = unwrapToFunction(rootNode); @@ -117,6 +139,7 @@ export function extractFunctionParts(rootNode: JsNode): { return { params: parseParams(functionNode), + paramDefaults: parseParamDefaults(functionNode), body: functionNode.body, }; } diff --git a/packages/tinyest-for-wgsl/src/parsers.ts b/packages/tinyest-for-wgsl/src/parsers.ts index 9d13f948d7..08d4dece77 100644 --- a/packages/tinyest-for-wgsl/src/parsers.ts +++ b/packages/tinyest-for-wgsl/src/parsers.ts @@ -121,9 +121,18 @@ function createParser(kind: 'acorn' | 'babel' | 'legacy') { return { transpileFn(rootNode: JsNode, options: TranspilationOptions): TranspilationResult { - const { params, body } = extractFunctionParts(rootNode); + const { params, paramDefaults, body } = extractFunctionParts(rootNode); const ctx = createContext(params, options); + // Transpiled with the param names already declared, so identifiers used in + // a default value only count as externals when they're not other params. + for (const [i, defaultNode] of paramDefaults.entries()) { + const param = params[i]; + if (defaultNode && param?.type === tinyest.FuncParameterType.identifier) { + param.default = transpile(ctx, defaultNode) as tinyest.Expression; + } + } + const tinyestBody = transpile(ctx, body); if (body.type === 'BlockStatement') { diff --git a/packages/tinyest-for-wgsl/tests/parsers.test.ts b/packages/tinyest-for-wgsl/tests/parsers.test.ts index d80855fe61..a283b3dbfd 100644 --- a/packages/tinyest-for-wgsl/tests/parsers.test.ts +++ b/packages/tinyest-for-wgsl/tests/parsers.test.ts @@ -97,7 +97,7 @@ describe('transpileBabelFn and transpileAcornFn', () => { it( 'supports default parameter values', - dualTest((p) => { + dualTest((p, transpileFn) => { const { params, body, externalNames } = transpileFn(p('(a, b = 2, c = e) => a + b + c')); expect(params).toStrictEqual([ @@ -119,7 +119,7 @@ describe('transpileBabelFn and transpileAcornFn', () => { it( 'does not treat params referenced in default values as external names', - dualTest((p) => { + dualTest((p, transpileFn) => { const { params, externalNames } = transpileFn(p('(a, b = a) => a + b')); expect(params).toStrictEqual([ From ef6a5135a47fc3aa8dc3b9577cbe538f1903ef1e Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Fri, 23 Oct 2026 22:03:32 +0200 Subject: [PATCH 3/3] feat: Inline comptime-known default parameters, support passing `undefined` - Default values of omitted parameters are passed to `functionDefinition` instead of being prepended to the tinyest body. Defaults known at comptime are inlined at every use; only runtime defaults are declared at the top of the function body. - Passing `undefined` to a shellless function's defaulted parameter is treated as omitting it (`fn(a, undefined)` resolves the same as `fn(a)`). - Default values must have a known WGSL data type. - Shellless functions whose every parameter has a default value can be resolved directly. - The parser no longer mutates params after creating them. Co-Authored-By: Claude Opus 5.5 --- packages/tinyest-for-wgsl/src/parsers.ts | 16 +- packages/typegpu-gl/src/glslGenerator.ts | 7 +- .../typegpu-gl/tests/glslGenerator.test.ts | 23 +++ packages/typegpu/src/resolutionCtx.ts | 35 +++-- packages/typegpu/src/tgsl/shaderGenerator.ts | 12 +- packages/typegpu/src/tgsl/shellless.ts | 28 +++- packages/typegpu/src/tgsl/wgslGenerator.ts | 59 ++++++- packages/typegpu/src/types.ts | 4 + packages/typegpu/tests/tgsl/shellless.test.ts | 148 ++++++++++++++++-- 9 files changed, 286 insertions(+), 46 deletions(-) diff --git a/packages/tinyest-for-wgsl/src/parsers.ts b/packages/tinyest-for-wgsl/src/parsers.ts index 08d4dece77..0e66e4d098 100644 --- a/packages/tinyest-for-wgsl/src/parsers.ts +++ b/packages/tinyest-for-wgsl/src/parsers.ts @@ -121,17 +121,17 @@ function createParser(kind: 'acorn' | 'babel' | 'legacy') { return { transpileFn(rootNode: JsNode, options: TranspilationOptions): TranspilationResult { - const { params, paramDefaults, body } = extractFunctionParts(rootNode); - const ctx = createContext(params, options); + const { params: rawParams, paramDefaults, body } = extractFunctionParts(rootNode); + const ctx = createContext(rawParams, options); // Transpiled with the param names already declared, so identifiers used in // a default value only count as externals when they're not other params. - for (const [i, defaultNode] of paramDefaults.entries()) { - const param = params[i]; - if (defaultNode && param?.type === tinyest.FuncParameterType.identifier) { - param.default = transpile(ctx, defaultNode) as tinyest.Expression; - } - } + const params = rawParams.map((param, i): tinyest.FuncParameter => { + const defaultNode = paramDefaults[i]; + return defaultNode && param.type === tinyest.FuncParameterType.identifier + ? { ...param, default: transpile(ctx, defaultNode) as tinyest.Expression } + : param; + }); const tinyestBody = transpile(ctx, body); diff --git a/packages/typegpu-gl/src/glslGenerator.ts b/packages/typegpu-gl/src/glslGenerator.ts index d36cd821cc..5ffe1b4c7d 100644 --- a/packages/typegpu-gl/src/glslGenerator.ts +++ b/packages/typegpu-gl/src/glslGenerator.ts @@ -1284,7 +1284,12 @@ The generated shader will omit it, so its runtime side effects will not occur.`) } try { - const body = this._block(options.body, /* allowInlining */ false).code; + const body = this._block( + options.body, + /* allowInlining */ false, + /* externalMap */ undefined, + options.paramDefaults, + ).code; const returnType = options.determineReturnType(); if (options.functionType !== 'normal') { diff --git a/packages/typegpu-gl/tests/glslGenerator.test.ts b/packages/typegpu-gl/tests/glslGenerator.test.ts index 80fef15f58..399c7f695e 100644 --- a/packages/typegpu-gl/tests/glslGenerator.test.ts +++ b/packages/typegpu-gl/tests/glslGenerator.test.ts @@ -82,6 +82,29 @@ describe('GlslGenerator - variable declarations', () => { expect(result.code).not.toContain('var '); }); + it('supports default parameter values', () => { + const shift = (v: d.v2f, amount: number = 2, by: d.v2f = v) => { + 'use gpu'; + return std.add(v, std.mul(by, d.f32(amount))); + }; + + const foo = () => { + 'use gpu'; + return shift(d.vec2f(1, 2)); + }; + + expect(tgpu.resolve([foo], glOptions())).toMatchInlineSnapshot(` + "vec2 shift(vec2 v) { + vec2 by = v; + return (v + (by * 2.0)); + } + + vec2 foo() { + return shift(vec2(1, 2)); + }" + `); + }); + it('translates f32 variable declaration to float', () => { const fragFn = tgpu.fragmentFn({ out: d.vec4f, diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index 173cf8ab8b..9481b2bdef 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -20,7 +20,7 @@ import { type Snippet, withValue, } from './data/snippet.ts'; -import { type BaseData, isPtr, isWgslArray, isWgslStruct, Void } from './data/wgslTypes.ts'; +import { type BaseData, isPtr, isVoid, isWgslArray, isWgslStruct, Void } from './data/wgslTypes.ts'; import { invariant, MissingSlotValueError, ResolutionError, WgslTypeError } from './errors.ts'; import { provideCtx, topLevelState } from './execMode.ts'; import { naturalsExcept } from './shared/generators.ts'; @@ -39,7 +39,7 @@ import { LogGeneratorImpl, LogGeneratorNullImpl } from './tgsl/consoleLog/logGen import type { LogGenerator, LogResources, SupportedLogOp } from './tgsl/consoleLog/types.ts'; import { getBestConversion } from './tgsl/conversion.ts'; import { coerceToSnippet, concretize, numericLiteralToSnippet } from './tgsl/generationHelpers.ts'; -import type { ShaderGenerator } from './tgsl/shaderGenerator.ts'; +import type { ParamDefault, ShaderGenerator } from './tgsl/shaderGenerator.ts'; import { WgslGenerator } from './tgsl/wgslGenerator.ts'; import type { BlockScopeLayer, @@ -58,7 +58,7 @@ import type { import { CodegenState, isSelfResolvable, NormalState, type FunctionArgument } from './types.ts'; import type { WgslEnableExtension } from './wgslExtensions.ts'; import { getName, hasTinyestMetadata, isNamable, setName } from './shared/meta.ts'; -import { FuncParameterType, NodeTypeCatalog, type Block, type Statement } from 'tinyest'; +import { FuncParameterType } from 'tinyest'; import { accessProp } from './tgsl/accessProp.ts'; import { createIoSchema } from './core/function/ioSchema.ts'; import { isShelllessImpl } from './core/function/shelllessImpl.ts'; @@ -674,6 +674,10 @@ export class ResolutionCtxImpl implements ResolutionCtx { } } else { for (const [i, argType] of options.argTypes.entries()) { + if (isVoid(argType)) { + // Omitted at the call site, the parameter's default value is used instead + continue; + } const astParam = options.params[i]; // We know if arguments are passed by reference or by value, because we // enforce that based on the whether the argument is a pointer or not. @@ -734,20 +738,20 @@ export class ResolutionCtxImpl implements ResolutionCtx { } } - // Parameters beyond the call site's arguments fall back to their - // default value expressions (`(a, b = 2) => ...`), prepended to the - // body as const statements so they generate in the function's scope. - let body = options.body; + // Parameters omitted at the call site fall back to their default values (`(a, b = 2) => ...`) + const paramDefaults: ParamDefault[] = []; if (!options.entryInput) { - const defaultStatements: Statement[] = []; - for (const param of options.params.slice(options.argTypes.length)) { - if (param.type === FuncParameterType.identifier && param.default !== undefined) { - defaultStatements.push([NodeTypeCatalog.const, param.name, param.default]); + for (const [i, param] of options.params.entries()) { + const argType = options.argTypes[i]; + const omitted = argType === undefined || isVoid(argType); + if ( + omitted && + param.type === FuncParameterType.identifier && + param.default !== undefined + ) { + paramDefaults.push({ name: param.name, expression: param.default }); } } - if (defaultStatements.length > 0) { - body = [NodeTypeCatalog.block, [...defaultStatements, ...body[1]]] as Block; - } } let returnType: BaseData | undefined; @@ -757,7 +761,8 @@ export class ResolutionCtxImpl implements ResolutionCtx { name: options.name, workgroupSize: options.workgroupSize, args, - body, + body: options.body, + paramDefaults, determineReturnType: () => { if (returnType) { // Already determined diff --git a/packages/typegpu/src/tgsl/shaderGenerator.ts b/packages/typegpu/src/tgsl/shaderGenerator.ts index b7acd356b7..b7b2da6d99 100644 --- a/packages/typegpu/src/tgsl/shaderGenerator.ts +++ b/packages/typegpu/src/tgsl/shaderGenerator.ts @@ -1,4 +1,4 @@ -import type { Block } from 'tinyest'; +import type { Block, Expression } from 'tinyest'; import type { BaseData } from '../data/wgslTypes.ts'; import type { ResolvedSnippet, Snippet } from '../data/snippet.ts'; import type { VariableScope } from '../core/variable/tgpuVariable.ts'; @@ -15,10 +15,20 @@ export interface FunctionDefinitionOptions { readonly workgroupSize?: readonly number[] | undefined; readonly args: readonly FunctionArgument[]; readonly body: Block; + /** + * Default values of parameters that were omitted at the call site + * (`(a, b = 2) => ...` called as `fn(a)` or `fn(a, undefined)`), in parameter order. + */ + readonly paramDefaults?: readonly ParamDefault[] | undefined; determineReturnType(): BaseData; } +export interface ParamDefault { + readonly name: string; + readonly expression: Expression; +} + export interface ConstantDefinitionOptions { readonly id: string; readonly dataType: BaseData; diff --git a/packages/typegpu/src/tgsl/shellless.ts b/packages/typegpu/src/tgsl/shellless.ts index f7f6bb6606..323aff7edf 100644 --- a/packages/typegpu/src/tgsl/shellless.ts +++ b/packages/typegpu/src/tgsl/shellless.ts @@ -1,8 +1,9 @@ +import { FuncParameterType } from 'tinyest'; import { createShelllessImpl, type ShelllessImpl } from '../core/function/shelllessImpl.ts'; import { type AnyData, UnknownData } from '../data/dataTypes.ts'; import { schemasEqualNominally } from '../data/schemasEqual.ts'; import type { Snippet } from '../data/snippet.ts'; -import { isPtr } from '../data/wgslTypes.ts'; +import { isPtr, isVoid, Void } from '../data/wgslTypes.ts'; import { getFunctionMetadata, getName } from '../shared/meta.ts'; import { concretize } from './generationHelpers.ts'; import { stringifySnippet } from './stringifySnippet.ts'; @@ -17,7 +18,13 @@ export class ShelllessRepository { if (!meta) { return undefined; } - if (!argSnippets && meta.ast.params.length > 0) { + const { params } = meta.ast; + const hasDefault = (index: number) => { + const param = params[index]; + return param?.type === FuncParameterType.identifier && param.default !== undefined; + }; + + if (!argSnippets && params.some((_, index) => !hasDefault(index))) { throw new Error( `Cannot resolve '${getName( fn, @@ -26,9 +33,19 @@ export class ShelllessRepository { } const argTypes = (argSnippets ?? []).map((s, index) => { + if (s.value === undefined) { + if (!hasDefault(index)) { + throw new Error( + `Passed undefined as the #${index} argument to ${getName(fn) ?? ''}(...), but the parameter has no default value.`, + ); + } + // Omitted, the parameter's default value is used instead + return Void; + } + if (s.dataType === UnknownData) { const snippetStr = stringifySnippet(s); - const valueStr = String(s.value); + const valueStr = String(s.value as unknown); const description = snippetStr === valueStr ? `'${snippetStr}'` : `'${snippetStr}' (${valueStr})`; throw new Error( @@ -56,6 +73,11 @@ export class ShelllessRepository { return type; }); + // `fn(a, undefined)` is equivalent to `fn(a)` + while (argTypes.length > 0 && isVoid(argTypes.at(-1))) { + argTypes.pop(); + } + let cache = this.cache.get(fn); if (cache) { const variant = cache.find( diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 46bc0a665e..177beedeef 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -48,6 +48,7 @@ import type { ShaderGenerator, ConstantDefinitionOptions, FunctionDefinitionOptions, + ParamDefault, VariableDefinitionOptions, BinaryOperator, ResolvedStatement, @@ -278,6 +279,7 @@ export class WgslGenerator implements ShaderGenerator { [_, statementNodes]: tinyest.Block, allowInlining: boolean, externalMap?: ExternalMap, + paramDefaults?: readonly ParamDefault[], ): ResolvedStatement { this.ctx.pushBlockScope(); @@ -297,6 +299,10 @@ export class WgslGenerator implements ShaderGenerator { let definesInNearestScope = false; let endsWithControlFlow: ResolvedStatement['endsWithControlFlow']; this.ctx.indent(); + for (const statement of this._paramDefaults(paramDefaults ?? [])) { + body += `${statement.code}\n`; + definesInNearestScope ||= statement.definesInNearestScope ?? false; + } for (const statementNode of statementNodes) { const statement = this._statement(statementNode); if (statement.code.length > 0) { @@ -342,6 +348,31 @@ export class WgslGenerator implements ShaderGenerator { } } + /** + * Makes the default values of omitted parameters available in the function body. + * Defaults known at comptime are inlined at every use, the rest are declared at + * the top of the function body. + */ + protected _paramDefaults(paramDefaults: readonly ParamDefault[]): ResolvedStatement[] { + return paramDefaults.flatMap(({ name, expression }) => { + const snippet = this._expression(expression); + + if (snippet.dataType === UnknownData || wgsl.isVoid(snippet.dataType)) { + const exprStr = stringifyNode(expression); + throw new WgslTypeError( + `Default value '${exprStr}' of parameter '${name}' is invalid, cannot determine WGSL type of '${exprStr}'${schemaWrappingSuggestion(name, exprStr, snippet.value)}`, + ); + } + + if (asComptime(snippet) !== undefined) { + this.ctx.defineVariable(name, snippet); + return []; + } + + return [this._constDeclaration(name, expression, snippet)]; + }); + } + protected _blockStatement(block: tinyest.Block, externalMap?: ExternalMap): ResolvedStatement { const { code, ...properties } = this._block(block, /* allowInlining */ true, externalMap); @@ -455,9 +486,13 @@ export class WgslGenerator implements ShaderGenerator { return undefined; } - const converted = args.map((s, idx) => { - const argType = shellless.argTypes[idx] as wgsl.BaseData; - return tryConvertSnippet(this.ctx, s, argType, /* verbose */ false); + const converted = args.flatMap((s, idx) => { + const argType = shellless.argTypes[idx]; + if (argType === undefined || wgsl.isVoid(argType)) { + // Omitted, the parameter's default value is used instead + return []; + } + return [tryConvertSnippet(this.ctx, s, argType, /* verbose */ false)]; }); return this.ctx.withResetIndentLevel(() => { @@ -1263,7 +1298,12 @@ The generated shader will evaluate them in declaration order.`, this.ctx.blockDepth === functionInitialBlockDepth - 1, `Expecting exactly ${functionInitialBlockDepth - 1} block(s) before going into the first function block scope`, ); - let body = this._block(options.body, /* allowInlining */ false); + let body = this._block( + options.body, + /* allowInlining */ false, + /* externalMap */ undefined, + options.paramDefaults, + ); const scope = this.ctx.topFunctionScope; invariant(scope, 'Expected function scope to be present'); const replacements = Object.fromEntries( @@ -1556,8 +1596,17 @@ Try 'return ${typeStr}(${str});' instead. ); } - const eq = this._expression(eqNode); + return this._constDeclaration(rawId, eqNode, this._expression(eqNode)); + } + /** + * Declares `const = ;`, where `eq` is the already generated `eqNode`. + */ + protected _constDeclaration( + rawId: string, + eqNode: tinyest.Expression, + eq: Snippet, + ): ResolvedStatement { if (eq.value instanceof RefOperator) { // We're assigning a newly created `d.ref()` if (eq.value.addressable) { diff --git a/packages/typegpu/src/types.ts b/packages/typegpu/src/types.ts index 3c66e1031f..4ba734a0fd 100644 --- a/packages/typegpu/src/types.ts +++ b/packages/typegpu/src/types.ts @@ -84,6 +84,10 @@ export interface ResolveFunctionOptions { functionType: 'normal' | ShaderStage; workgroupSize?: readonly number[] | undefined; name: string; + /** + * `Void` marks an argument omitted at the call site (passed as `undefined`), in which + * case the parameter's default value is used. Same goes for parameters past the end. + */ argTypes: BaseData[]; /** * The return type of the function. If undefined, the type should be inferred diff --git a/packages/typegpu/tests/tgsl/shellless.test.ts b/packages/typegpu/tests/tgsl/shellless.test.ts index 84098f3ba3..d4655068c8 100644 --- a/packages/typegpu/tests/tgsl/shellless.test.ts +++ b/packages/typegpu/tests/tgsl/shellless.test.ts @@ -512,7 +512,7 @@ describe('shellless', () => { `); }); - it('falls back to default parameter values when arguments are omitted', () => { + it('inlines comptime-known default values of omitted parameters', () => { const scale = (v: d.v2f, factor: number = 2) => { 'use gpu'; return std.mul(v, d.f32(factor)); @@ -525,8 +525,7 @@ describe('shellless', () => { expect(tgpu.resolve([foo])).toMatchInlineSnapshot(` "fn scale(v: vec2f) -> vec2f { - const factor = 2; - return (v * f32(factor)); + return (v * 2f); } fn scale_1(v: vec2f, factor: f32) -> vec2f { @@ -539,7 +538,7 @@ describe('shellless', () => { `); }); - it('supports comptime-style mode flags with default values', () => { + it('prunes branches based on comptime-known default values', () => { const interpolate = (t: number, cyclic: boolean = false) => { 'use gpu'; if (cyclic) { @@ -555,10 +554,6 @@ describe('shellless', () => { expect(tgpu.resolve([foo])).toMatchInlineSnapshot(` "fn interpolate(t: f32) -> f32 { - const cyclic = false; - if (cyclic) { - return fract(t); - } return clamp(t, 0f, 1f); } @@ -575,15 +570,81 @@ describe('shellless', () => { `); }); - it('evaluates default values referencing earlier parameters', () => { + it('treats undefined arguments as omitted', () => { + const interpolate = (t: number, cyclic: boolean = false) => { + 'use gpu'; + if (cyclic) { + return std.fract(t); + } + return std.clamp(t, 0, 1); + }; + + const offset = (base: number = 10, x: number) => { + 'use gpu'; + return base + x; + }; + + const foo = () => { + 'use gpu'; + return interpolate(1.5) + interpolate(2.5, undefined) + offset(undefined, 2); + }; + + expect(tgpu.resolve([foo])).toMatchInlineSnapshot(` + "fn interpolate(t: f32) -> f32 { + return clamp(t, 0f, 1f); + } + + fn offset(x: i32) -> i32 { + return (10i + x); + } + + fn foo() -> f32 { + return ((interpolate(1.5f) + interpolate(2.5f)) + f32(offset(2i))); + }" + `); + }); + + it('inlines comptime-known default values referencing other defaults', () => { + const foo = (a: number = 2, b: number = a * 3) => { + 'use gpu'; + return d.f32(a + b); + }; + + const main = () => { + 'use gpu'; + return foo() + foo(1); + }; + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "fn foo() -> f32 { + return 8f; + } + + fn foo_1(a: i32) -> f32 { + let b = (a * 3i); + return f32((a + b)); + } + + fn main() -> f32 { + return (foo() + foo_1(1i)); + }" + `); + }); + + it('declares runtime default values in the function body', () => { const madd = (a: d.v2f, b: d.v2f = a) => { 'use gpu'; return std.add(a, b); }; + const normalize = (v: d.v2f, len: number = std.length(v)) => { + 'use gpu'; + return std.div(v, len); + }; + const foo = () => { 'use gpu'; - return madd(d.vec2f(1, 2)); + return std.add(madd(d.vec2f(1, 2)), normalize(d.vec2f(3, 4))); }; expect(tgpu.resolve([foo])).toMatchInlineSnapshot(` @@ -592,8 +653,13 @@ describe('shellless', () => { return (a + b); } + fn normalize_1(v: vec2f) -> vec2f { + let len = length(v); + return (v / len); + } + fn foo() -> vec2f { - return madd(vec2f(1, 2)); + return (madd(vec2f(1, 2)) + normalize_1(vec2f(3, 4))); }" `); }); @@ -611,8 +677,7 @@ describe('shellless', () => { expect(tgpu.resolve([main])).toMatchInlineSnapshot(` "fn addBias(x: f32) -> f32 { - const bias = 10; - return (x + f32(bias)); + return (x + 10f); } fn main() -> f32 { @@ -620,4 +685,61 @@ describe('shellless', () => { }" `); }); + + it('can be resolved directly if every parameter has a default value', () => { + const foo = (a: number = 1, b: boolean = true) => { + 'use gpu'; + return b ? d.f32(a) : 0; + }; + + expect(tgpu.resolve([foo])).toMatchInlineSnapshot(` + "fn foo() -> f32 { + return 1f; + }" + `); + }); + + it('throws when passing undefined to a parameter without a default value', () => { + const foo = (a: number | undefined) => { + 'use gpu'; + return a; + }; + + const main = () => { + 'use gpu'; + return foo(undefined); + }; + + expect(() => tgpu.resolve([main])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn*:main + - fn*:main(): Passed undefined as the #0 argument to foo(...), but the parameter has no default value.] + `); + }); + + it('throws when a default value has no WGSL representation', () => { + const config = { scale: 2 }; + + const foo = (x: number, c = config) => { + 'use gpu'; + return x * c.scale; + }; + + const main = () => { + 'use gpu'; + return foo(1); + }; + + expect(() => tgpu.resolve([main])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn*:main + - fn*:main() + - fn*:foo(i32): Default value 'config' of parameter 'c' is invalid, cannot determine WGSL type of 'config' + ----- + - Try using or defining a schema that matches your desired value the most, and wrap the value with it: 'c = Schema(config)' + -----] + `); + }); });