fix: Auto dedent 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,12 @@ 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,
+  normalizeIndentation,
+  renameIdentifiers,
+} from '../../rawShaderCodeUtils.ts';
+import { $getNameForward, $internal } from '../../shared/symbols.ts';
 import type { ResolutionCtx, ShaderStage } from '../../types.ts';
 import {
   type ExternalMap,
@@ -119,7 +124,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 +141,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 = 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 = new Set(
+            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 = new Set(
+            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(
+            normalizeIndentation(implementation),
+            scope.localRenames,
+          );
+
+          for (const ident of uniqueIdentifiers) {
+            ctx.reserveIdentifier(ident, 'block');
           }
 
-          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 replacedImpl = replaceExternalsInWgsl(ctx, externalMap, renamedImpl);
 
-        ctx.addDeclaration(`${attributes}fn ${id}${header}${body}`, id);
+          let header = '';
+          let body = '';
 
-        return snip(id, returnType, /* origin */ 'runtime');
+          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);
+          }
+
+          ctx.addDeclaration(`${attributes}fn ${id}${header}${body}`, id);
+
+          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,7 @@ 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';
 
 // ----------
 // Public API
@@ -106,11 +107,12 @@ class TgpuRawCodeSnippetImpl<TDataType extends BaseData> implements TgpuRawCodeS
           return `raw(${String(this.dataType)}): "${this.#expression}"`;
         },
         resolve(ctx) {
-          const replacedExpression = replaceExternalsInWgsl(
-            ctx,
-            this.#externals ?? {},
-            this.#expression,
-          );
+          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/data/autoStruct.ts

@@ -1,5 +1,4 @@
 import { createIoSchema } from '../core/function/ioSchema.ts';
-import { validateProp } from '../nameUtils.ts';
 import { getName, setName } from '../shared/meta.ts';
 import { $internal, $repr, $resolve } from '../shared/symbols.ts';
 import type { ResolutionCtx, SelfResolvable } from '../types.ts';
@@ -75,10 +74,6 @@ export class AutoStruct implements BaseData, SelfResolvable {
           `Property name '${wgslKey}' causes naming clashes. Choose a different name.`,
         );
       }
-      const result = validateProp(wgslKey);
-      if (!result.success) {
-        throw new Error(`Invalid property key '${key}'${result.error ? `: ${result.error}` : ''}`);
-      }
 
       this.#usedWgslKeys.add(wgslKey);
       alloc = { prop: wgslKey, type: dataType };

packages/typegpu/src/data/struct.ts

@@ -1,4 +1,3 @@
-import { validateProp } from '../nameUtils.ts';
 import { getName, setName } from '../shared/meta.ts';
 import { $internal } from '../shared/symbols.ts';
 import { schemaCallWrapper } from './schemaCallWrapper.ts';
@@ -39,13 +38,6 @@ export function INTERNAL_createStruct<TProps extends Record<string, BaseData>>(
   props: TProps,
   isAbstruct: boolean,
 ): WgslStruct<TProps> {
-  Object.keys(props).forEach((key) => {
-    const result = validateProp(key);
-    if (!result.success) {
-      throw new Error(`Invalid property key '${key}'${result.error ? `: ${result.error}` : ''}`);
-    }
-  });
-
   // In the schema call, create and return a deep copy
   // by wrapping all the values in corresponding schema calls.
   const structSchema = (instanceProps?: TProps) =>