impr: Unify bufferShorthand and bufferUsage

apps/typegpu-docs/src/examples/simulation/fluid-double-buffering/index.ts

@@ -1,5 +1,5 @@
 import { randf } from '@typegpu/noise';
-import tgpu, { d, std, type TgpuBufferMutable, type TgpuBufferReadonly } from 'typegpu';
+import tgpu, { d, std, type TgpuMutable, type TgpuReadonly } from 'typegpu';
 import { defineControls } from '../../common/defineControls.ts';
 
 const MAX_GRID_SIZE = 1024;
@@ -22,8 +22,8 @@ const BoxObstacle = d.struct({
 
 const gridSize = 256;
 
-const inputGridSlot = tgpu.slot<TgpuBufferReadonly<GridData> | TgpuBufferMutable<GridData>>();
-const outputGridSlot = tgpu.slot<TgpuBufferMutable<GridData>>();
+const inputGridSlot = tgpu.slot<TgpuReadonly<GridData> | TgpuMutable<GridData>>();
+const outputGridSlot = tgpu.slot<TgpuMutable<GridData>>();
 
 const MAX_OBSTACLES = 4;
 const BoxObstacleArray = d.arrayOf(BoxObstacle, MAX_OBSTACLES);
@@ -409,8 +409,8 @@ const canvas = document.querySelector('canvas') as HTMLCanvasElement;
 const context = root.configureContext({ canvas, alphaMode: 'premultiplied' });
 
 function makePipelines(
-  inputGridReadonly: TgpuBufferReadonly<GridData>,
-  outputGridMutable: TgpuBufferMutable<GridData>,
+  inputGridReadonly: TgpuReadonly<GridData>,
+  outputGridMutable: TgpuMutable<GridData>,
 ) {
   const initWorldPipeline = root
     .with(outputGridSlot, outputGridMutable)

packages/typegpu/src/core/buffer/buffer.ts

@@ -22,17 +22,16 @@ import { $internal } from '../../shared/symbols.ts';
 import type { Prettify, UnionToIntersection } from '../../shared/utilityTypes.ts';
 import { isGPUBuffer } from '../../types.ts';
 import type { ExperimentalTgpuRoot } from '../root/rootTypes.ts';
+import { calculateOffsets, readFromArrayBuffer, writeToArrayBuffer } from '../../data/dataIO.ts';
+import { patchArrayBuffer } from '../../data/partialIO.ts';
 import {
   mutable,
   readonly,
-  type TgpuBufferMutable,
-  type TgpuBufferReadonly,
-  type TgpuBufferUniform,
-  type TgpuFixedBufferUsage,
   uniform,
-} from './bufferUsage.ts';
-import { calculateOffsets, readFromArrayBuffer, writeToArrayBuffer } from '../../data/dataIO.ts';
-import { patchArrayBuffer } from '../../data/partialIO.ts';
+  type TgpuMutable,
+  type TgpuReadonly,
+  type TgpuUniform,
+} from './bufferShorthand.ts';
 
 // ----------
 // Public API
@@ -82,10 +81,10 @@ type ViewUsages<TBuffer extends TgpuBuffer<BaseData>> =
   | (boolean extends TBuffer['usableAsUniform'] ? never : 'uniform')
   | (boolean extends TBuffer['usableAsStorage'] ? never : 'readonly' | 'mutable');
 
-type UsageTypeToBufferUsage<TData extends BaseData> = {
-  uniform: TgpuBufferUniform<TData> & TgpuFixedBufferUsage<TData>;
-  mutable: TgpuBufferMutable<TData> & TgpuFixedBufferUsage<TData>;
-  readonly: TgpuBufferReadonly<TData> & TgpuFixedBufferUsage<TData>;
+type UsageTypeToBufferShorthand<TData extends BaseData> = {
+  uniform: TgpuUniform<TData>;
+  mutable: TgpuMutable<TData>;
+  readonly: TgpuReadonly<TData>;
 };
 
 const usageToUsageConstructor = { uniform, mutable, readonly };
@@ -141,7 +140,7 @@ export interface TgpuBuffer<TData extends BaseData> extends TgpuNamable {
   ): this & UnionToIntersection<LiteralToUsageType<T[number]>>;
   $addFlags(flags: GPUBufferUsageFlags): this;
 
-  as<T extends ViewUsages<this>>(usage: T): UsageTypeToBufferUsage<TData>[T];
+  as<T extends ViewUsages<this>>(usage: T): UsageTypeToBufferShorthand<TData>[T];
 
   compileWriter(): void;
   write(data: InferInput<TData>, options?: BufferWriteOptions): void;
@@ -449,8 +448,8 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
     return res;
   }
 
-  as<T extends ViewUsages<this>>(usage: T): UsageTypeToBufferUsage<TData>[T] {
-    return usageToUsageConstructor[usage]?.(this as never) as UsageTypeToBufferUsage<TData>[T];
+  as<T extends ViewUsages<this>>(usage: T): UsageTypeToBufferShorthand<TData>[T] {
+    return usageToUsageConstructor[usage](this as never) as UsageTypeToBufferShorthand<TData>[T];
   }
 
   destroy() {

packages/typegpu/src/core/buffer/bufferShorthand.ts

@@ -1,12 +1,24 @@
-import type { ResolvedSnippet } from '../../data/snippet.ts';
-import type { BaseData } from '../../data/wgslTypes.ts';
-import type { StorageFlag } from '../../extension.ts';
+import { schemaCallWrapper } from '../../data/schemaCallWrapper.ts';
+import { snip, type ResolvedSnippet } from '../../data/snippet.ts';
+import type { AnyWgslData, BaseData } from '../../data/wgslTypes.ts';
+import { IllegalBufferAccessError } from '../../errors.ts';
+import { getExecMode, isInsideTgpuFn } from '../../execMode.ts';
+import { isUsableAsStorage, type StorageFlag } from '../../extension.ts';
 import { getName, setName, type TgpuNamable } from '../../shared/meta.ts';
 import type { Infer, InferGPU, InferInput, InferPatch, InferPartial } from '../../shared/repr.ts';
-import { $getNameForward, $gpuValueOf, $internal, $resolve } from '../../shared/symbols.ts';
+import {
+  $getNameForward,
+  $gpuValueOf,
+  $internal,
+  $ownSnippet,
+  $repr,
+  $resolve,
+} from '../../shared/symbols.ts';
+import { assertExhaustive } from '../../shared/utilityTypes.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
+import { valueProxyHandler } from '../valueProxyUtils.ts';
 import type { BufferWriteOptions, TgpuBuffer, UniformFlag } from './buffer.ts';
-import type { TgpuBufferUsage } from './bufferUsage.ts';
+import { isUsableAsUniform } from './bufferUsage.ts';
 
 // ----------
 // Public API
@@ -26,6 +38,9 @@ interface TgpuBufferShorthandBase<TData extends BaseData> extends TgpuNamable {
   // Accessible on the GPU
   readonly [$gpuValueOf]: InferGPU<TData>;
   // ---
+
+  /** Type-token, not available at runtime */
+  readonly [$repr]: Infer<TData>;
 }
 
 export interface TgpuMutable<out TData extends BaseData> extends TgpuBufferShorthandBase<TData> {
@@ -39,6 +54,9 @@ export interface TgpuMutable<out TData extends BaseData> extends TgpuBufferShort
   value: InferGPU<TData>;
   $: InferGPU<TData>;
   // ---
+
+  /** Type-token, not available at runtime */
+  readonly [$repr]: Infer<TData>;
 }
 
 export interface TgpuReadonly<out TData extends BaseData> extends TgpuBufferShorthandBase<TData> {
@@ -52,6 +70,9 @@ export interface TgpuReadonly<out TData extends BaseData> extends TgpuBufferShor
   readonly value: InferGPU<TData>;
   readonly $: InferGPU<TData>;
   // ---
+
+  /** Type-token, not available at runtime */
+  readonly [$repr]: Infer<TData>;
 }
 
 export interface TgpuUniform<out TData extends BaseData> extends TgpuBufferShorthandBase<TData> {
@@ -86,23 +107,22 @@ export class TgpuBufferShorthandImpl<
   TType extends 'mutable' | 'readonly' | 'uniform',
   TData extends BaseData,
 > implements SelfResolvable {
+  /** Type-token, not available at runtime */
+  declare readonly [$repr]: Infer<TData>;
+
   readonly [$internal] = true;
   readonly [$getNameForward]: object;
   readonly resourceType: TType;
   readonly buffer: TgpuBuffer<TData> &
     (TType extends 'mutable' | 'readonly' ? StorageFlag : UniformFlag);
 
-  readonly #usage: TgpuBufferUsage<TData, TType>;
-
   constructor(
-    resourceType: TType,
+    usage: TType,
     buffer: TgpuBuffer<TData> & (TType extends 'mutable' | 'readonly' ? StorageFlag : UniformFlag),
   ) {
-    this.resourceType = resourceType;
+    this.resourceType = usage;
     this.buffer = buffer;
     this[$getNameForward] = buffer;
-    // oxlint-disable-next-line typescript/no-explicit-any -- too complex a type
-    this.#usage = (this.buffer as any).as(this.resourceType);
   }
 
   $name(label: string): this {
@@ -128,22 +148,178 @@ export class TgpuBufferShorthandImpl<
   }
 
   get [$gpuValueOf](): InferGPU<TData> {
-    return this.#usage.$;
+    const dataType = this.buffer.dataType;
+    const usage = this.resourceType;
+
+    return new Proxy(
+      {
+        [$internal]: true,
+        get [$ownSnippet]() {
+          return snip(this, dataType, usage);
+        },
+        [$resolve]: (ctx) => ctx.resolve(this),
+        toString: () => `${usage}:${getName(this) ?? '<unnamed>'}.$`,
+      },
+      valueProxyHandler,
+    ) as InferGPU<TData>;
   }
 
   get $(): InferGPU<TData> {
-    return this.#usage.$;
+    const mode = getExecMode();
+    const insideTgpuFn = isInsideTgpuFn();
+
+    if (mode.type === 'normal') {
+      throw new IllegalBufferAccessError(
+        insideTgpuFn
+          ? `Cannot access ${String(
+              this.buffer,
+            )}. TypeGPU functions that depends on GPU resources need to be part of a compute dispatch, draw call or simulation`
+          : '.$ is inaccessible during normal JS execution. Try `.read()`',
+      );
+    }
+
+    if (mode.type === 'codegen') {
+      return this[$gpuValueOf];
+    }
+
+    if (mode.type === 'simulate') {
+      if (!mode.buffers.has(this.buffer)) {
+        // Not initialized yet
+        mode.buffers.set(this.buffer, schemaCallWrapper(this.buffer.dataType, this.buffer.initial));
+      }
+      return mode.buffers.get(this.buffer) as InferGPU<TData>;
+    }
+
+    return assertExhaustive(mode, 'bufferShorthand.ts#TgpuBufferShorthandImpl/$');
+  }
+
+  set $(value: InferGPU<TData>) {
+    const mode = getExecMode();
+    const insideTgpuFn = isInsideTgpuFn();
+
+    if (mode.type === 'normal') {
+      throw new IllegalBufferAccessError(
+        insideTgpuFn
+          ? `Cannot access ${String(
+              this.buffer,
+            )}. TypeGPU functions that depends on GPU resources need to be part of a compute dispatch, draw call or simulation`
+          : '.$ is inaccessible during normal JS execution. Try `.write()`',
+      );
+    }
+
+    if (mode.type === 'codegen') {
+      // The WGSL generator handles buffer assignment, and does not defer to
+      // whatever's being assigned to generate the WGSL.
+      throw new Error('Unreachable bufferShorthand.ts#TgpuBufferShorthandImpl/$');
+    }
+
+    if (mode.type === 'simulate') {
+      mode.buffers.set(this.buffer, value);
+      return;
+    }
+
+    assertExhaustive(mode, 'bufferShorthand.ts#TgpuBufferShorthandImpl/$');
   }
 
   get value(): InferGPU<TData> {
     return this.$;
   }
 
+  set value(value: InferGPU<TData>) {
+    this.$ = value;
+  }
+
   toString(): string {
     return `${this.resourceType}BufferShorthand:${getName(this) ?? '<unnamed>'}`;
   }
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
-    return ctx.resolve(this.#usage);
+    const dataType = this.buffer.dataType;
+    const id = ctx.makeUniqueIdentifier(getName(this), 'global');
+    const { group, binding } = ctx.allocateFixedEntry(
+      this.resourceType === 'uniform'
+        ? { uniform: dataType }
+        : { storage: dataType, access: this.resourceType },
+      this.buffer,
+    );
+
+    return ctx.gen.declareGlobalVar({
+      group,
+      binding,
+      scope: this.resourceType,
+      id,
+      dataType,
+      init: undefined,
+    });
+  }
+}
+
+// --------------
+// Constructors
+// --------------
+
+const mutableUsageMap = new WeakMap<
+  TgpuBuffer<BaseData>,
+  TgpuBufferShorthandImpl<'mutable', BaseData>
+>();
+
+export function mutable<TData extends AnyWgslData>(
+  buffer: TgpuBuffer<TData> & StorageFlag,
+): TgpuMutable<TData> {
+  if (!isUsableAsStorage(buffer)) {
+    throw new Error(
+      `Cannot call as('mutable') on ${buffer}, as it is not allowed to be used as storage. To allow it, call .$usage('storage') when creating the buffer.`,
+    );
+  }
+
+  let usage = mutableUsageMap.get(buffer);
+  if (!usage) {
+    usage = new TgpuBufferShorthandImpl('mutable', buffer);
+    mutableUsageMap.set(buffer, usage);
+  }
+  return usage as unknown as TgpuMutable<TData>;
+}
+
+const readonlyUsageMap = new WeakMap<
+  TgpuBuffer<BaseData>,
+  TgpuBufferShorthandImpl<'readonly', BaseData>
+>();
+
+export function readonly<TData extends AnyWgslData>(
+  buffer: TgpuBuffer<TData> & StorageFlag,
+): TgpuReadonly<TData> {
+  if (!isUsableAsStorage(buffer)) {
+    throw new Error(
+      `Cannot call as('readonly') on ${buffer}, as it is not allowed to be used as storage. To allow it, call .$usage('storage') when creating the buffer.`,
+    );
+  }
+
+  let usage = readonlyUsageMap.get(buffer);
+  if (!usage) {
+    usage = new TgpuBufferShorthandImpl('readonly', buffer);
+    readonlyUsageMap.set(buffer, usage);
+  }
+  return usage as unknown as TgpuReadonly<TData>;
+}
+
+const uniformUsageMap = new WeakMap<
+  TgpuBuffer<BaseData>,
+  TgpuBufferShorthandImpl<'uniform', BaseData>
+>();
+
+export function uniform<TData extends AnyWgslData>(
+  buffer: TgpuBuffer<TData> & UniformFlag,
+): TgpuUniform<BaseData> {
+  if (!isUsableAsUniform(buffer)) {
+    throw new Error(
+      `Cannot call as('uniform') on ${buffer}, as it is not allowed to be used as a uniform. To allow it, call .$usage('uniform') when creating the buffer.`,
+    );
+  }
+
+  let usage = uniformUsageMap.get(buffer);
+  if (!usage) {
+    usage = new TgpuBufferShorthandImpl('uniform', buffer);
+    uniformUsageMap.set(buffer, usage);
   }
+  return usage as unknown as TgpuUniform<TData>;
 }