impr: Less work on the shader generation happy path

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

@@ -20,6 +20,10 @@ function isResolvable(value: unknown) {
  */
 export function mergeFunctionExternals(fnExternals: FnExternals): ExternalMap {
   const base = fnExternals.pluginProvided ?? fnExternals.userProvided ?? {};
+  if (fnExternals.args === undefined && fnExternals.out === undefined) {
+    // Nothing to merge in, and consumers never mutate the map
+    return base;
+  }
   // avoid calling any of the getters
   const result: ExternalMap = Object.defineProperties({}, Object.getOwnPropertyDescriptors(base));
   for (const flatExternal of [fnExternals.args, fnExternals.out].filter((e) => e !== undefined)) {

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

@@ -1,12 +1,15 @@
 import type { ResolvedSnippet } from '../../data/snippet.ts';
-import { bannedTokens, builtins, templateEnumerants } from '../../nameUtils.ts';
 import { $internal } from '../../shared/symbols.ts';
 import { ShelllessRepository } from '../../tgsl/shellless.ts';
 import type { TgpuLazy, TgpuSlot } from '../slot/slotTypes.ts';
 
 type SlotToValueMap = Map<TgpuSlot<unknown>, unknown>;
 
 export interface NamespaceInternal {
+  /**
+   * Global identifiers taken by resolved items. Reserved WGSL words are checked separately,
+   * so that creating a namespace doesn't have to copy them.
+   */
   readonly takenGlobalIdentifiers: Set<string>;
   readonly shelllessRepo: ShelllessRepository;
   readonly strategy: 'random' | 'strict';
@@ -36,7 +39,7 @@ class NamespaceImpl implements Namespace {
   constructor(strategy: 'random' | 'strict') {
     this[$internal] = {
       strategy,
-      takenGlobalIdentifiers: new Set([...bannedTokens, ...builtins, ...templateEnumerants]),
+      takenGlobalIdentifiers: new Set(),
       shelllessRepo: new ShelllessRepository(),
       memoizedResolves: new WeakMap(),
       memoizedLazy: new WeakMap(),

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

@@ -4,6 +4,10 @@ import type { ResolutionCtx } from '../../types.ts';
 
 type ValueOrArray<T> = T | readonly T[];
 
+function resolvePart(ctx: ResolutionCtx, part: Snippet | string | number) {
+  return isSnippet(part) ? ctx.resolveSnippet(part).value : part;
+}
+
 /**
  * "The reverse of snipping"
  * Injects resolved snippets into a template string.
@@ -14,21 +18,22 @@ export function stitch(
 ) {
   const ctx = getResolutionCtx() as ResolutionCtx;
 
-  function resolveSnippet(maybeSnippet: Snippet | string | number) {
-    return isSnippet(maybeSnippet) ? ctx.resolveSnippet(maybeSnippet).value : maybeSnippet;
-  }
-
   let result = '';
   for (let i = 0; i < strings.length; ++i) {
     result += strings[i];
     const snippet = snippets[i];
     if (Array.isArray(snippet)) {
-      result += snippet
-        .filter((s) => s !== undefined)
-        .map(resolveSnippet)
-        .join(', ');
+      // Joining the defined elements with ', '
+      let first = true;
+      for (const element of snippet) {
+        if (element === undefined) {
+          continue;
+        }
+        result += (first ? '' : ', ') + resolvePart(ctx, element);
+        first = false;
+      }
     } else if (snippet) {
-      result += resolveSnippet(snippet as Snippet | string | number);
+      result += resolvePart(ctx, snippet as Snippet | string | number);
     }
   }
   return result;

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

@@ -743,10 +743,15 @@ class TgpuFixedTextureViewImpl<T extends WgslTexture | WgslStorageTexture>
     return this;
   }
 
+  // Created once, since shader code can access a view many times
+  #valueProxy: Infer<T> | undefined;
+
   get [$gpuValueOf](): Infer<T> {
+    if (this.#valueProxy) {
+      return this.#valueProxy;
+    }
     const schema = this.schema;
-
-    return new Proxy(
+    this.#valueProxy = new Proxy(
       {
         [$internal]: true,
         get [$ownSnippet]() {
@@ -758,6 +763,7 @@ class TgpuFixedTextureViewImpl<T extends WgslTexture | WgslStorageTexture>
       },
       valueProxyHandler,
     ) as unknown as Infer<T>;
+    return this.#valueProxy;
   }
 
   get $(): Infer<T> {
@@ -837,9 +843,15 @@ export class TgpuLaidOutTextureViewImpl<T extends WgslTexture | WgslStorageTextu
     return snip(id, this.schema, /* origin */ 'handle');
   }
 
+  // Created once, since shader code can access a view many times
+  #valueProxy: Infer<T> | undefined;
+
   get [$gpuValueOf](): Infer<T> {
+    if (this.#valueProxy) {
+      return this.#valueProxy;
+    }
     const schema = this.schema;
-    return new Proxy(
+    this.#valueProxy = new Proxy(
       {
         [$internal]: true,
         get [$ownSnippet]() {
@@ -851,6 +863,7 @@ export class TgpuLaidOutTextureViewImpl<T extends WgslTexture | WgslStorageTextu
       },
       valueProxyHandler,
     ) as unknown as Infer<T>;
+    return this.#valueProxy;
   }
 
   get $(): Infer<T> {

packages/typegpu/src/data/dataTypes.ts

@@ -197,7 +197,8 @@ export function isUnstruct(schema: unknown): schema is Unstruct {
 }
 
 export function isLooseDecorated(value: unknown): value is LooseDecorated {
-  return isMarkedInternal(value) && (value as LooseDecorated)?.type === 'loose-decorated';
+  // Testing the cheap `type` field first for performance reasons
+  return (value as LooseDecorated)?.type === 'loose-decorated' && isMarkedInternal(value);
 }
 
 export function getCustomAlignment(data: wgsl.BaseData): number | undefined {

packages/typegpu/src/data/wgslTypes.ts

@@ -1693,7 +1693,8 @@ export function isWgslStruct(schema: unknown): schema is WgslStruct {
  * isPtr(d.f32) // false
  */
 export function isPtr(schema: unknown): schema is Ptr {
-  return isMarkedInternal(schema) && (schema as Ptr)?.type === 'ptr';
+  // Testing the cheap `type` field first for performance reasons
+  return (schema as Ptr)?.type === 'ptr' && isMarkedInternal(schema);
 }
 
 /**
@@ -1733,7 +1734,8 @@ export function isInvariantAttrib(value: unknown): value is Invariant {
 }
 
 export function isDecorated(value: unknown): value is Decorated {
-  return isMarkedInternal(value) && (value as Decorated)?.type === 'decorated';
+  // Testing the cheap `type` field first for performance reasons
+  return (value as Decorated)?.type === 'decorated' && isMarkedInternal(value);
 }
 
 export function isAbstractFloat(value: unknown): value is AbstractFloat {

packages/typegpu/src/resolutionCtx.ts

@@ -66,7 +66,13 @@ import { isTgpuFn } from './core/function/tgpuFn.ts';
 import type { IOData } from './core/function/fnTypes.ts';
 import { AutoStruct } from './data/autoStruct.ts';
 import { EntryInputRouter } from './core/function/entryInputRouter.ts';
-import { validateIdentifier, sanitizePrimer, bannedTokens } from './nameUtils.ts';
+import {
+  validateIdentifier,
+  sanitizePrimer,
+  bannedTokens,
+  builtins,
+  templateEnumerants,
+} from './nameUtils.ts';
 import { minify } from './minify.ts';
 
 /**
@@ -155,6 +161,7 @@ class ItemStateStackImpl implements ItemStateStack {
       argAccess,
       returnType,
       externalMap,
+      externalSnippets: new Map(),
       reportedReturnTypes: new Set(),
       placeholderForVariable: new Map(),
       modifiedVariables: new Set(),
@@ -220,12 +227,18 @@ class ItemStateStackImpl implements ItemStateStack {
           return access();
         }
 
+        const cached = layer.externalSnippets.get(id);
+        if (cached) {
+          return cached;
+        }
         if (Object.hasOwn(layer.externalMap, id)) {
           const external = layer.externalMap[id];
           if (isNamable(external) && getName(external) === undefined) {
             setName(external, id.replaceAll('.', '_'));
           }
-          return coerceToSnippet(external);
+          const snippet = coerceToSnippet(external);
+          layer.externalSnippets.set(id, snippet);
+          return snippet;
         }
 
         return undefined;
@@ -491,6 +504,9 @@ export class ResolutionCtxImpl implements ResolutionCtx {
 
   isIdentifierTaken(name: string, scope: 'global' | 'block'): boolean {
     return (
+      bannedTokens.has(name) ||
+      builtins.has(name) ||
+      templateEnumerants.has(name) ||
       this.#namespaceInternal.takenGlobalIdentifiers.has(name) ||
       (scope === 'block'
         ? this._itemStateStack.isIdentifierTakenLocally(name)
@@ -911,6 +927,18 @@ export class ResolutionCtxImpl implements ResolutionCtx {
     return maybeEventual;
   }
 
+  /**
+   * Whether every slot in the map currently has a value equal to the one recorded in it.
+   */
+  #slotsHaveValues(slotToValueMap: Map<TgpuSlot<unknown>, unknown>): boolean {
+    for (const [slot, expectedValue] of slotToValueMap) {
+      if (!slot.areEqual(this._itemStateStack.readSlot(slot), expectedValue)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   _getOrCompute<T>(lazy: TgpuLazy<T>): T {
     // All memoized versions of `lazy`
     const instances = this.#namespaceInternal.memoizedLazy.get(lazy) ?? [];
@@ -919,13 +947,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
 
     try {
       for (const instance of instances) {
-        const slotValuePairs = [...instance.slotToValueMap.entries()];
-
-        if (
-          slotValuePairs.every(([slot, expectedValue]) =>
-            slot.areEqual(this._itemStateStack.readSlot(slot), expectedValue),
-          )
-        ) {
+        if (this.#slotsHaveValues(instance.slotToValueMap)) {
           return instance.result as T;
         }
       }
@@ -972,13 +994,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
 
     try {
       for (const instance of instances) {
-        const slotValuePairs = [...instance.slotToValueMap.entries()];
-
-        if (
-          slotValuePairs.every(([slot, expectedValue]) =>
-            slot.areEqual(this._itemStateStack.readSlot(slot), expectedValue),
-          )
-        ) {
+        if (this.#slotsHaveValues(instance.slotToValueMap)) {
           return instance.result;
         }
       }

packages/typegpu/src/tgsl/conversion.ts

@@ -145,10 +145,37 @@ function getImplicitConversionRank(src: BaseData, dest: BaseData): ConversionRan
   return INFINITE_RANK;
 }
 
+/**
+ * Conversion ranks only depend on immutable facts about the two schemas (their type tags,
+ * identity, inner types and pointer flags), so they're computed once per pair. Callers must
+ * treat the returned info as read-only.
+ */
+const autoRankCache = new WeakMap<BaseData, WeakMap<BaseData, ConversionRankInfo>>();
+const implicitRankCache = new WeakMap<BaseData, WeakMap<BaseData, ConversionRankInfo>>();
+
 function getConversionRank(
   src: BaseData,
   dest: BaseData,
   allowImplicit: boolean,
+): ConversionRankInfo {
+  const cache = allowImplicit ? implicitRankCache : autoRankCache;
+  let byDest = cache.get(src);
+  if (!byDest) {
+    byDest = new WeakMap();
+    cache.set(src, byDest);
+  }
+  let rank = byDest.get(dest);
+  if (!rank) {
+    rank = computeConversionRank(src, dest, allowImplicit);
+    byDest.set(dest, rank);
+  }
+  return rank;
+}
+
+function computeConversionRank(
+  src: BaseData,
+  dest: BaseData,
+  allowImplicit: boolean,
 ): ConversionRankInfo {
   const autoRank = getAutoConversionRank(src, dest);
   if (autoRank.rank < Number.POSITIVE_INFINITY) {

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -807,7 +807,7 @@ export class WgslGenerator implements ShaderGenerator {
         typeof expression[1] === 'string'
           ? numericLiteralToSnippet(parseNumericString(expression[1]))
           : numericLiteralToSnippet(expression[1]);
-      invariant(type, `Expected ${stringifyNode(expression)} to be valid numeric literal`);
+      invariant(type, () => `Expected ${stringifyNode(expression)} to be valid numeric literal`);
       return type;
     }
 
@@ -1266,17 +1266,17 @@ The generated shader will evaluate them in declaration order.`,
     let body = this._block(options.body, /* allowInlining */ false);
     const scope = this.ctx.topFunctionScope;
     invariant(scope, 'Expected function scope to be present');
-    const replacements = Object.fromEntries(
-      [...scope.placeholderForVariable.entries()].map(([variable, placeholder]) => [
-        placeholder,
-        scope.modifiedVariables.has(variable) ? 'var' : 'let',
-      ]),
-    );
-    if (Object.keys(replacements).length > 0) {
-      const regex = new RegExp(Object.keys(replacements).join('|'), 'gi');
+    if (scope.placeholderForVariable.size > 0) {
+      const keywordForPlaceholder = new Map<string, string>();
+      for (const [variable, placeholder] of scope.placeholderForVariable) {
+        keywordForPlaceholder.set(
+          placeholder,
+          scope.modifiedVariables.has(variable) ? 'var' : 'let',
+        );
+      }
       body.code = body.code.replace(
-        regex,
-        (match) => replacements[match as keyof typeof replacements] ?? '#ERR',
+        /#VAR_\d+#/g,
+        (match) => keywordForPlaceholder.get(match) ?? '#ERR',
       );
     }
 

packages/typegpu/src/types.ts

@@ -119,6 +119,8 @@ export type FunctionScopeLayer = {
   functionType: 'normal' | 'compute' | 'vertex' | 'fragment';
   argAccess: Record<string, FunctionArgumentAccess>;
   externalMap: Record<string, unknown>;
+  /** Snippets of externals already looked up in this function, so each is only coerced once */
+  externalSnippets: Map<string, Snippet>;
   /**
    * The return type of the function. If undefined, the type should be inferred
    * from the implementation (relevant for shellless functions).