feat: Dollarless de-ref (#2119)

Author: iwoplaza

packages/typegpu/src/data/index.ts

@@ -191,7 +191,7 @@ export { unstruct } from './unstruct.ts';
 export { mat2x2f, mat3x3f, mat4x4f, matToArray } from './matrix.ts';
 export * from './vertexFormatData.ts';
 export { atomic } from './atomic.ts';
-export { ref } from './ref.ts';
+export { _ref as ref } from './ref.ts';
 export {
   align,
   type AnyAttribute,

packages/typegpu/src/data/ref.ts

@@ -22,15 +22,8 @@ import {
 // Public API
 // ----------
 
-/**
- * A reference to a value `T`. Can be passed to other functions to give them
- * mutable access to the underlying value.
- *
- * Conceptually, it represents a WGSL pointer.
- */
-export interface ref<T> {
-  readonly [$internal]: unknown;
-  readonly type: 'ref';
+interface ref<T> {
+  readonly [$internal]: { type: 'ref' };
 
   /**
    * Derefences the reference, and gives access to the underlying value.
@@ -47,10 +40,18 @@ export interface ref<T> {
   $: T;
 }
 
-type RefFn = DualFn<(<T>(value: T) => ref<T>)> & { [$internal]: true };
+/**
+ * A reference to a value `T`. Can be passed to other functions to give them
+ * mutable access to the underlying value.
+ *
+ * Conceptually, it represents a WGSL pointer.
+ */
+export type _ref<T> = T extends object ? T & ref<T> : ref<T>;
+
+type RefFn = DualFn<(<T>(value: T) => _ref<T>)> & { [$internal]: true };
 
-export const ref = (() => {
-  const impl = (<T>(value: T) => new refImpl(value)) as unknown as RefFn;
+export const _ref = (() => {
+  const impl = (<T>(value: T) => INTERNAL_createRef(value)) as unknown as RefFn;
 
   setName(impl, 'ref');
   impl.toString = () => 'ref';
@@ -94,39 +95,60 @@ export const ref = (() => {
 })();
 
 export function isRef<T>(value: unknown | ref<T>): value is ref<T> {
-  return value instanceof refImpl;
+  return (value as ref<T>)?.[$internal]?.type === 'ref';
 }
 
 // --------------
 // Implementation
 // --------------
 
-class refImpl<T> implements ref<T> {
-  readonly [$internal]: true;
-  readonly type: 'ref';
-  #value: T;
+export function INTERNAL_createRef<T>(value: T): ref<T> {
+  const target = {
+    [$internal]: { type: 'ref' },
 
-  constructor(value: T) {
-    this[$internal] = true;
-    this.type = 'ref';
-    this.#value = value;
-  }
+    get $(): T {
+      return value;
+    },
+
+    set $(newValue: T) {
+      if (newValue && typeof newValue === 'object') {
+        // Setting an object means updating the properties of the original object.
+        // e.g.: foo.$ = Boid();
+        for (const key of Object.keys(newValue) as (keyof T)[]) {
+          value[key] = newValue[key];
+        }
+      } else {
+        value = newValue;
+      }
+    },
+  };
 
-  get $(): T {
-    return this.#value as T;
+  if (value === undefined || value === null) {
+    throw new Error('Cannot create a ref from undefined or null');
   }
 
-  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;
-    }
+  if (typeof value === 'object') {
+    return new Proxy(target, {
+      get(target, prop) {
+        if (prop in target) {
+          return target[prop as keyof typeof target];
+        }
+        return value[prop as keyof T];
+      },
+      set(_target, prop, propValue) {
+        if (prop === $internal) {
+          return false;
+        }
+        if (prop === '$') {
+          console.log('Setting ref value:', propValue);
+          return Reflect.set(target, prop, propValue);
+        }
+        return Reflect.set(value as object, prop, propValue);
+      },
+    }) as ref<T>;
   }
+
+  return target as ref<T>;
 }
 
 /**

packages/typegpu/src/data/wgslTypes.ts

@@ -31,7 +31,7 @@ import type {
   WgslTexture,
 } from './texture.ts';
 import type { WgslComparisonSampler, WgslSampler } from './sampler.ts';
-import type { ref } from './ref.ts';
+import type { _ref as ref } from './ref.ts';
 import type { DualFn } from '../types.ts';
 
 type DecoratedLocation<T extends BaseData> = Decorated<T, Location[]>;

packages/typegpu/src/std/array.ts

@@ -2,7 +2,7 @@ import { dualImpl } from '../core/function/dualImpl.ts';
 import { stitch } from '../core/resolve/stitch.ts';
 import { abstractInt, u32 } from '../data/numeric.ts';
 import { ptrFn } from '../data/ptr.ts';
-import { isRef, type ref } from '../data/ref.ts';
+import { type _ref as ref, isRef } from '../data/ref.ts';
 import { isPtr, isWgslArray, type StorableData } from '../data/wgslTypes.ts';
 
 const sizeOfPointedToArray = (dataType: unknown) =>

packages/typegpu/tests/ref.test.ts

@@ -108,6 +108,75 @@ describe('d.ref', () => {
     `);
   });
 
+  it('allows updating a struct property from another function', () => {
+    type Entity = d.Infer<typeof Entity>;
+    const Entity = d.struct({ pos: d.vec3f });
+
+    const clearPosition = (entity: d.ref<Entity>) => {
+      'use gpu';
+      entity.pos = d.vec3f();
+    };
+
+    const main = () => {
+      'use gpu';
+      const entity = Entity({ pos: d.vec3f(1, 2, 3) });
+      clearPosition(d.ref(entity));
+      // entity.pos should be vec3f(0, 0, 0)
+      return entity;
+    };
+
+    // Works in JS
+    expect(main().pos).toStrictEqual(d.vec3f(0, 0, 0));
+
+    // And on the GPU
+    expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
+      "struct Entity {
+        pos: vec3f,
+      }
+
+      fn clearPosition(entity: ptr<function, Entity>) {
+        (*entity).pos = vec3f();
+      }
+
+      fn main() -> Entity {
+        var entity = Entity(vec3f(1, 2, 3));
+        clearPosition((&entity));
+        return entity;
+      }"
+    `);
+  });
+
+  it('allows updating a vector component from another function', () => {
+    const clearX = (pos: d.ref<d.v3f>) => {
+      'use gpu';
+      pos.x = 0;
+    };
+
+    const main = () => {
+      'use gpu';
+      const pos = d.vec3f(1, 0, 0);
+      clearX(d.ref(pos));
+      // pos should be vec3f(0, 0, 0)
+      return pos;
+    };
+
+    // Works in JS
+    expect(main()).toStrictEqual(d.vec3f(0, 0, 0));
+
+    // And on the GPU
+    expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
+      "fn clearX(pos: ptr<function, vec3f>) {
+        (*pos).x = 0f;
+      }
+
+      fn main() -> vec3f {
+        var pos = vec3f(1, 0, 0);
+        clearX((&pos));
+        return pos;
+      }"
+    `);
+  });
+
   it('allows updating a number from another function', () => {
     const increment = (value: d.ref<number>) => {
       'use gpu';