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,10 +117,16 @@ 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};`,
+      ctx.gen.globalVarDefinition({
+        scope: this.usage,
+        name: id,
+        dataType,
+        group: group,
+        binding: binding,
+        init: undefined,
+      }),
     );
 
     return snip(id, dataType, this.usage);
@@ -243,12 +243,16 @@ 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
-      };`,
+      ctx.gen.globalVarDefinition({
+        scope: this.usage,
+        name: id,
+        dataType: this.dataType,
+        group: group,
+        binding: this.#membership.idx,
+        init: undefined,
+      }),
     );
 
     return snip(id, this.dataType, this.usage);

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

@@ -109,10 +109,14 @@ 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};`);
+    ctx.addDeclaration(
+      ctx.gen.globalConstDefinition(
+        id,
+        this.dataType,
+        snip(this.#value, this.dataType, 'constant'),
+      ),
+    );
 
     return snip(id, this.dataType, 'constant-immutable-def');
   }

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

@@ -88,13 +88,13 @@ 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};`);
-    }
+    ctx.addDeclaration(
+      ctx.gen.globalVarDefinition({ scope: this.#scope, name: id, dataType: this.#dataType, init }),
+    );
 
     return snip(id, this.#dataType, this.#scope);
   }

packages/typegpu/src/tgsl/shaderGenerator.ts

@@ -1,7 +1,10 @@
 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 {
+  FunctionDefinitionOptions,
+  VariableDefinitionOptions,
+} from './shaderGenerator_members.ts';
 
 /**
  * **NOTE: This is an unstable API and may change in the future.**
@@ -12,7 +15,10 @@ import type { FunctionDefinitionOptions } from './shaderGenerator_members.ts';
 export interface ShaderGenerator {
   initGenerator(ctx: GenerationCtx): void;
 
+  globalConstDefinition(id: string, schema: BaseData, init: Snippet): string;
+  globalVarDefinition(options: VariableDefinitionOptions): string;
   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,12 @@ export interface FunctionDefinitionOptions {
 
   determineReturnType(): BaseData;
 }
+
+export interface VariableDefinitionOptions {
+  readonly scope: VariableScope | BindableBufferUsage;
+  readonly name: 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,14 @@ 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 {
+  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 +202,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 +861,30 @@ ${this.ctx.pre}}`;
     assertExhaustive(expression);
   }
 
+  public globalConstDefinition(id: string, schema: wgsl.BaseData, init: Snippet): string {
+    const resolvedDataType = this.ctx.resolve(schema).value;
+    const resolvedValue = this.ctx.resolveSnippet(init).value;
+
+    return `const ${id}: ${resolvedDataType} = ${resolvedValue};`;
+  }
+
+  public globalVarDefinition(options: VariableDefinitionOptions): string {
+    let pre = `var<${usageToVarTemplateMap[options.scope]}> ${options.name}: ${this.ctx.resolve(options.dataType).value}`;
+
+    if (options.binding !== undefined) {
+      pre = `@binding(${options.binding}) ` + pre;
+    }
+
+    if (options.group !== undefined) {
+      pre = `@group(${options.group}) ` + pre;
+    }
+
+    if (options.init) {
+      return `${pre} = ${this.ctx.resolveSnippet(options.init).value};`;
+    }
+    return `${pre};`;
+  }
+
   public functionDefinition(options: FunctionDefinitionOptions): string {
     // Function body
     let body = this._block(options.body);