More useful refs

Author: iwoplaza

apps/typegpu-docs/src/examples/simulation/gravity/compute.ts

@@ -59,7 +59,7 @@ export const computeCollisionsShader = tgpu['~unstable'].computeFn({
         if (isSmaller(currentId, otherId)) {
           const dir = std.normalize(current.position.sub(other.position));
           current.position = other.position.add(
-            dir.mul(radiusOf(current) + radiusOf(other)),
+            dir.mul(radiusOf(d.ref(current)) + radiusOf(d.ref(other))),
           );
         }
 
@@ -118,7 +118,7 @@ export const computeGravityShader = tgpu['~unstable'].computeFn({
       }
 
       const dist = std.max(
-        radiusOf(current) + radiusOf(other),
+        radiusOf(d.ref(current)) + radiusOf(d.ref(other)),
         std.distance(current.position, other.position),
       );
       const gravityForce = (current.mass * other.mass) / dist / dist;

apps/typegpu-docs/src/examples/simulation/gravity/render.ts

@@ -56,7 +56,7 @@ export const mainVertex = tgpu['~unstable'].vertexFn({
   const currentBody = renderLayout.$.celestialBodies[input.instanceIndex];
 
   const worldPosition = currentBody.position.add(
-    input.position.xyz.mul(radiusOf(currentBody)),
+    input.position.xyz.mul(radiusOf(d.ref(currentBody))),
   );
 
   const camera = cameraAccess.$;

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

@@ -44,7 +44,9 @@ export function createShelllessImpl(
     },
 
     toString(): string {
-      return `fn*:${getName(core) ?? '<unnamed>'}`;
+      return `fn*:${getName(core) ?? '<unnamed>'}(${
+        argTypes.map((t) => t.toString()).join(', ')
+      })`;
     },
   };
 }

packages/typegpu/src/data/ptr.ts

@@ -1,4 +1,5 @@
 import { $internal } from '../shared/symbols.ts';
+import { Origin, OriginToPtrParams, originToPtrParams } from './snippet.ts';
 import type { Access, AddressSpace, Ptr, StorableData } from './wgslTypes.ts';
 
 export function ptrFn<T extends StorableData>(
@@ -58,3 +59,20 @@ export function INTERNAL_createPtr<
     toString: () => `ptr<${addressSpace}, ${inner}, ${access}>`,
   } as Ptr<TAddressSpace, TInner, TAccess>;
 }
+
+export function createPtrFromOrigin(
+  origin: Origin,
+  innerDataType: StorableData,
+): Ptr | undefined {
+  const ptrParams = originToPtrParams[origin as keyof OriginToPtrParams];
+
+  if (ptrParams) {
+    return INTERNAL_createPtr(
+      ptrParams.space,
+      innerDataType,
+      ptrParams.access,
+    );
+  }
+
+  return undefined;
+}

packages/typegpu/src/data/ref.ts

@@ -1,20 +1,13 @@
 import { stitch } from '../core/resolve/stitch.ts';
-import { invariant, WgslTypeError } from '../errors.ts';
+import { invariant } from '../errors.ts';
 import { inCodegenMode } from '../execMode.ts';
 import { setName } from '../shared/meta.ts';
 import { $internal, $isRef, $ownSnippet, $resolve } from '../shared/symbols.ts';
 import type { ResolutionCtx, SelfResolvable } from '../types.ts';
 import { UnknownData } from './dataTypes.ts';
 import type { DualFn } from './dualFn.ts';
-import { INTERNAL_createPtr } from './ptr.ts';
-import {
-  isEphemeralSnippet,
-  type OriginToPtrParams,
-  originToPtrParams,
-  type ResolvedSnippet,
-  snip,
-  type Snippet,
-} from './snippet.ts';
+import { createPtrFromOrigin } from './ptr.ts';
+import { type ResolvedSnippet, snip, type Snippet } from './snippet.ts';
 import {
   isNaturallyEphemeral,
   isPtr,
@@ -29,17 +22,25 @@ import {
 export interface ref<T> {
   readonly [$internal]: unknown;
   readonly [$isRef]: true;
+
+  /**
+   * Derefences the reference, and gives access to the underlying value.
+   *
+   * @example ```ts
+   * const boid = Boid({ pos: d.vec3f(3, 2, 1) });
+   * const posRef = d.ref(boid.pos);
+   *
+   * // Actually updates `boid.pos`
+   * posRef.$ = d.vec3f(1, 2, 3);
+   * console.log(boid.pos); // Output: vec3f(1, 2, 3)
+   * ```
+   */
   $: T;
 }
 
 // TODO: Restrict calls to this function only from within TypeGPU functions
 export const ref: DualFn<<T>(value: T) => ref<T>> = (() => {
   const gpuImpl = (value: Snippet) => {
-    if (!isEphemeralSnippet(value)) {
-      throw new WgslTypeError(
-        `Can't create refs from references. Copy the value first by wrapping it in its schema.`,
-      );
-    }
     return snip(new RefOnGPU(value), UnknownData, /* origin */ 'runtime');
   };
 
@@ -71,11 +72,11 @@ export const ref: DualFn<<T>(value: T) => ref<T>> = (() => {
 // --------------
 
 class refImpl<T> implements ref<T> {
-  readonly #value: T | string;
+  #value: T;
   readonly [$internal]: true;
   readonly [$isRef]: true;
 
-  constructor(value: T | string) {
+  constructor(value: T) {
     this.#value = value;
     this[$internal] = true;
     this[$isRef] = true;
@@ -84,20 +85,49 @@ class refImpl<T> implements ref<T> {
   get $(): T {
     return this.#value as T;
   }
+
+  set $(value: T) {
+    if (value && typeof value === 'object') {
+      // Setting an object means updating the properties of the original object.
+      // e.g.: foo.$ = Boid();
+      for (const key of Object.keys(value) as (keyof T)[]) {
+        this.#value[key] = value[key];
+      }
+    } else {
+      this.#value = value;
+    }
+  }
 }
 
 export class RefOnGPU {
-  readonly snippet: Snippet;
   readonly [$internal]: true;
 
+  readonly snippet: Snippet;
+  /**
+   * Pointer params only exist if the ref was created from a reference (buttery-butter).
+   */
+  readonly ptrType: Ptr | undefined;
+
   constructor(snippet: Snippet) {
-    this.snippet = snippet;
     this[$internal] = true;
+    this.snippet = snippet;
+    this.ptrType = createPtrFromOrigin(
+      snippet.origin,
+      snippet.dataType as StorableData,
+    );
   }
 
   toString(): string {
     return `ref:${this.snippet.value}`;
   }
+
+  [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
+    invariant(
+      !!this.ptrType,
+      'RefOnGPU must have a pointer type when resolved',
+    );
+    return snip(stitch`(&${this.snippet})`, this.ptrType, this.snippet.origin);
+  }
 }
 
 export class RefOperator implements SelfResolvable {
@@ -109,20 +139,18 @@ export class RefOperator implements SelfResolvable {
     this[$internal] = true;
     this.snippet = snippet;
 
-    const ptrParams =
-      originToPtrParams[this.snippet.origin as keyof OriginToPtrParams];
+    const ptrType = createPtrFromOrigin(
+      snippet.origin,
+      snippet.dataType as StorableData,
+    );
 
-    if (!ptrParams) {
+    if (!ptrType) {
       throw new Error(
         `Cannot take a reference of a value with origin ${this.snippet.origin}`,
       );
     }
 
-    this.#ptrType = INTERNAL_createPtr(
-      ptrParams.space,
-      this.snippet.dataType as StorableData,
-      ptrParams.access,
-    );
+    this.#ptrType = ptrType;
   }
 
   get [$ownSnippet](): Snippet {

packages/typegpu/src/tgsl/shellless.ts

@@ -3,8 +3,10 @@ import {
   type ShelllessImpl,
 } from '../core/function/shelllessImpl.ts';
 import type { AnyData } from '../data/dataTypes.ts';
+import { RefOnGPU } from '../data/ref.ts';
 import type { Snippet } from '../data/snippet.ts';
 import { isPtr } from '../data/wgslTypes.ts';
+import { WgslTypeError } from '../errors.ts';
 import { getResolutionCtx } from '../execMode.ts';
 import { getMetaData, getName } from '../shared/meta.ts';
 import { concretize } from './generationHelpers.ts';
@@ -47,7 +49,22 @@ export class ShelllessRepository {
       );
     }
 
-    const argTypes = (argSnippets ?? []).map((s) => {
+    const argTypes = (argSnippets ?? []).map((s, index) => {
+      if (s.value instanceof RefOnGPU) {
+        if (!s.value.ptrType) {
+          throw new WgslTypeError(
+            `d.ref() created with primitive types must be stored in a variable before use`,
+          );
+        }
+        return s.value.ptrType;
+      }
+
+      if (s.dataType.type === 'unknown') {
+        throw new Error(
+          `Passed illegal value ${s.value} as the #${index} argument to ${meta.name}(...)`,
+        );
+      }
+
       let type = concretize(s.dataType as AnyData);
 
       if (s.origin === 'constant-ref') {

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -782,6 +782,11 @@ ${this.ctx.pre}else ${alternate}`;
 
       if (eq.value instanceof RefOnGPU) {
         // We're assigning a newly created `d.ref()`
+        if (eq.value.ptrType) {
+          throw new WgslTypeError(
+            `Cannot store d.ref() in a variable if it references another value. Copy the value passed into d.ref() instead.`,
+          );
+        }
         const refSnippet = eq.value.snippet;
         const varName = this.refVariable(
           rawId,

packages/typegpu/tests/constant.test.ts

@@ -64,7 +64,7 @@ describe('tgpu.const', () => {
       [Error: Resolution of the following tree failed:
       - <root>
       - fn*:fn2
-      - fn*:fn2: Cannot pass constant references as function arguments. Explicitly copy them by wrapping them in a schema: 'vec3f(...)']
+      - fn*:fn2(): Cannot pass constant references as function arguments. Explicitly copy them by wrapping them in a schema: 'vec3f(...)']
     `);
   });
 
@@ -84,7 +84,7 @@ describe('tgpu.const', () => {
       [Error: Resolution of the following tree failed:
       - <root>
       - fn*:fn
-      - fn*:fn: 'boid.pos = vec3f()' is invalid, because boid.pos is a constant.]
+      - fn*:fn(): 'boid.pos = vec3f()' is invalid, because boid.pos is a constant.]
     `);
 
     // Since we freeze the object, we cannot mutate when running the function in JS either