feat: tgpu.comptime, tgpu.rawCodeSnippet and this allowed in TypeGPU shader functions

apps/typegpu-docs/src/examples/algorithms/mnist-inference/index.html

@@ -2,8 +2,6 @@
 <div class="example">
   <canvas></canvas>
   <div class="predictions-container">
-    <div class="predictions-label">Predictions (-.--ms)</div>
-
     <div class="bars-container">
       <div class="bar">O</div>
       <div class="bar">1</div>

packages/tinyest-for-wgsl/src/parsers.ts

@@ -143,6 +143,14 @@ const Transpilers: Partial<
     return node.name;
   },
 
+  ThisExpression(ctx) {
+    if (ctx.ignoreExternalDepth === 0) {
+      ctx.externalNames.add('this');
+    }
+
+    return 'this';
+  },
+
   BinaryExpression(ctx, node) {
     const wgslOp = BINARY_OP_MAP[node.operator];
     const left = transpile(ctx, node.left) as tinyest.Expression;

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

@@ -0,0 +1,76 @@
+import type { DualFn } from '../../data/dualFn.ts';
+import type { MapValueToSnippet } from '../../data/snippet.ts';
+import { WgslTypeError } from '../../errors.ts';
+import { inCodegenMode } from '../../execMode.ts';
+import { setName, type TgpuNamable } from '../../shared/meta.ts';
+import { $getNameForward, $internal } from '../../shared/symbols.ts';
+import { coerceToSnippet } from '../../tgsl/generationHelpers.ts';
+import { isKnownAtComptime } from '../../types.ts';
+
+export type TgpuComptime<T extends (...args: never[]) => unknown> =
+  & DualFn<T>
+  & TgpuNamable
+  & { [$getNameForward]: unknown };
+
+/**
+ * Creates a version of `func` that can called safely in a TypeGPU function to
+ * precompute and inject a value into the final shader code.
+ *
+ * Note how the function passed into `comptime` doesn't have to be marked with
+ * 'use gpu'. That's because the function doesn't execute on the GPU, it gets
+ * executed before the shader code gets sent to the GPU.
+ *
+ * @example
+ * ```ts
+ * const injectRand01 = tgpu['~unstable']
+ *   .comptime(() => Math.random());
+ *
+ * const getColor = (diffuse: d.v3f) => {
+ *   'use gpu';
+ *   const albedo = hsvToRgb(injectRand01(), 1, 0.5);
+ *   return albedo.mul(diffuse);
+ * };
+ * ```
+ */
+export function comptime<T extends (...args: never[]) => unknown>(
+  func: T,
+): TgpuComptime<T> {
+  const gpuImpl = (...args: MapValueToSnippet<Parameters<T>>) => {
+    const argSnippets = args as MapValueToSnippet<Parameters<T>>;
+
+    if (!argSnippets.every((s) => isKnownAtComptime(s))) {
+      throw new WgslTypeError(
+        `Called comptime function with runtime-known values: ${
+          argSnippets.filter((s) => !isKnownAtComptime(s)).map((s) =>
+            `'${s.value}'`
+          ).join(', ')
+        }`,
+      );
+    }
+
+    return coerceToSnippet(func(...argSnippets.map((s) => s.value) as never[]));
+  };
+
+  const impl = ((...args: Parameters<T>) => {
+    if (inCodegenMode()) {
+      return gpuImpl(...args as MapValueToSnippet<Parameters<T>>);
+    }
+    return func(...args);
+  }) as TgpuComptime<T>;
+
+  impl.toString = () => 'comptime';
+  impl[$getNameForward] = func;
+  impl.$name = (label: string) => {
+    setName(func, label);
+    return impl;
+  };
+  Object.defineProperty(impl, $internal, {
+    value: {
+      jsImpl: func,
+      gpuImpl,
+      argConversionHint: 'keep',
+    },
+  });
+
+  return impl as TgpuComptime<T>;
+}

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

@@ -142,9 +142,13 @@ export function createFnCore(
       // get data generated by the plugin
       const pluginData = getMetaData(implementation);
 
-      if (pluginData?.externals) {
+      const pluginExternals = typeof pluginData?.externals === 'function'
+        ? pluginData.externals()
+        : pluginData?.externals;
+
+      if (pluginExternals) {
         const missing = Object.fromEntries(
-          Object.entries(pluginData.externals).filter(
+          Object.entries(pluginExternals).filter(
             ([name]) => !(name in externalMap),
           ),
         );

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

@@ -0,0 +1,166 @@
+import type { AnyData } from '../../data/dataTypes.ts';
+import { type Origin, type ResolvedSnippet, snip } from '../../data/snippet.ts';
+import { inCodegenMode } from '../../execMode.ts';
+import type { InferGPU } from '../../shared/repr.ts';
+import {
+  $gpuValueOf,
+  $internal,
+  $ownSnippet,
+  $resolve,
+} from '../../shared/symbols.ts';
+import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
+import {
+  applyExternals,
+  type ExternalMap,
+  replaceExternalsInWgsl,
+} from '../resolve/externals.ts';
+import { valueProxyHandler } from '../valueProxyUtils.ts';
+
+// ----------
+// Public API
+// ----------
+
+/**
+ * Extra declaration that shall be included in final WGSL code,
+ * when resolving objects that use it.
+ */
+export interface TgpuRawCodeSnippet<TDataType extends AnyData> {
+  $: InferGPU<TDataType>;
+  value: InferGPU<TDataType>;
+
+  $uses(dependencyMap: Record<string, unknown>): this;
+}
+
+// The origin 'function' refers to values passed in from the calling scope, which means
+// we would have access to this value anyway. Same goes for 'argument' and 'this-function',
+// the values literally exist in the function we're writing.
+//
+// 'constant-ref' was excluded because it's a special origin reserved for tgpu.const values.
+export type RawCodeSnippetOrigin = Exclude<
+  Origin,
+  'function' | 'this-function' | 'argument' | 'constant-ref'
+>;
+
+/**
+ * An advanced API that creates a typed shader expression which
+ * can be injected into the final shader bundle upon use.
+ *
+ * @param expression The code snippet that will be injected in place of `foo.$`
+ * @param type The type of the expression
+ * @param [origin='runtime'] Where the value originates from.
+ *
+ * **-- Which origin to choose?**
+ *
+ * Usually 'runtime' (the default) is a safe bet, but if you're sure that the expression or
+ * computation is constant (either a reference to a constant, a numeric literal,
+ * or an operation on constants), then pass 'constant' as it might lead to better
+ * optimisations.
+ *
+ * If what the expression is a direct reference to an existing value (e.g. a uniform, a
+ * storage binding, ...), then choose from 'uniform', 'mutable', 'readonly', 'workgroup',
+ * 'private' or 'handle' depending on the address space of the referred value.
+ *
+ * @example
+ * ```ts
+ * // An identifier that we know will be in the
+ * // final shader bundle, but we cannot
+ * // refer to it in any other way.
+ * const existingGlobal = tgpu['~unstable']
+ *   .rawCodeSnippet('EXISTING_GLOBAL', d.f32, 'constant');
+ *
+ * const foo = () => {
+ *   'use gpu';
+ *   return existingGlobal.$ * 2;
+ * };
+ *
+ * const wgsl = tgpu.resolve([foo]);
+ * // fn foo() -> f32 {
+ * //   return EXISTING_GLOBAL * 2;
+ * // }
+ * ```
+ */
+export function rawCodeSnippet<TDataType extends AnyData>(
+  expression: string,
+  type: TDataType,
+  origin: RawCodeSnippetOrigin | undefined = 'runtime',
+): TgpuRawCodeSnippet<TDataType> {
+  return new TgpuRawCodeSnippetImpl(expression, type, origin);
+}
+
+// --------------
+// Implementation
+// --------------
+
+class TgpuRawCodeSnippetImpl<TDataType extends AnyData>
+  implements TgpuRawCodeSnippet<TDataType>, SelfResolvable {
+  readonly [$internal]: true;
+  readonly dataType: TDataType;
+  readonly origin: RawCodeSnippetOrigin;
+
+  #expression: string;
+  #externalsToApply: ExternalMap[];
+
+  constructor(
+    expression: string,
+    type: TDataType,
+    origin: RawCodeSnippetOrigin,
+  ) {
+    this[$internal] = true;
+    this.dataType = type;
+    this.origin = origin;
+
+    this.#expression = expression;
+    this.#externalsToApply = [];
+  }
+
+  $uses(dependencyMap: Record<string, unknown>): this {
+    this.#externalsToApply.push(dependencyMap);
+    return this;
+  }
+
+  [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
+    const externalMap: ExternalMap = {};
+
+    for (const externals of this.#externalsToApply) {
+      applyExternals(externalMap, externals);
+    }
+
+    const replacedExpression = replaceExternalsInWgsl(
+      ctx,
+      externalMap,
+      this.#expression,
+    );
+
+    return snip(replacedExpression, this.dataType, this.origin);
+  }
+
+  toString() {
+    return `raw(${String(this.dataType)}): "${this.#expression}"`;
+  }
+
+  get [$gpuValueOf](): InferGPU<TDataType> {
+    const dataType = this.dataType;
+    const origin = this.origin;
+
+    return new Proxy({
+      [$internal]: true,
+      get [$ownSnippet]() {
+        return snip(this, dataType, origin);
+      },
+      [$resolve]: (ctx) => ctx.resolve(this),
+      toString: () => `raw(${String(this.dataType)}): "${this.#expression}".$`,
+    }, valueProxyHandler) as InferGPU<TDataType>;
+  }
+
+  get $(): InferGPU<TDataType> {
+    if (!inCodegenMode()) {
+      throw new Error('Raw code snippets can only be used on the GPU.');
+    }
+
+    return this[$gpuValueOf];
+  }
+
+  get value(): InferGPU<TDataType> {
+    return this.$;
+  }
+}

packages/typegpu/src/index.ts

@@ -6,8 +6,10 @@ import { constant } from './core/constant/tgpuConstant.ts';
 import { declare } from './core/declare/tgpuDeclare.ts';
 import { computeFn } from './core/function/tgpuComputeFn.ts';
 import { fn } from './core/function/tgpuFn.ts';
+import { rawCodeSnippet } from './core/rawCodeSnippet/tgpuRawCodeSnippet.ts';
 import { fragmentFn } from './core/function/tgpuFragmentFn.ts';
 import { vertexFn } from './core/function/tgpuVertexFn.ts';
+import { comptime } from './core/function/comptime.ts';
 import { resolve, resolveWithContext } from './core/resolve/tgpuResolve.ts';
 import { simulate } from './core/simulate/tgpuSimulate.ts';
 import { init, initFromDevice } from './core/root/init.ts';
@@ -44,6 +46,7 @@ export const tgpu = {
     fragmentFn,
     vertexFn,
     computeFn,
+    comptime,
     /**
      * @deprecated This feature is now stable, use tgpu.vertexLayout.
      */
@@ -68,6 +71,7 @@ export const tgpu = {
      */
     const: constant,
     declare,
+    rawCodeSnippet,
 
     simulate,
   },
@@ -139,6 +143,10 @@ export type {
   TgpuDerived,
   TgpuSlot,
 } from './core/slot/slotTypes.ts';
+export type {
+  RawCodeSnippetOrigin,
+  TgpuRawCodeSnippet,
+} from './core/rawCodeSnippet/tgpuRawCodeSnippet.ts';
 export type { TgpuTexture, TgpuTextureView } from './core/texture/texture.ts';
 export type { TextureProps } from './core/texture/textureProps.ts';
 export type { RenderFlag, SampledFlag } from './core/texture/usageExtension.ts';
@@ -167,6 +175,7 @@ export type {
   TgpuLayoutUniform,
 } from './tgpuBindGroupLayout.ts';
 export type { TgpuFn, TgpuFnShell } from './core/function/tgpuFn.ts';
+export type { TgpuComptime } from './core/function/comptime.ts';
 export type {
   TgpuVertexFn,
   TgpuVertexFnShell,

packages/typegpu/src/shared/meta.ts

@@ -10,7 +10,12 @@ export interface MetaData {
     body: Block;
     externalNames: string[];
   } | undefined;
-  externals?: Record<string, unknown> | undefined;
+  externals?:
+    // Passing a record happens prior to version 0.9.0
+    // TODO: Support for this can be removed down the line
+    | Record<string, unknown>
+    | (() => Record<string, unknown>)
+    | undefined;
 }
 
 /**