feat: Route more definitions through the shader generator

Author: iwoplaza

packages/typegpu/src/core/buffer/bufferUsage.ts

@@ -86,12 +86,6 @@ export function isUsableAsUniform<T extends TgpuBuffer<BaseData>>(
 // Implementation
 // --------------
 
-const usageToVarTemplateMap: Record<BindableBufferUsage, string> = {
-  uniform: 'uniform',
-  mutable: 'storage, read_write',
-  readonly: 'storage, read',
-};
-
 class TgpuFixedBufferImpl<TData extends BaseData, TUsage extends BindableBufferUsage>
   implements TgpuBufferUsage<TData, TUsage>, SelfResolvable, TgpuFixedBufferUsage<TData>
 {
@@ -123,13 +117,15 @@ class TgpuFixedBufferImpl<TData extends BaseData, TUsage extends BindableBufferU
       this.usage === 'uniform' ? { uniform: dataType } : { storage: dataType, access: this.usage },
       this.buffer,
     );
-    const usage = usageToVarTemplateMap[this.usage];
-
-    ctx.addDeclaration(
-      `@group(${group}) @binding(${binding}) var<${usage}> ${id}: ${ctx.resolve(dataType).value};`,
-    );
 
-    return snip(id, dataType, this.usage);
+    return ctx.gen.declareGlobalVar({
+      group,
+      binding,
+      scope: this.usage,
+      id,
+      dataType,
+      init: undefined,
+    });
   }
 
   toString(): string {
@@ -243,15 +239,15 @@ export class TgpuLaidOutBufferImpl<TData extends BaseData, TUsage extends Bindab
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
     const id = ctx.makeUniqueIdentifier(getName(this), 'global');
     const group = ctx.allocateLayoutEntry(this.#membership.layout);
-    const usage = usageToVarTemplateMap[this.usage];
-
-    ctx.addDeclaration(
-      `@group(${group}) @binding(${this.#membership.idx}) var<${usage}> ${id}: ${
-        ctx.resolve(this.dataType).value
-      };`,
-    );
 
-    return snip(id, this.dataType, this.usage);
+    return ctx.gen.declareGlobalVar({
+      group,
+      binding: this.#membership.idx,
+      scope: this.usage,
+      id,
+      dataType: this.dataType,
+      init: undefined,
+    });
   }
 
   toString(): string {

packages/typegpu/src/core/constant/tgpuConstant.ts

@@ -109,12 +109,12 @@ class TgpuConstImpl<TDataType extends BaseData> implements TgpuConst<TDataType>,
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
     const id = ctx.makeUniqueIdentifier(getName(this), 'global');
-    const resolvedDataType = ctx.resolve(this.dataType).value;
-    const resolvedValue = ctx.resolve(this.#value, this.dataType).value;
 
-    ctx.addDeclaration(`const ${id}: ${resolvedDataType} = ${resolvedValue};`);
-
-    return snip(id, this.dataType, 'constant-immutable-def');
+    return ctx.gen.declareGlobalConst({
+      id,
+      dataType: this.dataType,
+      init: snip(this.#value, this.dataType, 'constant'),
+    });
   }
 
   toString() {

packages/typegpu/src/core/texture/texture.ts

@@ -613,11 +613,14 @@ class TgpuFixedTextureViewImpl<T extends WgslTexture | WgslStorageTexture>
       this,
     );
 
-    ctx.addDeclaration(
-      `@group(${group}) @binding(${binding}) var ${id}: ${ctx.resolve(this.schema).value};`,
-    );
-
-    return snip(id, this.schema, /* origin */ 'handle');
+    return ctx.gen.declareGlobalVar({
+      group,
+      binding,
+      id,
+      dataType: this.schema,
+      scope: 'handle',
+      init: undefined,
+    });
   }
 }
 

packages/typegpu/src/core/variable/tgpuVariable.ts

@@ -88,15 +88,11 @@ class TgpuVarImpl<TScope extends VariableScope, TDataType extends BaseData>
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
     const id = ctx.makeUniqueIdentifier(getName(this), 'global');
-    const pre = `var<${this.#scope}> ${id}: ${ctx.resolve(this.#dataType).value}`;
+    const init = this.#initialValue
+      ? snip(this.#initialValue, this.#dataType, 'constant')
+      : undefined;
 
-    if (this.#initialValue) {
-      ctx.addDeclaration(`${pre} = ${ctx.resolve(this.#initialValue, this.#dataType).value};`);
-    } else {
-      ctx.addDeclaration(`${pre};`);
-    }
-
-    return snip(id, this.#dataType, this.#scope);
+    return ctx.gen.declareGlobalVar({ scope: this.#scope, id, dataType: this.#dataType, init });
   }
 
   $name(label: string) {

packages/typegpu/src/tgsl/shaderGenerator.ts

@@ -1,7 +1,11 @@
 import type { BaseData } from '../data/wgslTypes.ts';
 import type { GenerationCtx } from './generationHelpers.ts';
 import type { ResolvedSnippet, Snippet } from '../data/snippet.ts';
-import type { FunctionDefinitionOptions } from './shaderGenerator_members.ts';
+import type {
+  ConstantDefinitionOptions,
+  FunctionDefinitionOptions,
+  VariableDefinitionOptions,
+} from './shaderGenerator_members.ts';
 
 /**
  * **NOTE: This is an unstable API and may change in the future.**
@@ -12,7 +16,10 @@ import type { FunctionDefinitionOptions } from './shaderGenerator_members.ts';
 export interface ShaderGenerator {
   initGenerator(ctx: GenerationCtx): void;
 
+  declareGlobalConst(options: ConstantDefinitionOptions): ResolvedSnippet;
+  declareGlobalVar(options: VariableDefinitionOptions): ResolvedSnippet;
   functionDefinition(options: FunctionDefinitionOptions): string;
+
   typeInstantiation(schema: BaseData, args: readonly Snippet[]): ResolvedSnippet;
   typeAnnotation(schema: BaseData): string;
 }

packages/typegpu/src/tgsl/shaderGenerator_members.ts

@@ -1,6 +1,8 @@
 import type { Block } from 'tinyest';
 import type { BaseData } from '../data/wgslTypes.ts';
-import type { FunctionArgument, TgpuShaderStage } from '../types.ts';
+import type { BindableBufferUsage, FunctionArgument, TgpuShaderStage } from '../types.ts';
+import type { VariableScope } from '../core/variable/tgpuVariable.ts';
+import type { Snippet } from '../data/snippet.ts';
 
 export { UnknownData } from '../data/dataTypes.ts';
 export { getName } from '../shared/meta.ts';
@@ -19,3 +21,18 @@ export interface FunctionDefinitionOptions {
 
   determineReturnType(): BaseData;
 }
+
+export interface ConstantDefinitionOptions {
+  readonly id: string;
+  readonly dataType: BaseData;
+  readonly init: Snippet;
+}
+
+export interface VariableDefinitionOptions {
+  readonly scope: VariableScope | BindableBufferUsage | 'handle';
+  readonly id: string;
+  readonly dataType: BaseData;
+  readonly init: Snippet | undefined;
+  readonly group?: string | undefined;
+  readonly binding?: number | undefined;
+}

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -19,7 +19,12 @@ import { $gpuCallable, $internal, $providing, isMarkedInternal } from '../shared
 import { safeStringify } from '../shared/stringify.ts';
 import { pow } from '../std/numeric.ts';
 import { add, div, mul, neg, sub } from '../std/operators.ts';
-import { isGPUCallable, isKnownAtComptime, type DualFn } from '../types.ts';
+import {
+  isGPUCallable,
+  isKnownAtComptime,
+  type BindableBufferUsage,
+  type DualFn,
+} from '../types.ts';
 import { convertStructValues, convertToCommonType, tryConvertSnippet } from './conversion.ts';
 import {
   ArrayExpression,
@@ -44,10 +49,15 @@ import type { ExternalMap } from '../core/resolve/externals.ts';
 import * as forOfUtils from './forOfUtils.ts';
 import { isTgpuRange } from '../std/range.ts';
 import { stringifyNode } from '../shared/tseynit.ts';
-import type { FunctionDefinitionOptions } from './shaderGenerator_members.ts';
+import type {
+  ConstantDefinitionOptions,
+  FunctionDefinitionOptions,
+  VariableDefinitionOptions,
+} from './shaderGenerator_members.ts';
 import { getAttributesString } from '../data/attributes.ts';
 import { validSelectBranchTypes } from '../std/boolean.ts';
 import { isInfixDispatch } from './infixDispatch.ts';
+import type { VariableScope } from '../core/variable/tgpuVariable.ts';
 
 const { NodeTypeCatalog: NODE } = tinyest;
 
@@ -193,6 +203,14 @@ const binaryOpCodeToCodegen = {
   '**': pow[$gpuCallable].call.bind(pow),
 } satisfies Partial<Record<tinyest.BinaryOperator, (...args: never[]) => unknown>>;
 
+const usageToVarTemplateMap: Record<VariableScope | BindableBufferUsage, string> = {
+  private: 'private',
+  workgroup: 'workgroup',
+  uniform: 'uniform',
+  mutable: 'storage, read_write',
+  readonly: 'storage, read',
+};
+
 export class WgslGenerator implements ShaderGenerator {
   #ctx: GenerationCtx | undefined = undefined;
   // used to detect `continue` and `break` nodes in loop body
@@ -844,6 +862,41 @@ ${this.ctx.pre}}`;
     assertExhaustive(expression);
   }
 
+  public declareGlobalConst(options: ConstantDefinitionOptions): ResolvedSnippet {
+    const resolvedDataType = this.ctx.resolve(options.dataType).value;
+    const resolvedValue = this.ctx.resolveSnippet(options.init).value;
+
+    this.ctx.addDeclaration(`const ${options.id}: ${resolvedDataType} = ${resolvedValue};`);
+
+    return snip(options.id, options.dataType, 'constant-immutable-def');
+  }
+
+  public declareGlobalVar(options: VariableDefinitionOptions): ResolvedSnippet {
+    let pre = '';
+
+    if (options.group !== undefined) {
+      pre += `@group(${options.group}) `;
+    }
+
+    if (options.binding !== undefined) {
+      pre += `@binding(${options.binding}) `;
+    }
+
+    if (options.scope in usageToVarTemplateMap) {
+      pre += `var<${usageToVarTemplateMap[options.scope as keyof typeof usageToVarTemplateMap]}> `;
+    } else {
+      pre += `var `;
+    }
+
+    pre += `${options.id}: ${this.ctx.resolve(options.dataType).value}`;
+
+    this.ctx.addDeclaration(
+      options.init ? `${pre} = ${this.ctx.resolveSnippet(options.init).value};` : `${pre};`,
+    );
+
+    return snip(options.id, options.dataType, options.scope);
+  }
+
   public functionDefinition(options: FunctionDefinitionOptions): string {
     // Function body
     let body = this._block(options.body);