Fix for referencing implicit pointers

Author: iwoplaza

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

@@ -43,7 +43,7 @@ export const computeCollisionsShader = tgpu['~unstable'].computeFn({
         current.collisionBehavior === none || // ...when current behavior is none
         other.collisionBehavior === none || // ...when other behavior is none
         std.distance(current.position, other.position) >=
-          radiusOf(d.ref(current)) + radiusOf(d.ref(other)) // ...when other is too far away
+          radiusOf(current) + radiusOf(other) // ...when other is too far away
       ) {
         // no collision occurs...
         continue;
@@ -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(d.ref(current)) + radiusOf(d.ref(other))),
+            dir.mul(radiusOf(current) + radiusOf(other)),
           );
         }
 
@@ -118,7 +118,7 @@ export const computeGravityShader = tgpu['~unstable'].computeFn({
       }
 
       const dist = std.max(
-        radiusOf(d.ref(current)) + radiusOf(d.ref(other)),
+        radiusOf(current) + radiusOf(other),
         std.distance(current.position, other.position),
       );
       const gravityForce = (current.mass * other.mass) / dist / dist;

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

@@ -166,7 +166,7 @@ export async function loadSphereTextures(root: TgpuRoot) {
   return texture;
 }
 
-export const radiusOf = (body: d.ref<CelestialBody>): number => {
+export const radiusOf = (body: 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;
 };

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(d.ref(currentBody))),
+    input.position.xyz.mul(radiusOf(currentBody)),
   );
 
   const camera = cameraAccess.$;

apps/typegpu-docs/src/examples/tests/tgsl-parsing-test/pointers.ts

@@ -16,11 +16,6 @@ const modifyStructFn = tgpu.fn([d.ptrFn(SimpleStruct)])((ptr) => {
   ptr.$.vec.x += 1;
 });
 
-const privateNum = tgpu.privateVar(d.u32);
-const modifyNumPrivate = tgpu.fn([d.ptrPrivate(d.u32)])((ptr) => {
-  ptr.$ += 1;
-});
-
 const privateVec = tgpu.privateVar(d.vec2f);
 const modifyVecPrivate = tgpu.fn([d.ptrPrivate(d.vec2f)])((ptr) => {
   ptr.$.x += 1;
@@ -49,9 +44,6 @@ export const pointersTest = tgpu.fn([], d.bool)(() => {
   s = s && std.allEq(myStruct.$.vec, d.vec2f(1, 0));
 
   // private pointers
-  modifyNumPrivate(d.ref(privateNum.$));
-  s = s && (privateNum.$ === 1);
-
   modifyVecPrivate(d.ref(privateVec.$));
   s = s && std.allEq(privateVec.$, d.vec2f(1, 0));
 

packages/typegpu/src/data/ptr.ts

@@ -80,3 +80,21 @@ export function createPtrFromOrigin(
 
   return undefined;
 }
+
+export function implicitFrom(ptr: Ptr): Ptr {
+  return INTERNAL_createPtr(
+    ptr.addressSpace,
+    ptr.inner,
+    ptr.access,
+    /* implicit */ true,
+  );
+}
+
+export function explicitFrom(ptr: Ptr): Ptr {
+  return INTERNAL_createPtr(
+    ptr.addressSpace,
+    ptr.inner,
+    ptr.access,
+    /* implicit */ false,
+  );
+}

packages/typegpu/src/data/ref.ts

@@ -6,7 +6,7 @@ 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';
-import { createPtrFromOrigin } from './ptr.ts';
+import { createPtrFromOrigin, explicitFrom } from './ptr.ts';
 import { type ResolvedSnippet, snip, type Snippet } from './snippet.ts';
 import {
   isNaturallyEphemeral,
@@ -52,6 +52,12 @@ export const ref: DualFn<<T>(value: T) => ref<T>> = (() => {
       );
     }
 
+    if (value.dataType.type === 'ptr') {
+      // This can happen if we take a reference of an *implicit* pointer, one
+      // made by assigning a reference to a `const`.
+      return snip(value.value, explicitFrom(value.dataType), value.origin);
+    }
+
     /**
      * Pointer type only exists if the ref was created from a reference (buttery-butter).
      *

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -41,7 +41,7 @@ import {
 } from './generationHelpers.ts';
 import type { ShaderGenerator } from './shaderGenerator.ts';
 import type { DualFn } from '../data/dualFn.ts';
-import { INTERNAL_createPtr, ptrFn } from '../data/ptr.ts';
+import { createPtrFromOrigin, implicitFrom, ptrFn } from '../data/ptr.ts';
 import { RefOperator } from '../data/ref.ts';
 import { constant } from '../core/constant/tgpuConstant.ts';
 
@@ -869,18 +869,21 @@ ${this.ctx.pre}else ${alternate}`;
         } else {
           varType = 'let';
           if (!wgsl.isPtr(dataType)) {
-            dataType = ptrFn(concretize(dataType) as wgsl.StorableData);
+            const ptrType = createPtrFromOrigin(
+              eq.origin,
+              concretize(dataType) as wgsl.StorableData,
+            );
+            invariant(
+              ptrType !== undefined,
+              `Creating pointer type from origin ${eq.origin}`,
+            );
+            dataType = ptrType;
           }
 
           if (!(eq.value instanceof RefOperator)) {
             // If what we're assigning is something preceded by `&`, then it's a value
             // created using `d.ref()`. Otherwise, it's an implicit pointer
-            dataType = INTERNAL_createPtr(
-              dataType.addressSpace,
-              dataType.inner,
-              dataType.access,
-              /* implicit */ true,
-            );
+            dataType = implicitFrom(dataType);
           }
         }
       } else {
@@ -932,13 +935,11 @@ ${this.ctx.pre}else ${alternate}`;
       const [_, init, condition, update, body] = statement;
 
       const [initStatement, conditionExpr, updateStatement] = this.ctx
-        .withResetIndentLevel(
-          () => [
-            init ? this.statement(init) : undefined,
-            condition ? this.typedExpression(condition, bool) : undefined,
-            update ? this.statement(update) : undefined,
-          ],
-        );
+        .withResetIndentLevel(() => [
+          init ? this.statement(init) : undefined,
+          condition ? this.typedExpression(condition, bool) : undefined,
+          update ? this.statement(update) : undefined,
+        ]);
 
       const initStr = initStatement ? initStatement.slice(0, -1) : '';
       const updateStr = updateStatement ? updateStatement.slice(0, -1) : '';

packages/typegpu/tests/examples/individual/gravity.test.ts

@@ -42,15 +42,11 @@ describe('gravity example', () => {
 
       @group(0) @binding(0) var<uniform> celestialBodiesCount_3: i32;
 
-      fn radiusOf_4(body: ptr<function, CelestialBody_2>) -> f32 {
-        return (pow((((*body).mass * 0.75f) / 3.141592653589793f), 0.333f) * (*body).radiusMultiplier);
+      fn radiusOf_4(body: CelestialBody_2) -> f32 {
+        return (pow(((body.mass * 0.75f) / 3.141592653589793f), 0.333f) * body.radiusMultiplier);
       }
 
-      fn radiusOf_5(body: ptr<storage, ptr<function, CelestialBody_2>, read>) -> f32 {
-        return (pow((((*(*body)).mass * 0.75f) / 3.141592653589793f), 0.333f) * (*(*body)).radiusMultiplier);
-      }
-
-      fn isSmaller_6(currentId: u32, otherId: u32) -> bool {
+      fn isSmaller_5(currentId: u32, otherId: u32) -> bool {
         let current = (&inState_1[currentId]);
         let other = (&inState_1[otherId]);
         if (((*current).mass < (*other).mass)) {
@@ -62,33 +58,33 @@ describe('gravity example', () => {
         return false;
       }
 
-      @group(0) @binding(2) var<storage, read_write> outState_7: array<CelestialBody_2>;
+      @group(0) @binding(2) var<storage, read_write> outState_6: array<CelestialBody_2>;
 
-      struct computeCollisionsShader_Input_8 {
+      struct computeCollisionsShader_Input_7 {
         @builtin(global_invocation_id) gid: vec3u,
       }
 
-      @compute @workgroup_size(1) fn computeCollisionsShader_0(input: computeCollisionsShader_Input_8) {
+      @compute @workgroup_size(1) fn computeCollisionsShader_0(input: computeCollisionsShader_Input_7) {
         let currentId = input.gid.x;
         var current = inState_1[currentId];
         if ((current.destroyed == 0u)) {
           for (var otherId = 0u; (otherId < u32(celestialBodiesCount_3)); otherId++) {
             let other = (&inState_1[otherId]);
-            if ((((((otherId == currentId) || ((*other).destroyed == 1u)) || (current.collisionBehavior == 0u)) || ((*other).collisionBehavior == 0u)) || (distance(current.position, (*other).position) >= (radiusOf_4((&current)) + radiusOf_5((&other)))))) {
+            if ((((((otherId == currentId) || ((*other).destroyed == 1u)) || (current.collisionBehavior == 0u)) || ((*other).collisionBehavior == 0u)) || (distance(current.position, (*other).position) >= (radiusOf_4(current) + radiusOf_4((*other)))))) {
               continue;
             }
             if (((current.collisionBehavior == 1u) && ((*other).collisionBehavior == 1u))) {
-              if (isSmaller_6(currentId, otherId)) {
+              if (isSmaller_5(currentId, otherId)) {
                 var dir = normalize((current.position - (*other).position));
-                current.position = ((*other).position + (dir * (radiusOf_4((&current)) + radiusOf_5((&other)))));
+                current.position = ((*other).position + (dir * (radiusOf_4(current) + radiusOf_4((*other)))));
               }
               var posDiff = (current.position - (*other).position);
               var velDiff = (current.velocity - (*other).velocity);
               let posDiffFactor = ((((2f * (*other).mass) / (current.mass + (*other).mass)) * dot(velDiff, posDiff)) / dot(posDiff, posDiff));
               current.velocity = ((current.velocity - (posDiff * posDiffFactor)) * 0.99);
             }
             else {
-              let isCurrentAbsorbed = ((current.collisionBehavior == 1u) || ((current.collisionBehavior == 2u) && isSmaller_6(currentId, otherId)));
+              let isCurrentAbsorbed = ((current.collisionBehavior == 1u) || ((current.collisionBehavior == 2u) && isSmaller_5(currentId, otherId)));
               if (isCurrentAbsorbed) {
                 current.destroyed = 1u;
               }
@@ -101,7 +97,7 @@ describe('gravity example', () => {
             }
           }
         }
-        outState_7[currentId] = current;
+        outState_6[currentId] = current;
       }
 
       struct Time_2 {
@@ -126,21 +122,17 @@ describe('gravity example', () => {
 
       @group(1) @binding(0) var<uniform> celestialBodiesCount_5: i32;
 
-      fn radiusOf_6(body: ptr<function, CelestialBody_4>) -> f32 {
-        return (pow((((*body).mass * 0.75f) / 3.141592653589793f), 0.333f) * (*body).radiusMultiplier);
+      fn radiusOf_6(body: CelestialBody_4) -> f32 {
+        return (pow(((body.mass * 0.75f) / 3.141592653589793f), 0.333f) * body.radiusMultiplier);
       }
 
-      fn radiusOf_7(body: ptr<storage, ptr<function, CelestialBody_4>, read>) -> f32 {
-        return (pow((((*(*body)).mass * 0.75f) / 3.141592653589793f), 0.333f) * (*(*body)).radiusMultiplier);
-      }
-
-      @group(1) @binding(2) var<storage, read_write> outState_8: array<CelestialBody_4>;
+      @group(1) @binding(2) var<storage, read_write> outState_7: array<CelestialBody_4>;
 
-      struct computeGravityShader_Input_9 {
+      struct computeGravityShader_Input_8 {
         @builtin(global_invocation_id) gid: vec3u,
       }
 
-      @compute @workgroup_size(1) fn computeGravityShader_0(input: computeGravityShader_Input_9) {
+      @compute @workgroup_size(1) fn computeGravityShader_0(input: computeGravityShader_Input_8) {
         let dt = (time_1.passed * time_1.multiplier);
         let currentId = input.gid.x;
         var current = inState_3[currentId];
@@ -150,14 +142,14 @@ describe('gravity example', () => {
             if (((otherId == currentId) || ((*other).destroyed == 1u))) {
               continue;
             }
-            let dist = max((radiusOf_6((&current)) + radiusOf_7((&other))), distance(current.position, (*other).position));
+            let dist = max((radiusOf_6(current) + radiusOf_6((*other))), distance(current.position, (*other).position));
             let gravityForce = (((current.mass * (*other).mass) / dist) / dist);
             var direction = normalize(((*other).position - current.position));
             current.velocity = (current.velocity + (direction * ((gravityForce / current.mass) * dt)));
           }
           current.position = (current.position + (current.velocity * dt));
         }
-        outState_8[currentId] = current;
+        outState_7[currentId] = current;
       }
 
       struct Camera_2 {
@@ -209,8 +201,8 @@ describe('gravity example', () => {
 
       @group(1) @binding(1) var<storage, read> celestialBodies_1: array<CelestialBody_2>;
 
-      fn radiusOf_3(body: ptr<storage, ptr<function, CelestialBody_2>, read>) -> f32 {
-        return (pow((((*(*body)).mass * 0.75f) / 3.141592653589793f), 0.333f) * (*(*body)).radiusMultiplier);
+      fn radiusOf_3(body: CelestialBody_2) -> f32 {
+        return (pow(((body.mass * 0.75f) / 3.141592653589793f), 0.333f) * body.radiusMultiplier);
       }
 
       struct Camera_5 {
@@ -241,7 +233,7 @@ describe('gravity example', () => {
 
       @vertex fn mainVertex_0(input: mainVertex_Input_7) -> mainVertex_Output_6 {
         let currentBody = (&celestialBodies_1[input.instanceIndex]);
-        var worldPosition = ((*currentBody).position + (input.position.xyz * radiusOf_3((&currentBody))));
+        var worldPosition = ((*currentBody).position + (input.position.xyz * radiusOf_3((*currentBody))));
         let camera = (&camera_4);
         var positionOnCanvas = (((*camera).projection * (*camera).view) * vec4f(worldPosition, 1f));
         return mainVertex_Output_6(positionOnCanvas, input.uv, input.normal, worldPosition, (*currentBody).textureIndex, (*currentBody).destroyed, (*currentBody).ambientLightFactor);

packages/typegpu/tests/examples/individual/tgsl-parsing-test.test.ts

@@ -178,24 +178,18 @@ describe('tgsl parsing test example', () => {
         (*ptr).vec.x += 1f;
       }
 
-      var<private> privateNum_16: u32;
-
-      fn modifyNumPrivate_17(ptr: ptr<private, u32>) {
-        (*ptr) += 1u;
-      }
-
-      var<private> privateVec_18: vec2f;
-
-      fn modifyVecPrivate_19(ptr: ptr<private, vec2f>) {
+      fn modifyVecPrivate_16(ptr: ptr<private, vec2f>) {
         (*ptr).x += 1f;
       }
 
-      var<private> privateStruct_20: SimpleStruct_14;
+      var<private> privateVec_17: vec2f;
 
-      fn modifyStructPrivate_21(ptr: ptr<private, SimpleStruct_14>) {
+      fn modifyStructPrivate_18(ptr: ptr<private, SimpleStruct_14>) {
         (*ptr).vec.x += 1f;
       }
 
+      var<private> privateStruct_19: SimpleStruct_14;
+
       fn pointersTest_11() -> bool {
         var s = true;
         var num = 0u;
@@ -207,16 +201,14 @@ describe('tgsl parsing test example', () => {
         var myStruct = SimpleStruct_14();
         modifyStructFn_15((&myStruct));
         s = (s && all(myStruct.vec == vec2f(1, 0)));
-        modifyNumPrivate_17(privateNum_16);
-        s = (s && (privateNum_16 == 1u));
-        modifyVecPrivate_19(privateVec_18);
-        s = (s && all(privateVec_18 == vec2f(1, 0)));
-        modifyStructPrivate_21(privateStruct_20);
-        s = (s && all(privateStruct_20.vec == vec2f(1, 0)));
+        modifyVecPrivate_16((&privateVec_17));
+        s = (s && all(privateVec_17 == vec2f(1, 0)));
+        modifyStructPrivate_18((&privateStruct_19));
+        s = (s && all(privateStruct_19.vec == vec2f(1, 0)));
         return s;
       }
 
-      @group(0) @binding(0) var<storage, read_write> result_22: i32;
+      @group(0) @binding(0) var<storage, read_write> result_20: i32;
 
       @compute @workgroup_size(1) fn computeRunTests_0() {
         var s = true;
@@ -226,10 +218,10 @@ describe('tgsl parsing test example', () => {
         s = (s && arrayAndStructConstructorsTest_8());
         s = (s && pointersTest_11());
         if (s) {
-          result_22 = 1i;
+          result_20 = 1i;
         }
         else {
-          result_22 = 0i;
+          result_20 = 0i;
         }
       }"
     `);