feat: Name collision avoidance in raw-WGSL implemented functions

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

@@ -1,4 +1,4 @@
-import { blankSpaces, lineBreaks } from '../whitespaces.ts';
+import { blankSpaces, lineBreaks } from '../../rawShaderCodeUtils.ts';
 
 interface FunctionArgsInfo {
   args: ArgInfo[];

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

@@ -4,7 +4,8 @@ import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
 import { type BaseData, isWgslData, isWgslStruct, Void } from '../../data/wgslTypes.ts';
 import { validateIdentifier } from '../../nameUtils.ts';
 import { getFunctionMetadata, getName } from '../../shared/meta.ts';
-import { $getNameForward } from '../../shared/symbols.ts';
+import { extractIdentifierLikeTokens, renameIdentifiers } from '../../rawShaderCodeUtils.ts';
+import { $getNameForward, $internal } from '../../shared/symbols.ts';
 import type { ResolutionCtx, ShaderStage } from '../../types.ts';
 import {
   type ExternalMap,
@@ -14,6 +15,7 @@ import {
 import { ResolvableString } from '../resolve/resolvableString.ts';
 import { extractArgs } from './extractArgs.ts';
 import type { Implementation, SeparatedEntryArgs } from './fnTypes.ts';
+import { parentFunctionNameSlot } from '../slot/internalSlots.ts';
 
 export type FnExternals = {
   /**
@@ -119,7 +121,7 @@ export function createFnCore(
                 `Invalid argument name "${arg.schemaKey}"${result.error ? `: ${result.error}` : ''}`,
               );
             }
-            if (ctx.isIdentifierBanned(arg.schemaKey)) {
+            if (ctx.gen.isBannedToken(arg.schemaKey)) {
               throw new Error(
                 `Invalid argument name "${arg.schemaKey}", the identifier is a reserved keyword.`,
               );
@@ -136,75 +138,119 @@ export function createFnCore(
           });
         }
 
-        const replacedImpl = replaceExternalsInWgsl(
-          ctx,
-          mergeFunctionExternals(externals),
-          implementation,
-        );
+        const externalMap = mergeFunctionExternals(externals);
+        try {
+          const scope = ctx[$internal].itemStateStack.pushFunctionScope(
+            functionType,
+            {},
+            returnType,
+            externalMap,
+          );
+          // Pushing a block scope as well, so that any identifiers declared at this point will be scoped to the function body.
+          ctx.pushBlockScope();
+
+          const externalKeys = Object.keys(externalMap);
+          const identifiers = [
+            ...new Set(
+              extractIdentifierLikeTokens(implementation).filter((ident) => {
+                return (
+                  !ctx.gen.isBannedToken(ident) &&
+                  !ctx.gen.isBuiltinGlobal(ident) &&
+                  !externalKeys.some((key) => key === ident || key.startsWith(`${ident}.`))
+                );
+              }),
+            ),
+          ];
 
-        let header = '';
-        let body = '';
+          const clashingIdentifiers = identifiers.filter((ident) =>
+            ctx.isIdentifierTaken(ident, 'block'),
+          );
 
-        if (functionType !== 'normal' && entryInput) {
-          const { dataSchema, positionalArgs } = entryInput;
-          const parts: string[] = [];
-          if (dataSchema && isArgUsedInBody('in', replacedImpl)) {
-            parts.push(`in: ${ctx.resolve(dataSchema).value}`);
-          }
-          for (const a of positionalArgs) {
-            const argName = a.schemaKey;
-            if (isArgUsedInBody(argName, replacedImpl)) {
-              parts.push(`${getAttributesString(a.type)}${argName}: ${ctx.resolve(a.type).value}`);
-            }
+          const uniqueIdentifiers = identifiers.filter(
+            (ident) => !ctx.isIdentifierTaken(ident, 'block'),
+          );
+
+          for (const ident of clashingIdentifiers) {
+            const renamed = ctx.makeUniqueIdentifier(ident, 'block');
+            scope.localRenames.set(ident, renamed);
           }
-          const input = `(${parts.join(', ')})`;
-
-          const attributes = isWgslData(returnType) ? getAttributesString(returnType) : '';
-          const output =
-            returnType !== Void
-              ? isWgslStruct(returnType)
-                ? ` -> ${ctx.resolve(returnType).value} `
-                : ` -> ${attributes !== '' ? attributes : '@location(0)'} ${
-                    ctx.resolve(returnType).value
-                  } `
-              : ' ';
-
-          header = `${input}${output}`;
-          body = replacedImpl;
-        } else {
-          const providedArgs = extractArgs(replacedImpl);
-
-          if (providedArgs.args.length !== argTypes.length) {
-            throw new Error(
-              `WGSL implementation has ${providedArgs.args.length} arguments, while the shell has ${argTypes.length} arguments.`,
-            );
+          const renamedImpl = renameIdentifiers(implementation, scope.localRenames);
+          for (const ident of uniqueIdentifiers) {
+            ctx.reserveIdentifier(ident, 'block');
           }
+          const replacedImpl = ctx.withSlots([[parentFunctionNameSlot, id]], () =>
+            replaceExternalsInWgsl(ctx, externalMap, renamedImpl),
+          );
+
+          let header = '';
+          let body = '';
+
+          if (functionType !== 'normal' && entryInput) {
+            const { dataSchema, positionalArgs } = entryInput;
+            const parts: string[] = [];
+            if (dataSchema && isArgUsedInBody('in', replacedImpl)) {
+              parts.push(`in: ${ctx.resolve(dataSchema).value}`);
+            }
+            for (const a of positionalArgs) {
+              const argName = a.schemaKey;
+              if (isArgUsedInBody(argName, replacedImpl)) {
+                parts.push(
+                  `${getAttributesString(a.type)}${argName}: ${ctx.resolve(a.type).value}`,
+                );
+              }
+            }
+            const input = `(${parts.join(', ')})`;
+
+            const attributes = isWgslData(returnType) ? getAttributesString(returnType) : '';
+            const output =
+              returnType !== Void
+                ? isWgslStruct(returnType)
+                  ? ` -> ${ctx.resolve(returnType).value} `
+                  : ` -> ${attributes !== '' ? attributes : '@location(0)'} ${
+                      ctx.resolve(returnType).value
+                    } `
+                : ' ';
+
+            header = `${input}${output}`;
+            body = replacedImpl;
+          } else {
+            const providedArgs = extractArgs(replacedImpl);
+
+            if (providedArgs.args.length !== argTypes.length) {
+              throw new Error(
+                `WGSL implementation has ${providedArgs.args.length} arguments, while the shell has ${argTypes.length} arguments.`,
+              );
+            }
 
-          const input = providedArgs.args
-            .map(
-              (argInfo, i) =>
-                `${argInfo.identifier}: ${checkAndReturnType(
-                  ctx,
-                  `parameter ${argInfo.identifier}`,
-                  argInfo.type,
-                  argTypes[i],
-                )}`,
-            )
-            .join(', ');
-
-          const output =
-            returnType === Void
-              ? ' '
-              : ` -> ${checkAndReturnType(ctx, 'return type', providedArgs.ret?.type, returnType)} `;
-
-          header = `(${input})${output}`;
-
-          body = replacedImpl.slice(providedArgs.range.end);
-        }
+            const input = providedArgs.args
+              .map(
+                (argInfo, i) =>
+                  `${argInfo.identifier}: ${checkAndReturnType(
+                    ctx,
+                    `parameter ${argInfo.identifier}`,
+                    argInfo.type,
+                    argTypes[i],
+                  )}`,
+              )
+              .join(', ');
+
+            const output =
+              returnType === Void
+                ? ' '
+                : ` -> ${checkAndReturnType(ctx, 'return type', providedArgs.ret?.type, returnType)} `;
+
+            header = `(${input})${output}`;
+
+            body = replacedImpl.slice(providedArgs.range.end);
+          }
 
-        ctx.addDeclaration(`${attributes}fn ${id}${header}${body}`, id);
+          ctx.addDeclaration(`${attributes}fn ${id}${header}${body}`, id);
 
-        return snip(id, returnType, /* origin */ 'runtime');
+          return snip(id, returnType, /* origin */ 'runtime');
+        } finally {
+          ctx[$internal].itemStateStack.pop('blockScope');
+          ctx[$internal].itemStateStack.pop('functionScope');
+        }
       }
 
       // get data generated by the plugin

packages/typegpu/src/core/rawCodeSnippet/tgpuRawCodeSnippet.ts

@@ -6,6 +6,8 @@ import { makeResolvable } from '../../tgsl/makeResolvable.ts';
 import type { InferGPU } from '../../shared/repr.ts';
 import { $gpuValueOf, $internal } from '../../shared/symbols.ts';
 import { type ExternalMap, replaceExternalsInWgsl } from '../resolve/externals.ts';
+import { renameIdentifiers } from '../../rawShaderCodeUtils.ts';
+import { parentFunctionNameSlot } from '../slot/internalSlots.ts';
 
 // ----------
 // Public API
@@ -106,11 +108,17 @@ class TgpuRawCodeSnippetImpl<TDataType extends BaseData> implements TgpuRawCodeS
           return `raw(${String(this.dataType)}): "${this.#expression}"`;
         },
         resolve(ctx) {
-          const replacedExpression = replaceExternalsInWgsl(
-            ctx,
-            this.#externals ?? {},
-            this.#expression,
-          );
+          // The code of the snippet can actually change depending on which function it's referenced in.
+          // For example an identifier can refer to a locally defined variable or argument, or a global,
+          // and therefore can be renamed or left as is.
+          void parentFunctionNameSlot.$;
+
+          let expression = this.#expression;
+          if (ctx.topFunctionScope) {
+            expression = renameIdentifiers(expression, ctx.topFunctionScope.localRenames);
+          }
+
+          const replacedExpression = replaceExternalsInWgsl(ctx, this.#externals ?? {}, expression);
 
           return snip(replacedExpression, this.dataType, this.origin, this.possibleSideEffects);
         },

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

@@ -1,6 +1,7 @@
 import { isLooseData } from '../../data/dataTypes.ts';
 import { isWgslStruct } from '../../data/wgslTypes.ts';
 import { getName, hasTinyestMetadata, isNamable, setName } from '../../shared/meta.ts';
+import { anyIdent } from '../../rawShaderCodeUtils.ts';
 import { logger } from '../../tgpuLogger.ts';
 import { isWgsl, type ResolutionCtx } from '../../types.ts';
 import type { FnExternals } from '../function/fnCore.ts';
@@ -67,7 +68,6 @@ export function addReturnTypeToExternals(
   }
 }
 
-export const anyIdent = /([$_\p{XID_Start}][$\p{XID_Continue}]*)/u; // WGSL ident, modified to include $
 const anyPropChain = new RegExp(`(${anyIdent.source})(\\.${anyIdent.source})*`, 'ug');
 export const boundedPropChain = new RegExp(
   `(?<![\\p{XID_Continue}\\$.])${anyPropChain.source}(?![\\p{XID_Continue}\\$])`,

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

@@ -1,5 +1,4 @@
 import type { ResolvedSnippet } from '../../data/snippet.ts';
-import { bannedTokens, builtins } from '../../nameUtils.ts';
 import { $internal } from '../../shared/symbols.ts';
 import { ShelllessRepository } from '../../tgsl/shellless.ts';
 import type { TgpuLazy, TgpuSlot } from '../slot/slotTypes.ts';
@@ -36,7 +35,7 @@ class NamespaceImpl implements Namespace {
   constructor(strategy: 'random' | 'strict') {
     this[$internal] = {
       strategy,
-      takenGlobalIdentifiers: new Set([...bannedTokens, ...builtins]),
+      takenGlobalIdentifiers: new Set(),
       shelllessRepo: new ShelllessRepository(),
       memoizedResolves: new WeakMap(),
       memoizedLazy: new WeakMap(),

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

@@ -38,6 +38,7 @@ import type {
   WgslArray,
   WgslStruct,
 } from '../../data/wgslTypes.ts';
+import { validateProp } from '../../nameUtils.ts';
 import { getName } from '../../shared/meta.ts';
 import { $internal } from '../../shared/symbols.ts';
 import { assertExhaustive } from '../../shared/utilityTypes.ts';
@@ -125,6 +126,13 @@ function resolveStructProperty(ctx: ResolutionCtx, [key, property]: [string, Bas
  * @returns The resolved struct name.
  */
 function resolveStruct(ctx: ResolutionCtx, struct: WgslStruct) {
+  Object.keys(struct.propTypes).forEach((key) => {
+    const result = validateProp(ctx, key);
+    if (!result.success) {
+      throw new Error(`Invalid property key '${key}'${result.error ? `: ${result.error}` : ''}`);
+    }
+  });
+
   if (struct[$internal].isAbstruct) {
     throw new Error('Cannot resolve abstract struct types to WGSL.');
   }

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

@@ -1,6 +1,7 @@
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
 import { Void } from '../../data/wgslTypes.ts';
 import { getName } from '../../internal.ts';
+import { extractIdentifierLikeTokens, renameIdentifiers } from '../../rawShaderCodeUtils.ts';
 import { type ResolutionResult, resolve as resolveImpl } from '../../resolutionCtx.ts';
 import { $internal, $resolve, $soul } from '../../shared/symbols.ts';
 import { isBindGroupLayout } from '../../tgpuBindGroupLayout.ts';
@@ -211,11 +212,55 @@ function resolveFromTemplate(options: TgpuExtendedResolveOptions): ResolutionRes
   const resolutionObj: SelfResolvable = {
     [$internal]: true,
     [$resolve](ctx): ResolvedSnippet {
-      return snip(
-        replaceExternalsInWgsl(ctx, externals, template ?? ''),
-        Void,
-        /* origin */ 'runtime',
-      );
+      try {
+        // It's technically not a function we're resolving, but we need a place to store local renames, so we treat the whole
+        const scope = ctx[$internal].itemStateStack.pushFunctionScope(
+          'normal',
+          {},
+          Void,
+          externals,
+        );
+        // Pushing a block scope as well, so that any identifiers declared at this point will be scoped to the function body.
+        ctx.pushBlockScope();
+
+        const identifiers = [
+          ...new Set(
+            extractIdentifierLikeTokens(template ?? '').filter((ident) => {
+              return (
+                !ctx.gen.isBannedToken(ident) &&
+                !ctx.gen.isBuiltinGlobal(ident) &&
+                externals[ident] === undefined
+              );
+            }),
+          ),
+        ];
+
+        const clashingIdentifiers = identifiers.filter((ident) =>
+          ctx.isIdentifierTaken(ident, 'global'),
+        );
+
+        const uniqueIdentifiers = identifiers.filter(
+          (ident) => !ctx.isIdentifierTaken(ident, 'global'),
+        );
+
+        for (const ident of clashingIdentifiers) {
+          const renamed = ctx.makeUniqueIdentifier(ident, 'global');
+          scope.localRenames.set(ident, renamed);
+        }
+        const renamedImpl = renameIdentifiers(template ?? '', scope.localRenames);
+        for (const ident of uniqueIdentifiers) {
+          ctx.reserveIdentifier(ident, 'global');
+        }
+
+        return snip(
+          replaceExternalsInWgsl(ctx, externals, renamedImpl),
+          Void,
+          /* origin */ 'runtime',
+        );
+      } finally {
+        ctx[$internal].itemStateStack.pop('blockScope');
+        ctx[$internal].itemStateStack.pop('functionScope');
+      }
     },
     toString: () => '<root>',
   };

packages/typegpu/src/core/slot/internalSlots.ts

@@ -2,3 +2,4 @@ import type { ShaderStage } from '../../types.ts';
 import { slot } from './slot.ts';
 
 export const shaderStageSlot = slot<ShaderStage | null>(null);
+export const parentFunctionNameSlot = slot<string | null>(null);