feat: tgpu.resolveWithContext

packages/typegpu/src/core/pipeline/computePipeline.ts

@@ -74,8 +74,8 @@ type TgpuComputePipelinePriors = {
 
 type Memo = {
   pipeline: GPUComputePipeline;
-  bindGroupLayouts: TgpuBindGroupLayout[];
-  catchall: [number, TgpuBindGroup] | null;
+  usedBindGroupLayouts: TgpuBindGroupLayout[];
+  catchall: [number, TgpuBindGroup] | undefined;
 };
 
 class TgpuComputePipelineImpl implements TgpuComputePipeline {
@@ -156,9 +156,9 @@ class TgpuComputePipelineImpl implements TgpuComputePipeline {
 
     pass.setPipeline(memo.pipeline);
 
-    const missingBindGroups = new Set(memo.bindGroupLayouts);
+    const missingBindGroups = new Set(memo.usedBindGroupLayouts);
 
-    memo.bindGroupLayouts.forEach((layout, idx) => {
+    memo.usedBindGroupLayouts.forEach((layout, idx) => {
       if (memo.catchall && idx === memo.catchall[0]) {
         // Catch-all
         pass.setBindGroup(idx, branch.unwrap(memo.catchall[1]));
@@ -207,7 +207,7 @@ class ComputePipelineCore {
       const device = this.branch.device;
 
       // Resolving code
-      const { code, bindGroupLayouts, catchall } = resolve(
+      const { code, usedBindGroupLayouts, catchall } = resolve(
         {
           '~resolve': (ctx) => {
             ctx.withSlots(this._slotBindings, () => {
@@ -224,8 +224,8 @@ class ComputePipelineCore {
         },
       );
 
-      if (catchall !== null) {
-        bindGroupLayouts[catchall[0]]?.$name(
+      if (catchall !== undefined) {
+        usedBindGroupLayouts[catchall[0]]?.$name(
           `${getName(this) ?? '<unnamed>'} - Automatic Bind Group & Layout`,
         );
       }
@@ -235,7 +235,7 @@ class ComputePipelineCore {
           label: getName(this) ?? '<unnamed>',
           layout: device.createPipelineLayout({
             label: `${getName(this) ?? '<unnamed>'} - Pipeline Layout`,
-            bindGroupLayouts: bindGroupLayouts.map((l) =>
+            bindGroupLayouts: usedBindGroupLayouts.map((l) =>
               this.branch.unwrap(l)
             ),
           }),
@@ -246,7 +246,7 @@ class ComputePipelineCore {
             }),
           },
         }),
-        bindGroupLayouts,
+        usedBindGroupLayouts,
         catchall,
       };
     }

packages/typegpu/src/core/pipeline/renderPipeline.ts

@@ -232,8 +232,8 @@ type TgpuRenderPipelinePriors = {
 
 type Memo = {
   pipeline: GPURenderPipeline;
-  bindGroupLayouts: TgpuBindGroupLayout[];
-  catchall: [number, TgpuBindGroup] | null;
+  usedBindGroupLayouts: TgpuBindGroupLayout[];
+  catchall: [number, TgpuBindGroup] | undefined;
 };
 
 class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
@@ -390,9 +390,9 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
 
     pass.setPipeline(memo.pipeline);
 
-    const missingBindGroups = new Set(memo.bindGroupLayouts);
+    const missingBindGroups = new Set(memo.usedBindGroupLayouts);
 
-    memo.bindGroupLayouts.forEach((layout, idx) => {
+    memo.usedBindGroupLayouts.forEach((layout, idx) => {
       if (memo.catchall && idx === memo.catchall[0]) {
         // Catch-all
         pass.setBindGroup(idx, branch.unwrap(memo.catchall[1]));
@@ -473,7 +473,7 @@ class RenderPipelineCore {
       } = this.options;
 
       // Resolving code
-      const { code, bindGroupLayouts, catchall } = resolve(
+      const { code, usedBindGroupLayouts, catchall } = resolve(
         {
           '~resolve': (ctx) => {
             ctx.withSlots(slotBindings, () => {
@@ -491,8 +491,8 @@ class RenderPipelineCore {
         },
       );
 
-      if (catchall !== null) {
-        bindGroupLayouts[catchall[0]]?.$name(
+      if (catchall !== undefined) {
+        usedBindGroupLayouts[catchall[0]]?.$name(
           `${getName(this) ?? '<unnamed>'} - Automatic Bind Group & Layout`,
         );
       }
@@ -507,7 +507,7 @@ class RenderPipelineCore {
       const descriptor: GPURenderPipelineDescriptor = {
         layout: device.createPipelineLayout({
           label: `${getName(this) ?? '<unnamed>'} - Pipeline Layout`,
-          bindGroupLayouts: bindGroupLayouts.map((l) => branch.unwrap(l)),
+          bindGroupLayouts: usedBindGroupLayouts.map((l) => branch.unwrap(l)),
         }),
         vertex: {
           module,
@@ -538,7 +538,7 @@ class RenderPipelineCore {
 
       this._memo = {
         pipeline: device.createRenderPipeline(descriptor),
-        bindGroupLayouts,
+        usedBindGroupLayouts,
         catchall,
       };
     }

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

@@ -1,6 +1,9 @@
 import type { JitTranspiler } from '../../jitTranspiler.ts';
 import { RandomNameRegistry, StrictNameRegistry } from '../../nameRegistry.ts';
-import { resolve as resolveImpl } from '../../resolutionCtx.ts';
+import {
+  type ResolutionResult,
+  resolve as resolveImpl,
+} from '../../resolutionCtx.ts';
 import type { SelfResolvable, Wgsl } from '../../types.ts';
 import { applyExternals, replaceExternalsInWgsl } from './externals.ts';
 
@@ -31,7 +34,7 @@ export interface TgpuResolveOptions {
  * Any dependencies of the externals will also be resolved and included in the output.
  * @param options - The options for the resolution.
  *
- * @returns The resolved code.
+ * @returns {ResolutionResult}
  *
  * @example
  * ```ts
@@ -40,7 +43,7 @@ export interface TgpuResolveOptions {
  *   to: d.vec3f,
  * });
  *
- * const resolved = tgpu.resolve({
+ * const { code, usedBindGroupLayouts, catchall } = tgpu.resolveWithContext({
  *   template: `
  *     fn getGradientAngle(gradient: Gradient) -> f32 {
  *       return atan(gradient.to.y - gradient.from.y, gradient.to.x - gradient.from.x);
@@ -51,7 +54,7 @@ export interface TgpuResolveOptions {
  *   },
  * });
  *
- * console.log(resolved);
+ * console.log(code);
  * // struct Gradient_0 {
  * //   from: vec3f,
  * //   to: vec3f,
@@ -61,7 +64,9 @@ export interface TgpuResolveOptions {
  * // }
  * ```
  */
-export function resolve(options: TgpuResolveOptions): string {
+export function resolveWithContext(
+  options: TgpuResolveOptions,
+): ResolutionResult {
   const {
     externals,
     template,
@@ -80,12 +85,49 @@ export function resolve(options: TgpuResolveOptions): string {
     toString: () => '<root>',
   };
 
-  const { code } = resolveImpl(resolutionObj, {
+  return resolveImpl(resolutionObj, {
     names: names === 'strict'
       ? new StrictNameRegistry()
       : new RandomNameRegistry(),
     jitTranspiler,
   });
+}
 
-  return code;
+/**
+ * Resolves a template with external values. Each external will get resolved to a code string and replaced in the template.
+ * Any dependencies of the externals will also be resolved and included in the output.
+ * @param options - The options for the resolution.
+ *
+ * @returns The resolved code.
+ *
+ * @example
+ * ```ts
+ * const Gradient = d.struct({
+ *   from: d.vec3f,
+ *   to: d.vec3f,
+ * });
+ *
+ * const resolved = tgpu.resolve({
+ *   template: `
+ *     fn getGradientAngle(gradient: Gradient) -> f32 {
+ *       return atan(gradient.to.y - gradient.from.y, gradient.to.x - gradient.from.x);
+ *     }
+ *   `,
+ *   externals: {
+ *     Gradient,
+ *   },
+ * });
+ *
+ * console.log(resolved);
+ * // struct Gradient_0 {
+ * //   from: vec3f,
+ * //   to: vec3f,
+ * // }
+ * // fn getGradientAngle(gradient: Gradient_0) -> f32 {
+ * //   return atan(gradient.to.y - gradient.from.y, gradient.to.x - gradient.from.x);
+ * // }
+ * ```
+ */
+export function resolve(options: TgpuResolveOptions): string {
+  return resolveWithContext(options).code;
 }

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

@@ -510,8 +510,8 @@ class TgpuRootImpl extends WithBindingImpl
 
       pass.setPipeline(memo.pipeline);
 
-      const missingBindGroups = new Set(memo.bindGroupLayouts);
-      memo.bindGroupLayouts.forEach((layout, idx) => {
+      const missingBindGroups = new Set(memo.usedBindGroupLayouts);
+      memo.usedBindGroupLayouts.forEach((layout, idx) => {
         if (memo.catchall && idx === memo.catchall[0]) {
           // Catch-all
           pass.setBindGroup(idx, this.unwrap(memo.catchall[1]));

packages/typegpu/src/index.ts

@@ -9,7 +9,7 @@ import { computeFn } from './core/function/tgpuComputeFn.ts';
 import { fn } from './core/function/tgpuFn.ts';
 import { fragmentFn } from './core/function/tgpuFragmentFn.ts';
 import { vertexFn } from './core/function/tgpuVertexFn.ts';
-import { resolve } from './core/resolve/tgpuResolve.ts';
+import { resolve, resolveWithContext } from './core/resolve/tgpuResolve.ts';
 import { init, initFromDevice } from './core/root/init.ts';
 import { comparisonSampler, sampler } from './core/sampler/sampler.ts';
 import { accessor } from './core/slot/accessor.ts';
@@ -27,6 +27,7 @@ export const tgpu = {
   initFromDevice,
 
   resolve,
+  resolveWithContext,
 
   '~unstable': {
     fn,

packages/typegpu/src/resolutionCtx.ts

@@ -641,10 +641,17 @@ export class ResolutionCtxImpl implements ResolutionCtx {
   }
 }
 
+/**
+ * The results of a WGSL resolution.
+ *
+ * @param code - The resolved code.
+ * @param usedBindGroupLayouts - List of used `tgpu.bindGroupLayout`s.
+ * @param catchall - Automatically constructed bind group for buffer usages and buffer shorthands, preceded by its index.
+ */
 export interface ResolutionResult {
   code: string;
-  bindGroupLayouts: TgpuBindGroupLayout[];
-  catchall: [number, TgpuBindGroup] | null;
+  usedBindGroupLayouts: TgpuBindGroupLayout[];
+  catchall: [number, TgpuBindGroup] | undefined;
 }
 
 export function resolve(
@@ -655,7 +662,7 @@ export function resolve(
   let code = ctx.resolve(item);
 
   const memoMap = ctx.bindGroupLayoutsToPlaceholderMap;
-  const bindGroupLayouts: TgpuBindGroupLayout[] = [];
+  const usedBindGroupLayouts: TgpuBindGroupLayout[] = [];
   const takenIndices = new Set<number>(
     [...memoMap.keys()]
       .map((layout) => layout.index)
@@ -672,7 +679,7 @@ export function resolve(
   const createCatchallGroup = () => {
     const catchallIdx = automaticIds.next().value;
     const catchallLayout = bindGroupLayout(Object.fromEntries(layoutEntries));
-    bindGroupLayouts[catchallIdx] = catchallLayout;
+    usedBindGroupLayouts[catchallIdx] = catchallLayout;
     code = code.replaceAll(CATCHALL_BIND_GROUP_IDX_MARKER, String(catchallIdx));
 
     return [
@@ -692,17 +699,17 @@ export function resolve(
 
   // Retrieving the catch-all binding index first, because it's inherently
   // the least swapped bind group (fixed and cannot be swapped).
-  const catchall = layoutEntries.length > 0 ? createCatchallGroup() : null;
+  const catchall = layoutEntries.length > 0 ? createCatchallGroup() : undefined;
 
   for (const [layout, placeholder] of memoMap.entries()) {
     const idx = layout.index ?? automaticIds.next().value;
-    bindGroupLayouts[idx] = layout;
+    usedBindGroupLayouts[idx] = layout;
     code = code.replaceAll(placeholder, String(idx));
   }
 
   return {
     code,
-    bindGroupLayouts,
+    usedBindGroupLayouts,
     catchall,
   };
 }