feat: Declare default parameters like local variables, support optional parameters

Adapts default parameter values to the approach described in #3277:
- Default values are declared at the top of the function body instead of being inlined when comptime-known.
- Parameters omitted without a default (optional parameters) are comptime-known `undefined`.
- Comparing a runtime value against `undefined` is folded at comptime, allowing branches to be pruned.

Author: iwoplaza

packages/typegpu-gl/src/glslGenerator.ts

@@ -1288,7 +1288,7 @@ The generated shader will omit it, so its runtime side effects will not occur.`)
         options.body,
         /* allowInlining */ false,
         /* externalMap */ undefined,
-        options.paramDefaults,
+        options.omittedParams,
       ).code;
       const returnType = options.determineReturnType();
 

packages/typegpu-gl/tests/glslGenerator.test.ts

@@ -95,8 +95,9 @@ describe('GlslGenerator - variable declarations', () => {
 
     expect(tgpu.resolve([foo], glOptions())).toMatchInlineSnapshot(`
       "vec2 shift(vec2 v) {
+        int amount = 2;
         vec2 by = v;
-        return (v + (by * 2.0));
+        return (v + (by * float(amount)));
       }
 
       vec2 foo() {

packages/typegpu/src/resolutionCtx.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 { ParamDefault, ShaderGenerator } from './tgsl/shaderGenerator.ts';
+import type { OmittedParam, ShaderGenerator } from './tgsl/shaderGenerator.ts';
 import { WgslGenerator } from './tgsl/wgslGenerator.ts';
 import type {
   BlockScopeLayer,
@@ -738,18 +738,15 @@ export class ResolutionCtxImpl implements ResolutionCtx {
         }
       }
 
-      // Parameters omitted at the call site fall back to their default values (`(a, b = 2) => ...`)
-      const paramDefaults: ParamDefault[] = [];
+      // Parameters omitted at the call site fall back to their default values (`(a, b = 2) => ...`),
+      // or become comptime-known `undefined` if they don't have one (`(a, b?: number) => ...`)
+      const omittedParams: OmittedParam[] = [];
       if (!options.entryInput) {
         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 (omitted && param.type === FuncParameterType.identifier) {
+            omittedParams.push({ name: param.name, default: param.default });
           }
         }
       }
@@ -762,7 +759,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
         workgroupSize: options.workgroupSize,
         args,
         body: options.body,
-        paramDefaults,
+        omittedParams,
         determineReturnType: () => {
           if (returnType) {
             // Already determined

packages/typegpu/src/tgsl/shaderGenerator.ts

@@ -16,17 +16,20 @@ export interface FunctionDefinitionOptions {
   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.
+   * Parameters that were omitted at the call site (`fn(a)` or `fn(a, undefined)`), in parameter order.
    */
-  readonly paramDefaults?: readonly ParamDefault[] | undefined;
+  readonly omittedParams?: readonly OmittedParam[] | undefined;
 
   determineReturnType(): BaseData;
 }
 
-export interface ParamDefault {
+export interface OmittedParam {
   readonly name: string;
-  readonly expression: Expression;
+  /**
+   * The parameter's default value (`(a, b = 2) => ...`), declared at the top of the function body.
+   * Optional parameters without a default (`(a, b?: number) => ...`) are comptime-known `undefined`.
+   */
+  readonly default?: Expression | undefined;
 }
 
 export interface ConstantDefinitionOptions {

packages/typegpu/src/tgsl/shellless.ts

@@ -34,12 +34,12 @@ export class ShelllessRepository {
 
     const argTypes = (argSnippets ?? []).map((s, index) => {
       if (s.value === undefined) {
-        if (!hasDefault(index)) {
+        if (params[index]?.type !== FuncParameterType.identifier) {
           throw new Error(
-            `Passed undefined as the #${index} argument to ${getName(fn) ?? '<unnamed>'}(...), but the parameter has no default value.`,
+            `Passed undefined as the #${index} argument to ${getName(fn) ?? '<unnamed>'}(...), but only identifier parameters can be omitted.`,
           );
         }
-        // Omitted, the parameter's default value is used instead
+        // Omitted, the parameter's default value (or `undefined`) is used instead
         return Void;
       }
 

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -49,7 +49,7 @@ import type {
   ShaderGenerator,
   ConstantDefinitionOptions,
   FunctionDefinitionOptions,
-  ParamDefault,
+  OmittedParam,
   VariableDefinitionOptions,
   BinaryOperator,
   ResolvedStatement,
@@ -280,7 +280,7 @@ export class WgslGenerator implements ShaderGenerator {
     [_, statementNodes]: tinyest.Block,
     allowInlining: boolean,
     externalMap?: ExternalMap,
-    paramDefaults?: readonly ParamDefault[],
+    omittedParams?: readonly OmittedParam[],
   ): ResolvedStatement {
     this.ctx.pushBlockScope();
 
@@ -300,8 +300,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`;
+      for (const statement of this._omittedParams(omittedParams ?? [])) {
+        if (statement.code.length > 0) {
+          body += `${statement.code}\n`;
+        }
         definesInNearestScope ||= statement.definesInNearestScope ?? false;
       }
       for (const statementNode of statementNodes) {
@@ -350,12 +352,17 @@ 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.
+   * Makes parameters omitted at the call site available in the function body.
+   * Default values are declared at the top of the function body, just like local
+   * variables would be, while parameters without a default are comptime-known `undefined`.
    */
-  protected _paramDefaults(paramDefaults: readonly ParamDefault[]): ResolvedStatement[] {
-    return paramDefaults.flatMap(({ name, expression }) => {
+  protected _omittedParams(omittedParams: readonly OmittedParam[]): ResolvedStatement[] {
+    return omittedParams.map(({ name, default: expression }) => {
+      if (expression === undefined) {
+        this.ctx.defineVariable(name, snip(undefined, wgsl.Void, 'constant', false));
+        return { code: '', definesInNearestScope: false };
+      }
+
       const snippet = this._expression(expression);
 
       if (snippet.dataType === UnknownData || wgsl.isVoid(snippet.dataType)) {
@@ -365,12 +372,7 @@ export class WgslGenerator implements ShaderGenerator {
         );
       }
 
-      if (asComptime(snippet) !== undefined) {
-        this.ctx.defineVariable(name, snippet);
-        return [];
-      }
-
-      return [this._constDeclaration(name, expression, snippet)];
+      return this._constDeclaration(name, expression, snippet);
     });
   }
 
@@ -636,6 +638,20 @@ export class WgslGenerator implements ShaderGenerator {
         throw new Error('Please use the !== operator instead of !=');
       }
 
+      if (op === '===' || op === '!==') {
+        // Comparing against `undefined` (e.g. an omitted optional parameter) is known at comptime,
+        // since values with a WGSL type are never `undefined`.
+        const isUndefined = (s: Snippet) => s.value === undefined && wgsl.isVoid(s.dataType);
+        const isRuntime = (s: Snippet) =>
+          asComptime(s) === undefined && !wgsl.isVoid(s.dataType) && !s.possibleSideEffects;
+        if (
+          (isUndefined(lhsExpr) && isRuntime(rhsExpr)) ||
+          (isRuntime(lhsExpr) && isUndefined(rhsExpr))
+        ) {
+          return snip(op === '!==', bool, 'constant', false);
+        }
+      }
+
       const stdBinaryRelationalOp = binaryRelationalOpToStdMap[op];
       const knownLhs = stdBinaryRelationalOp && asComptime(lhsExpr);
       const knownRhs = knownLhs && asComptime(rhsExpr);
@@ -1303,7 +1319,7 @@ The generated shader will evaluate them in declaration order.`,
       options.body,
       /* allowInlining */ false,
       /* externalMap */ undefined,
-      options.paramDefaults,
+      options.omittedParams,
     );
     const scope = this.ctx.topFunctionScope;
     invariant(scope, 'Expected function scope to be present');

packages/typegpu/tests/tgsl/shellless.test.ts

@@ -512,7 +512,7 @@ describe('shellless', () => {
     `);
   });
 
-  it('inlines comptime-known default values of omitted parameters', () => {
+  it('declares default values of omitted parameters at the top of the function body', () => {
     const scale = (v: d.v2f, factor: number = 2) => {
       'use gpu';
       return std.mul(v, d.f32(factor));
@@ -525,7 +525,8 @@ describe('shellless', () => {
 
     expect(tgpu.resolve([foo])).toMatchInlineSnapshot(`
       "fn scale(v: vec2f) -> vec2f {
-        return (v * 2f);
+        const factor = 2;
+        return (v * f32(factor));
       }
 
       fn scale_1(v: vec2f, factor: f32) -> vec2f {
@@ -538,7 +539,7 @@ describe('shellless', () => {
     `);
   });
 
-  it('prunes branches based on comptime-known default values', () => {
+  it('does not treat default values as comptime-known', () => {
     const interpolate = (t: number, cyclic: boolean = false) => {
       'use gpu';
       if (cyclic) {
@@ -554,6 +555,10 @@ 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);
       }
 
@@ -591,11 +596,16 @@ 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);
       }
 
       fn offset(x: i32) -> i32 {
-        return (10i + x);
+        const base = 10;
+        return (base + x);
       }
 
       fn foo() -> f32 {
@@ -604,7 +614,7 @@ describe('shellless', () => {
     `);
   });
 
-  it('inlines comptime-known default values referencing other defaults', () => {
+  it('supports default values referencing other parameters', () => {
     const foo = (a: number = 2, b: number = a * 3) => {
       'use gpu';
       return d.f32(a + b);
@@ -617,7 +627,9 @@ describe('shellless', () => {
 
     expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
       "fn foo() -> f32 {
-        return 8f;
+        const a = 2;
+        let b = (a * 3i);
+        return f32((a + b));
       }
 
       fn foo_1(a: i32) -> f32 {
@@ -677,7 +689,8 @@ describe('shellless', () => {
 
     expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
       "fn addBias(x: f32) -> f32 {
-        return (x + 10f);
+        const bias = 10;
+        return (x + f32(bias));
       }
 
       fn main() -> f32 {
@@ -694,27 +707,87 @@ describe('shellless', () => {
 
     expect(tgpu.resolve([foo])).toMatchInlineSnapshot(`
       "fn foo() -> f32 {
-        return 1f;
+        const a = 1;
+        const b = true;
+        return select(0f, f32(a), b);
+      }"
+    `);
+  });
+
+  it('treats omitted optional parameters as comptime-known undefined', () => {
+    const interpolate = (t: number, period?: number) => {
+      'use gpu';
+      if (period !== undefined) {
+        return std.fract(t / period);
+      }
+      return std.saturate(t);
+    };
+
+    const main = () => {
+      'use gpu';
+      return interpolate(0.5) + interpolate(1.5, 2) + interpolate(2.5, undefined);
+    };
+
+    expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
+      "fn interpolate(t: f32) -> f32 {
+        return saturate(t);
+      }
+
+      fn interpolate_1(t: f32, period: i32) -> f32 {
+        return fract((t / f32(period)));
+      }
+
+      fn main() -> f32 {
+        return ((interpolate(0.5f) + interpolate_1(1.5f, 2i)) + interpolate(2.5f));
+      }"
+    `);
+  });
+
+  it('folds comparisons of runtime values against undefined', () => {
+    const foo = (a: number, b?: number) => {
+      'use gpu';
+      if (a === undefined || b === undefined) {
+        return a;
+      }
+      return a + b;
+    };
+
+    const main = () => {
+      'use gpu';
+      return foo(1) + foo(1, 2);
+    };
+
+    expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
+      "fn foo(a: i32) -> i32 {
+        return a;
+      }
+
+      fn foo_1(a: i32, b: i32) -> i32 {
+        return (a + b);
+      }
+
+      fn main() -> i32 {
+        return (foo(1i) + foo_1(1i, 2i));
       }"
     `);
   });
 
-  it('throws when passing undefined to a parameter without a default value', () => {
-    const foo = (a: number | undefined) => {
+  it('throws when passing undefined to a destructured parameter', () => {
+    const foo = ({ a }: { a: number }) => {
       'use gpu';
       return a;
     };
 
     const main = () => {
       'use gpu';
-      return foo(undefined);
+      return foo(undefined as unknown as { a: number });
     };
 
     expect(() => tgpu.resolve([main])).toThrowErrorMatchingInlineSnapshot(`
       [Error: Resolution of the following tree failed:
       - <root>
       - fn*:main
-      - fn*:main(): Passed undefined as the #0 argument to foo(...), but the parameter has no default value.]
+      - fn*:main(): Passed undefined as the #0 argument to foo(...), but only identifier parameters can be omitted.]
     `);
   });