More tests and restrictions

Author: iwoplaza

packages/typegpu/src/core/variable/tgpuVariable.ts

@@ -10,7 +10,6 @@ import type { InferGPU } from '../../shared/repr.ts';
 import {
   $gpuValueOf,
   $internal,
-  $isRef,
   $ownSnippet,
   $resolve,
 } from '../../shared/symbols.ts';
@@ -78,7 +77,6 @@ export function isVariable<T extends TgpuVar>(
 class TgpuVarImpl<TScope extends VariableScope, TDataType extends AnyData>
   implements TgpuVar<TScope, TDataType>, SelfResolvable {
   readonly [$internal] = {};
-  readonly [$isRef]: true;
   readonly #scope: TScope;
   readonly #dataType: TDataType;
   readonly #initialValue: InferGPU<TDataType> | undefined;
@@ -88,7 +86,6 @@ class TgpuVarImpl<TScope extends VariableScope, TDataType extends AnyData>
     dataType: TDataType,
     initialValue?: InferGPU<TDataType> | undefined,
   ) {
-    this[$isRef] = true;
     this.#scope = scope;
     this.#dataType = dataType;
     this.#initialValue = initialValue;

packages/typegpu/src/data/ref.ts

@@ -2,7 +2,7 @@ import { stitch } from '../core/resolve/stitch.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 { $internal, $ownSnippet, $resolve } from '../shared/symbols.ts';
 import type { ResolutionCtx, SelfResolvable } from '../types.ts';
 import { UnknownData } from './dataTypes.ts';
 import type { DualFn } from './dualFn.ts';
@@ -19,9 +19,15 @@ 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 [$isRef]: true;
+  readonly type: 'ref';
 
   /**
    * Derefences the reference, and gives access to the underlying value.
@@ -38,7 +44,6 @@ export interface ref<T> {
   $: 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) => {
     /**
@@ -84,19 +89,23 @@ export const ref: DualFn<<T>(value: T) => ref<T>> = (() => {
   return impl as unknown as DualFn<<T>(value: T) => ref<T>>;
 })();
 
+export function isRef<T>(value: unknown | ref<T>): value is ref<T> {
+  return value instanceof refImpl;
+}
+
 // --------------
 // Implementation
 // --------------
 
 class refImpl<T> implements ref<T> {
-  #value: T;
   readonly [$internal]: true;
-  readonly [$isRef]: true;
+  readonly type: 'ref';
+  #value: T;
 
   constructor(value: T) {
-    this.#value = value;
     this[$internal] = true;
-    this[$isRef] = true;
+    this.type = 'ref';
+    this.#value = value;
   }
 
   get $(): T {

packages/typegpu/src/shared/symbols.ts

@@ -69,16 +69,6 @@ export const $invalidSchemaReason = Symbol(
   `typegpu:${version}:$invalidSchemaReason`,
 );
 
-/**
- * A symbol that identifies objects that act as references
- * (values returned from d.ref(), buffer usages, ...)
- */
-export const $isRef = Symbol(`typegpu:${version}:$isRef`);
-
-export function isRef(value: unknown): value is { [$isRef]: true } {
-  return !!(value as { [$isRef]: true })?.[$isRef];
-}
-
 export function isMarkedInternal(
   value: unknown,
 ): value is { [$internal]: Record<string, unknown> | true } {

packages/typegpu/src/std/array.ts

@@ -2,9 +2,8 @@ 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 type { ref } from '../data/ref.ts';
+import { isRef, type ref } from '../data/ref.ts';
 import { isPtr, isWgslArray, type StorableData } from '../data/wgslTypes.ts';
-import { isRef } from '../shared/symbols.ts';
 
 const sizeOfPointedToArray = (dataType: unknown) =>
   isPtr(dataType) && isWgslArray(dataType.inner)

packages/typegpu/src/tgsl/generationHelpers.ts

@@ -62,7 +62,7 @@ import type { ShelllessRepository } from './shellless.ts';
 import { add, div, mul, sub } from '../std/operators.ts';
 import { $internal } from '../shared/symbols.ts';
 import { stitch } from '../core/resolve/stitch.ts';
-import { derefSnippet } from '../data/ref.ts';
+import { derefSnippet, isRef } from '../data/ref.ts';
 
 type SwizzleableType = 'f' | 'h' | 'i' | 'u' | 'b';
 type SwizzleLength = 1 | 2 | 3 | 4;
@@ -428,6 +428,10 @@ export function coerceToSnippet(value: unknown): Snippet {
     return value;
   }
 
+  if (isRef(value)) {
+    throw new Error('Cannot use refs (d.ref(...)) from the outer scope.');
+  }
+
   // Maybe the value can tell us what snippet it is
   const ownSnippet = getOwnSnippet(value);
   if (ownSnippet) {

packages/typegpu/tests/ref.test.ts

@@ -4,8 +4,20 @@ import { it } from './utils/extendedIt.ts';
 import { asWgsl } from './utils/parseResolved.ts';
 
 describe('ref', () => {
-  it.skip('fails when created outside of a TypeGPU function', () => {
-    expect(() => d.ref(0)).toThrowErrorMatchingInlineSnapshot();
+  it('fails when using a ref as an external', () => {
+    const sup = d.ref(0);
+
+    const foo = () => {
+      'use gpu';
+      sup.$ += 1;
+    };
+
+    expect(() => asWgsl(foo)).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:foo
+      - fn*:foo(): Cannot use refs (d.ref(...)) from the outer scope.]
+    `);
   });
 
   it('creates a regular looking variable in WGSL', () => {

packages/typegpu/tests/tgsl/shellless.test.ts

@@ -186,6 +186,32 @@ describe('shellless', () => {
     `);
   });
 
+  it('generates private pointer params when passing a private variable ref to a function', ({ root }) => {
+    const foo = tgpu.privateVar(d.vec3f);
+
+    const sumComponents = (vec: d.ref<d.v3f>) => {
+      'use gpu';
+      return vec.$.x + vec.$.y + vec.$.z;
+    };
+
+    const main = () => {
+      'use gpu';
+      sumComponents(d.ref(foo.$));
+    };
+
+    expect(asWgsl(main)).toMatchInlineSnapshot(`
+      "fn sumComponents(vec: ptr<private, vec3f>) -> f32 {
+        return (((*vec).x + (*vec).y) + (*vec).z);
+      }
+
+      var<private> foo: vec3f;
+
+      fn main() {
+        sumComponents((&foo));
+      }"
+    `);
+  });
+
   it('generates uniform pointer params when passing a fixed uniform ref to a function', ({ root }) => {
     const posUniform = root.createUniform(d.vec3f);