fix: Pointers for reference types (#1591)

Author: aleksanderkatan

apps/typegpu-docs/src/content/docs/fundamentals/data-schemas.mdx

@@ -321,7 +321,6 @@ To define arrays of known constant length, use the `d.arrayOf` function with the
 
 ```ts twoslash
 import * as d from 'typegpu/data';
-
 // ---cut---
 const MyArray = d.arrayOf(d.f32, 4);
 
@@ -344,6 +343,57 @@ const array = ArrayPartialSchema(2)([1.2, 19.29]);
 //    ^?
 ```
 
+## Atomics
+
+To create a schema corresponding to an atomic data type, wrap `d.i32` or `d.u32` with `d.atomic`.
+
+```ts twoslash
+import * as d from 'typegpu/data';
+// ---cut---
+const AtomicI32 = d.atomic(d.i32);
+//    ^?
+```
+
+:::tip
+The `std` module provides functions for manipulating atomic data in TGSL.
+
+```ts twoslash
+import tgpu from 'typegpu';
+import * as d from 'typegpu/data';
+import * as std from 'typegpu/std';
+// ---cut---
+const count = tgpu['~unstable'].workgroupVar(d.atomic(d.u32));
+
+const incrementAndGet = tgpu.fn([], d.u32)(() => {
+  return std.atomicAdd(count.$, 1);
+});
+```
+
+:::
+
+## Pointers
+
+To create a schema corresponding to a pointer data type, wrap the inner schema with a pointer constructor corresponding to the appropriate address space.
+
+For address spaces `function`, `private`, `workgroup`, `storage` and `uniform`, use `d.ptrFn`, `d.ptrPrivate`, `d.ptrWorkgroup`, `d.ptrStorage` and `d.ptrUniform` respectively.
+
+```ts twoslash
+import * as d from 'typegpu/data';
+// ---cut---
+const vecPtrFn = d.ptrFn(d.vec3f);
+//    ^?
+
+const vecPtrPrivate = d.ptrPrivate(d.vec2f);
+//    ^?
+```
+
+:::note
+In order to use `workgroup`, `storage` and `uniform` pointers as function parameters,
+WGSL needs the `unrestricted_pointer_parameters` extension.
+
+This extension is enabled automatically when available.
+:::
+
 ## Copy constructors
 
 One of the biggest differences between JavaScript and WGSL is the existence of value/reference types.

apps/typegpu-docs/src/content/docs/fundamentals/tgsl.mdx

@@ -18,7 +18,7 @@ For the TGSL functions to work, you need to use the dedicated build plugin -- [u
 ## Usage
 
 Instead of using a WGSL code string, you can pass TGSL to the tgpu function shell as an argument instead.
-Functions from the WGSL standard library (*distance, arrayLength, workgroupBarrier, [etc.](https://github.com/software-mansion/TypeGPU/blob/release/packages/typegpu/src/std/index.ts)*) are accessible through the `typegpu/std` endpoint. 
+Functions from the WGSL standard library (*distance, arrayLength, workgroupBarrier, [etc.](https://github.com/software-mansion/TypeGPU/blob/release/packages/typegpu/src/std/index.ts)*) are accessible through the `typegpu/std` entrypoint.
 The package also includes functions for vector and matrix operators (*add, eq, lt...*).
 
 ```ts
@@ -103,21 +103,27 @@ const patternSolid = patternFn(() => {
 });
 ```
 
-TGSL-implemented functions can also be invoked on the CPU, as along as they do not use any GPU-exclusive functionalities, like buffers or textures (regardless of whether they are marked as *"kernel"* or not).
+## Executing TGSL functions in JS
+
+TGSL-implemented functions can be invoked on the CPU,
+as along as they do not use any GPU-exclusive functionalities,
+like buffers or textures (regardless of whether they are marked as *"kernel"* or not).
+
+Keep in mind that you cannot execute entry-point functions in JavaScript.
 
 ## What to keep in mind
 
 * **TGSL limitations** --
-For a function to be valid TGSL, it must consist only of supported JS syntax (again, see [tinyest-for-wgsl repository](https://github.com/software-mansion/TypeGPU/blob/release/packages/tinyest-for-wgsl/src/parsers.ts)), possibly including references to bound buffers, constant variables defined outside of the function, other TGSL functions etc.
+For a function to be valid TGSL, it must consist only of supported JS syntax (again, see [tinyest-for-wgsl repository](https://github.com/software-mansion/TypeGPU/blob/release/packages/tinyest-for-wgsl/src/parsers.ts)), possibly including references to bound buffers, constant variables defined outside the function, other TGSL functions etc.
 This means that, for example, `console.log()` calls will not work on the GPU.
 
 * **Differences between JS on the CPU and GPU** --
-TGSL is developed to work on the GPU the same as on the CPU as much as possible, 
+TGSL is developed to work on the GPU the same as on the CPU as much as possible,
 however because of the fundamental differences between the JavaScript and WGSL languages, it is not guaranteed to always be the case.
 
-    Currently the biggest known difference is that vectors, matrices and structs are treated as reference types in JavaScript and value types in WGSL.
-That is, on the WGSL side, the assignment operator copies the value instead of the reference, and two different vectors can be equal to each other if only they store the same values, unlike in JS, where they need to point to the same reference. 
-To somehow alleviate this issue, when passing arguments to tgpu functions on JS side, we perform a deep copy of them (note that in WGSL arguments are immutable by default). 
+    Currently, the biggest known difference is that vectors, matrices and structs are treated as reference types in JavaScript and value types in WGSL.
+That is, on the WGSL side, the assignment operator copies the value instead of the reference, and two different vectors can be equal to each other if only they store the same values, unlike in JS, where they need to point to the same reference.
+To somehow alleviate this issue, when passing arguments to tgpu functions on JS side, we perform a deep copy of them (note that in WGSL arguments are immutable by default).
 
     When using TGSL on the GPU, the behavior is that of WGSL, not JS, as one would expect.
 Therefore some WGSL knowledge is still required, even when opting out for TGSL.
@@ -146,15 +152,32 @@ const fragmentWgsl = tgpu['~unstable'].fragmentFn({ out: d.vec4f })`{
 ```
 
 * **Operators** --
-JavaScript does not support operator overloading. 
-This means that, while you can still use operators for numbers, 
+JavaScript does not support operator overloading.
+This means that, while you can still use operators for numbers,
 you have to use supplementary functions from `typegpu/std` (*add, mul, eq, lt, ge...*) for operations involving vectors and matrices.
 
+* **Pointers** --
+Since WGSL pointers can only point to a narrow set of items called Storable Types,
+TypeGPU tries to dereference pointers used in TGSL automatically when it seems appropriate.
+For example, function arguments that are pointers to reference types are passed just as the pointed object.
+Pointers to primitive types (numbers and booleans) are currently not supported.
+
+```ts twoslash
+import tgpu from 'typegpu';
+import * as d from 'typegpu/data';
+// ---cut---
+const fn = tgpu.fn([d.ptrFn(d.vec3f)], d.vec3f)((ptr) => {
+  ptr.x += 1;
+//^?
+  return ptr;
+}); 
+```
+
 * **When to use TGSL instead of WGSL** --
-Writing the code using TGSL has a few significant advantages. 
+Writing the code using TGSL has a few significant advantages.
 It allows defining utils only once and using them both as a kernel and host functions,
 as well as enables complete syntax highlighting and autocomplete in TypeGPU function definitions, leading to a better developer UX.
-However, it sometimes might be better to choose WGSL for certain functions. 
+However, it sometimes might be better to choose WGSL for certain functions.
 Since JavaScript doesn't support operator overloading, functions including complex matrix operations can be more readable in WGSL.
 Writing WGSL becomes a necessity whenever TGSL does not support some feature or standard library function quite yet.
 Luckily, you don't have to choose one or the other for the entire project. It is possible to mix and match WGSL and TGSL at every step of the way.

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

@@ -4,6 +4,7 @@ import { arrayAndStructConstructorsTest } from './array-and-struct-constructors.
 import { infixOperatorsTests } from './infix-operators.ts';
 import { logicalExpressionTests } from './logical-expressions.ts';
 import { matrixOpsTests } from './matrix-ops.ts';
+import { pointersTest } from './pointers.ts';
 
 const root = await tgpu.init();
 const result = root.createMutable(d.i32, 0);
@@ -15,6 +16,7 @@ const computeRunTests = tgpu['~unstable']
     s = s && matrixOpsTests();
     s = s && infixOperatorsTests();
     s = s && arrayAndStructConstructorsTest();
+    s = s && pointersTest();
 
     if (s) {
       result.value = 1;

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

@@ -0,0 +1,64 @@
+import tgpu from 'typegpu';
+import * as d from 'typegpu/data';
+import * as std from 'typegpu/std';
+
+const SimpleStruct = d.struct({ vec: d.vec2f });
+
+const modifyNumFn = tgpu.fn([d.ptrFn(d.u32)])((ptr) => {
+  // biome-ignore lint/style/noParameterAssign: it's just a test
+  ptr += 1;
+});
+
+const modifyVecFn = tgpu.fn([d.ptrFn(d.vec2f)])((ptr) => {
+  ptr.x += 1;
+});
+
+const modifyStructFn = tgpu.fn([d.ptrFn(SimpleStruct)])((ptr) => {
+  ptr.vec.x += 1;
+});
+
+const privateNum = tgpu['~unstable'].privateVar(d.u32);
+const modifyNumPrivate = tgpu.fn([d.ptrPrivate(d.u32)])((ptr) => {
+  // biome-ignore lint/style/noParameterAssign: it's just a test
+  ptr += 1;
+});
+
+const privateVec = tgpu['~unstable'].privateVar(d.vec2f);
+const modifyVecPrivate = tgpu.fn([d.ptrPrivate(d.vec2f)])((ptr) => {
+  ptr.x += 1;
+});
+
+const privateStruct = tgpu['~unstable'].privateVar(SimpleStruct);
+const modifyStructPrivate = tgpu.fn([d.ptrPrivate(SimpleStruct)])((ptr) => {
+  ptr.vec.x += 1;
+});
+
+// TODO: replace `s = s &&` with `s &&=` when implemented
+export const pointersTest = tgpu.fn([], d.bool)(() => {
+  let s = true;
+
+  // function pointers
+  const num = d.u32();
+  modifyNumFn(num);
+  s = s && (num === 1);
+
+  const vec = d.vec2f();
+  modifyVecFn(vec);
+  s = s && std.allEq(vec, d.vec2f(1, 0));
+
+  const myStruct = SimpleStruct();
+  modifyStructFn(myStruct);
+  s = s && std.allEq(myStruct.vec, d.vec2f(1, 0));
+
+  // private pointers
+  modifyNumPrivate(privateNum.$);
+  s = s && (privateNum.$ === 1);
+
+  modifyVecPrivate(privateVec.$);
+  s = s && std.allEq(privateVec.$, d.vec2f(1, 0));
+
+  modifyStructPrivate(privateStruct.$);
+  s = s && std.allEq(privateStruct.$.vec, d.vec2f(1, 0));
+
+  return s;
+});

packages/typegpu/src/data/index.ts

@@ -73,6 +73,7 @@ export type {
   Mat4x4f,
   Ptr,
   Size,
+  StorableData,
   U16,
   U32,
   v2b,

packages/typegpu/src/data/ptr.ts

@@ -1,76 +1,57 @@
 import { $internal } from '../shared/symbols.ts';
-import type { AnyData } from './dataTypes.ts';
-import type { Ptr } from './wgslTypes.ts';
+import type { Access, AddressSpace, Ptr, StorableData } from './wgslTypes.ts';
 
-export function ptrFn<T extends AnyData>(
+export function ptrFn<T extends StorableData>(
   inner: T,
 ): Ptr<'function', T, 'read-write'> {
-  return {
-    [$internal]: true,
-    type: 'ptr',
-    inner,
-    addressSpace: 'function',
-    access: 'read-write',
-  } as Ptr<'function', T, 'read-write'>;
+  return INTERNAL_createPtr('function', inner, 'read-write');
 }
 
-export function ptrPrivate<T extends AnyData>(
+export function ptrPrivate<T extends StorableData>(
   inner: T,
 ): Ptr<'private', T, 'read-write'> {
-  return {
-    [$internal]: true,
-    type: 'ptr',
-    inner,
-    addressSpace: 'private',
-    access: 'read-write',
-  } as Ptr<'private', T, 'read-write'>;
+  return INTERNAL_createPtr('private', inner, 'read-write');
 }
 
-export function ptrWorkgroup<T extends AnyData>(
+export function ptrWorkgroup<T extends StorableData>(
   inner: T,
 ): Ptr<'workgroup', T, 'read-write'> {
-  return {
-    [$internal]: true,
-    type: 'ptr',
-    inner,
-    addressSpace: 'workgroup',
-    access: 'read-write',
-  } as Ptr<'workgroup', T, 'read-write'>;
+  return INTERNAL_createPtr('workgroup', inner, 'read-write');
 }
 
 export function ptrStorage<
-  T extends AnyData,
+  T extends StorableData,
   TAccess extends 'read' | 'read-write' = 'read',
 >(inner: T, access: TAccess = 'read' as TAccess): Ptr<'storage', T, TAccess> {
-  return {
-    [$internal]: true,
-    type: 'ptr',
-    inner,
-    addressSpace: 'storage',
-    access,
-  } as Ptr<'storage', T, TAccess>;
+  return INTERNAL_createPtr('storage', inner, access);
 }
 
-export function ptrUniform<T extends AnyData>(
+export function ptrUniform<T extends StorableData>(
   inner: T,
 ): Ptr<'uniform', T, 'read'> {
-  return {
-    [$internal]: true,
-    type: 'ptr',
-    inner,
-    addressSpace: 'uniform',
-    access: 'read',
-  } as Ptr<'uniform', T, 'read'>;
+  return INTERNAL_createPtr('uniform', inner, 'read');
 }
 
-export function ptrHandle<T extends AnyData>(
+export function ptrHandle<T extends StorableData>(
   inner: T,
 ): Ptr<'handle', T, 'read'> {
+  return INTERNAL_createPtr('handle', inner, 'read');
+}
+
+function INTERNAL_createPtr<
+  TAddressSpace extends AddressSpace,
+  TInner extends StorableData,
+  TAccess extends Access,
+>(
+  addressSpace: TAddressSpace,
+  inner: TInner,
+  access: TAccess,
+): Ptr<TAddressSpace, TInner, TAccess> {
   return {
     [$internal]: true,
     type: 'ptr',
+    addressSpace,
     inner,
-    addressSpace: 'handle',
-    access: 'read',
-  } as Ptr<'handle', T, 'read'>;
+    access,
+  } as Ptr<TAddressSpace, TInner, TAccess>;
 }

packages/typegpu/src/data/schemaCallWrapper.ts

@@ -2,9 +2,11 @@ import type { AnyData } from './dataTypes.ts';
 import { formatToWGSLType } from './vertexFormatData.ts';
 
 /**
- * A wrapper for `schema(item)` or `schema()` call.
- * If the schema is a TgpuVertexFormatData, it instead calls the corresponding constructible schema.
- * Throws an error if the schema is not callable.
+ * A wrapper for `schema(item)` or `schema()` call on JS side.
+ * If the schema is a pointer, returns the value pointed to without copying.
+ * If the schema is a TgpuVertexFormatData, calls the corresponding constructible schema instead.
+ * If the schema is not callable, throws an error.
+ * Otherwise, returns `schema(item)` or `schema()`.
  */
 export function schemaCallWrapper<T>(schema: AnyData, item?: T): T {
   const maybeType = (schema as { type: string })?.type;