Working on umiform refs

Author: iwoplaza

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

@@ -37,15 +37,15 @@ export const computeCollisionsShader = tgpu['~unstable'].computeFn({
       otherId++
     ) {
       const other = computeLayout.$.inState[otherId];
-      // no collision occurs...
       if (
         otherId === currentId || // ...with itself
         other.destroyed === 1 || // ...when other is destroyed
         current.collisionBehavior === none || // ...when current behavior is none
         other.collisionBehavior === none || // ...when other behavior is none
         std.distance(current.position, other.position) >=
-          radiusOf(current) + radiusOf(other) // ...when other is too far away
+          radiusOf(d.ref(current)) + radiusOf(d.ref(other)) // ...when other is too far away
       ) {
+        // no collision occurs...
         continue;
       }
 
@@ -58,8 +58,9 @@ export const computeCollisionsShader = tgpu['~unstable'].computeFn({
         // push the smaller object outside
         if (isSmaller(currentId, otherId)) {
           const dir = std.normalize(current.position.sub(other.position));
-          current.position = other.position
-            .add(dir.mul(radiusOf(current) + radiusOf(other)));
+          current.position = other.position.add(
+            dir.mul(radiusOf(current) + radiusOf(other)),
+          );
         }
 
         // bounce with tiny damping

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

@@ -3,11 +3,7 @@ import { OBJLoader } from '@loaders.gl/obj';
 import type { TgpuRoot } from 'typegpu';
 import * as d from 'typegpu/data';
 import { sphereTextureNames } from './enums.ts';
-import {
-  type CelestialBody,
-  renderVertexLayout,
-  SkyBoxVertex,
-} from './schemas.ts';
+import { CelestialBody, renderVertexLayout, SkyBoxVertex } from './schemas.ts';
 
 function vert(
   position: [number, number, number],
@@ -166,7 +162,7 @@ export async function loadSphereTextures(root: TgpuRoot) {
   return texture;
 }
 
-export const radiusOf = (body: CelestialBody): number => {
+export const radiusOf = (body: d.ref<CelestialBody>): number => {
   'use gpu';
-  return (((body.mass * 0.75) / Math.PI) ** 0.333) * body.radiusMultiplier;
+  return (((body.$.mass * 0.75) / Math.PI) ** 0.333) * body.$.radiusMultiplier;
 };

packages/typegpu/tests/ref.test.ts

@@ -43,4 +43,20 @@ describe('ref', () => {
       - fn*:hello: Cannot assign a ref to an existing variable '(&foo)', define a new variable instead.]
     `);
   });
+
+  it('fails when creating a ref with a reference', () => {
+    const hello = () => {
+      'use gpu';
+      const position = d.vec3f(1, 2, 3);
+      const foo = d.ref(position);
+    };
+
+    expect(() => asWgsl(hello)).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:hello
+      - fn*:hello
+      - fn:ref: Can't create refs from references. Copy the value first by wrapping it in its schema.]
+    `);
+  });
 });

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

@@ -166,10 +166,11 @@ describe('shellless', () => {
       pos.$.z += vel.z;
     };
 
-    const main = tgpu.fn([])(() => {
+    const main = () => {
+      'use gpu';
       const pos = d.ref(d.vec3f(0, 0, 0));
       advance(pos, d.vec3f(1, 2, 3));
-    });
+    };
 
     expect(asWgsl(main)).toMatchInlineSnapshot(`
       "fn advance(pos: ptr<function, vec3f>, vel: vec3f) {
@@ -185,6 +186,36 @@ describe('shellless', () => {
     `);
   });
 
+  it('generates uniform pointer params when passing a fixed uniform directly to a function', ({ root }) => {
+    const posUniform = root.createUniform(d.vec3f);
+
+    const sumComponents = (vec: d.ref<d.v3f>) => {
+      'use gpu';
+      return vec.$.x + vec.$.y + vec.$.z;
+    };
+
+    const main = () => {
+      'use gpu';
+      sumComponents(posUniform);
+      sumComponents(d.ref(posUniform.$.zyx));
+      // sumComponents(&posUniform);
+    };
+
+    expect(asWgsl(main)).toMatchInlineSnapshot(`
+      "fn advance(pos: ptr<function, vec3f>, vel: vec3f) {
+        (*pos).x += vel.x;
+        (*pos).y += vel.y;
+        (*pos).z += vel.z;
+      }
+
+      fn main() {
+        var pos = vec3f();
+        var vel = vec3f(1, 2, 3);
+        advance((&pos), (&vel));
+      }"
+    `);
+  });
+
   it('resolves when accepting no arguments', () => {
     const main = () => {
       'use gpu';