fix: Giving local declarations unique names if they clash with global declarations, and vice-versa + 'strict' the default naming scheme in pipelines (#2000)

Author: iwoplaza

packages/typegpu/src/core/function/fnCore.ts

@@ -1,18 +1,7 @@
-import { FuncParameterType } from 'tinyest';
 import { getAttributesString } from '../../data/attributes.ts';
 import { type AnyData, undecorate } from '../../data/dataTypes.ts';
-import {
-  type ResolvedSnippet,
-  snip,
-  type Snippet,
-} from '../../data/snippet.ts';
-import {
-  isPtr,
-  isWgslData,
-  isWgslStruct,
-  Void,
-  type WgslStruct,
-} from '../../data/wgslTypes.ts';
+import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
+import { isWgslData, isWgslStruct, Void } from '../../data/wgslTypes.ts';
 import { MissingLinksError } from '../../errors.ts';
 import { getMetaData, getName, setName } from '../../shared/meta.ts';
 import type { ResolutionCtx } from '../../types.ts';
@@ -190,54 +179,6 @@ export function createFnCore(
 
       // generate wgsl string
 
-      const args: Snippet[] = [];
-      const argAliases: [string, Snippet][] = [];
-
-      for (const [i, argType] of argTypes.entries()) {
-        const astParam = ast.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.
-        //
-        // It still applies for shell-less functions, since we determine the type
-        // of the argument based on the argument's referentiality.
-        // In other words, if we pass a reference to a function, it's typed as a pointer,
-        // otherwise it's typed as a value.
-        const origin = isPtr(argType)
-          ? argType.addressSpace === 'storage'
-            ? argType.access === 'read' ? 'readonly' : 'mutable'
-            : argType.addressSpace
-          : 'argument';
-
-        switch (astParam?.type) {
-          case FuncParameterType.identifier: {
-            const rawName = astParam.name;
-            const snippet = snip(ctx.makeNameValid(rawName), argType, origin);
-            args.push(snippet);
-            if (snippet.value !== rawName) {
-              argAliases.push([rawName, snippet]);
-            }
-            break;
-          }
-          case FuncParameterType.destructuredObject: {
-            args.push(snip(`_arg_${i}`, argType, origin));
-            argAliases.push(...astParam.props.map(({ name, alias }) => {
-              // Undecorating, as the struct type can contain builtins
-              const destrType = undecorate(
-                (argTypes[i] as WgslStruct).propTypes[name],
-              );
-
-              return [
-                alias,
-                snip(`_arg_${i}.${name}`, destrType, 'argument'),
-              ] as [string, Snippet];
-            }));
-            break;
-          }
-          case undefined:
-            args.push(snip(`_arg_${i}`, argType, origin));
-        }
-      }
-
       const { head, body, returnType: actualReturnType } = ctx.fnToWgsl({
         functionType: fnAttribute.includes('@compute')
           ? 'compute'
@@ -246,8 +187,8 @@ export function createFnCore(
           : fnAttribute.includes('@fragment')
           ? 'fragment'
           : 'normal',
-        args,
-        argAliases: Object.fromEntries(argAliases),
+        argTypes,
+        params: ast.params,
         returnType,
         body: ast.body,
         externalMap,

packages/typegpu/src/core/resolve/namespace.ts

@@ -87,7 +87,7 @@ export function getUniqueName(
   namespace: NamespaceInternal,
   resource: object,
 ): string {
-  const name = namespace.nameRegistry.makeUnique(getName(resource));
+  const name = namespace.nameRegistry.makeUnique(getName(resource), true);
   for (const listener of namespace.listeners.name) {
     listener({ target: resource, name });
   }

packages/typegpu/src/core/root/init.ts

@@ -848,7 +848,7 @@ export async function init(options?: InitOptions): Promise<TgpuRoot> {
   const {
     adapter: adapterOpt,
     device: deviceOpt,
-    unstable_names: names = 'random',
+    unstable_names: names = 'strict',
     unstable_logOptions,
   } = options ?? {};
 
@@ -907,7 +907,7 @@ export async function init(options?: InitOptions): Promise<TgpuRoot> {
 export function initFromDevice(options: InitFromDeviceOptions): TgpuRoot {
   const {
     device,
-    unstable_names: names = 'random',
+    unstable_names: names = 'strict',
     unstable_logOptions,
   } = options ?? {};
 

packages/typegpu/src/nameRegistry.ts

@@ -173,6 +173,8 @@ const bannedTokens = new Set([
   'with',
   'writeonly',
   'yield',
+  // Keywords that should be reserved
+  'sampler',
 ]);
 
 export interface NameRegistry {
@@ -181,8 +183,9 @@ export interface NameRegistry {
    * in the lifetime of a single resolution process.
    * Should append "_" to primer, followed by some id.
    * @param primer Used in the generation process, makes the identifier more recognizable.
+   * @param global Whether the name should be registered in the global scope (true), or in the current function scope (false)
    */
-  makeUnique(primer?: string): string;
+  makeUnique(primer: string | undefined, global: boolean): string;
 
   /**
    * Creates a valid WGSL identifier.
@@ -196,6 +199,9 @@ export interface NameRegistry {
    * makeValid("_"); // ERROR (too difficult to make valid to care)
    */
   makeValid(primer: string): string;
+
+  pushFunctionScope(): void;
+  popFunctionScope(): void;
 }
 
 function sanitizePrimer(primer: string | undefined) {
@@ -229,45 +235,78 @@ export function isValidIdentifier(ident: string): boolean {
 }
 
 abstract class NameRegistryImpl implements NameRegistry {
-  abstract makeUnique(primer?: string): string;
+  abstract getUniqueVariant(base: string): string;
+
+  readonly #usedNames: Set<string>;
+  readonly #usedFunctionScopeNamesStack: Set<string>[];
+
+  constructor() {
+    this.#usedNames = new Set<string>(bannedTokens);
+    this.#usedFunctionScopeNamesStack = [];
+  }
+
+  get usedFunctionScopeNames(): Set<string> | undefined {
+    return this
+      .#usedFunctionScopeNamesStack[
+        this.#usedFunctionScopeNamesStack.length - 1
+      ];
+  }
+
+  makeUnique(primer: string | undefined, global: boolean): string {
+    const sanitizedPrimer = sanitizePrimer(primer);
+    const name = this.getUniqueVariant(sanitizedPrimer);
+
+    if (global) {
+      this.#usedNames.add(name);
+    } else {
+      this.usedFunctionScopeNames?.add(name);
+    }
+
+    return name;
+  }
 
   makeValid(primer: string): string {
-    if (isValidIdentifier(primer)) {
+    if (isValidIdentifier(primer) && !this.#usedNames.has(primer)) {
+      this.usedFunctionScopeNames?.add(primer);
       return primer;
     }
-    return this.makeUnique(primer);
+    return this.makeUnique(primer, false);
   }
-}
 
-export class RandomNameRegistry extends NameRegistryImpl {
-  private lastUniqueId = 0;
+  isUsed(name: string): boolean {
+    return this.#usedNames.has(name) ||
+      !!this.usedFunctionScopeNames?.has(name);
+  }
 
-  makeUnique(primer?: string | undefined): string {
-    const sanitizedPrimer = sanitizePrimer(primer);
+  pushFunctionScope(): void {
+    this.#usedFunctionScopeNamesStack.push(new Set<string>());
+  }
 
-    return `${sanitizedPrimer}_${this.lastUniqueId++}`;
+  popFunctionScope(): void {
+    this.#usedFunctionScopeNamesStack.pop();
   }
 }
 
-export class StrictNameRegistry extends NameRegistryImpl {
-  /**
-   * Allows to provide a good fallback for instances of the
-   * same function that are bound to different slot values.
-   */
-  private readonly _usedNames = new Set<string>(bannedTokens);
+export class RandomNameRegistry extends NameRegistryImpl {
+  #lastUniqueId = 0;
 
-  // TODO: optimize this with a map
-  makeUnique(primer?: string | undefined): string {
-    const sanitizedPrimer = sanitizePrimer(primer);
+  getUniqueVariant(base: string): string {
+    let name = `${base}_${this.#lastUniqueId++}`;
+    while (this.isUsed(name)) {
+      name = `${base}_${this.#lastUniqueId++}`;
+    }
+    return name;
+  }
+}
 
+export class StrictNameRegistry extends NameRegistryImpl {
+  getUniqueVariant(base: string): string {
     let index = 0;
-    let label = sanitizedPrimer;
-    while (this._usedNames.has(label)) {
+    let name = base;
+    while (this.isUsed(name)) {
       index++;
-      label = `${sanitizedPrimer}_${index}`;
+      name = `${base}_${index}`;
     }
-
-    this._usedNames.add(label);
-    return label;
+    return name;
   }
 }