fix: Pointers for reference types

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);
 
@@ -331,6 +330,34 @@ const myDefaultArray = MyArray(); // [0, 0, 0, 0]
 const myInitializedArray = MyArray([1, 0, 0.5, 20]);
 ```
 
+## 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);
+//    ^?
+```
+
+## 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`, `uniform`, `handle`, use `d.ptrFn`, `d.ptrPrivate`, `d.ptrWorkgroup`, `d.ptrStorage`, `d.ptrUniform`, `d.ptrHandle` respectively.
+
+```ts twoslash
+import * as d from 'typegpu/data';
+// ---cut---
+const VecPtrFn = d.ptrFn(d.vec3f);
+//    ^?
+
+const VecPtrUniform = d.ptrUniform(d.vec2f);
+//    ^?
+```
+
+
 ## 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

@@ -103,7 +103,13 @@ 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
 
@@ -150,6 +156,22 @@ 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, 
+TGPU tries to dereference all pointers used in TGSL automatically. 
+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. 
 It allows defining utils only once and using them both as a kernel and host functions,

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,37 @@
+import tgpu from 'typegpu';
+import * as d from 'typegpu/data';
+import * as std from 'typegpu/std';
+
+const SimpleStruct = d.struct({ vec: d.vec2f });
+
+const modifyNum = tgpu.fn([d.ptrFn(d.u32)])((ptr) => {
+  // biome-ignore lint/style/noParameterAssign: it's just a test
+  ptr += 1;
+});
+
+const modifyVec = tgpu.fn([d.ptrFn(d.vec2f)])((ptr) => {
+  ptr.x += 1;
+});
+
+const modifyStruct = tgpu.fn([d.ptrFn(SimpleStruct)])((ptr) => {
+  ptr.vec.x += 1;
+});
+
+// TODO: replace `s = s &&` with `s &&=` when implemented
+export const pointersTest = tgpu.fn([], d.bool)(() => {
+  let s = true;
+
+  const num = d.u32(0);
+  modifyNum(num);
+  s = s && (num === 1);
+
+  const vec = d.vec2f(1, 2);
+  modifyVec(vec);
+  s = s && std.allEq(vec, d.vec2f(2, 2));
+
+  const myStruct = SimpleStruct({ vec: d.vec2f(3, 4) });
+  modifyStruct(myStruct);
+  s = s && std.allEq(myStruct.vec, d.vec2f(4, 4));
+
+  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/utils.ts

@@ -1,15 +1,21 @@
-import type { AnyData } from './index.ts';
+import type { AnyData } from './dataTypes.ts';
 import { formatToWGSLType } from './vertexFormatData.ts';
+import { isPtr } from './wgslTypes.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;
 
   try {
+    if (item !== undefined && isPtr(schema)) {
+      return item;
+    }
     // TgpuVertexFormatData are not callable
     const callSchema = (maybeType in formatToWGSLType
       ? formatToWGSLType[maybeType as keyof typeof formatToWGSLType]

packages/typegpu/src/data/wgslTypes.ts

@@ -1,3 +1,5 @@
+import type { TgpuSampler } from '../core/sampler/sampler.ts';
+import type { TgpuTexture } from '../core/texture/texture.ts';
 import type { TgpuNamable } from '../shared/meta.ts';
 import type {
   ExtractInvalidSchemaError,
@@ -1339,7 +1341,7 @@ export type Access = 'read' | 'write' | 'read-write';
 
 export interface Ptr<
   TAddr extends AddressSpace = AddressSpace,
-  TInner extends BaseData = BaseData, // can also be sampler or texture (╯'□')╯︵ ┻━┻
+  TInner extends StorableData = StorableData, // can also be sampler or texture (╯'□')╯︵ ┻━┻
   TAccess extends Access = Access,
 > extends BaseData {
   readonly type: 'ptr';
@@ -1521,6 +1523,38 @@ export type ScalarData =
   | AbstractInt
   | AbstractFloat;
 
+export type VecData =
+  | Vec2f
+  | Vec2h
+  | Vec2i
+  | Vec2u
+  | Vec2b
+  | Vec3f
+  | Vec3h
+  | Vec3i
+  | Vec3u
+  | Vec3b
+  | Vec4f
+  | Vec4h
+  | Vec4i
+  | Vec4u
+  | Vec4b;
+
+export type MatData =
+  | Mat2x2f
+  | Mat3x3f
+  | Mat4x4f;
+
+export type StorableData =
+  | ScalarData
+  | VecData
+  | MatData
+  | Atomic
+  | WgslArray
+  | WgslStruct
+  | TgpuTexture
+  | TgpuSampler;
+
 export type AnyWgslData =
   | Bool
   | F32
@@ -1687,7 +1721,7 @@ export function isWgslStruct<T extends WgslStruct>(
 }
 
 /**
- * Checks whether passed in value is a pointer ('function' scope) schema.
+ * Checks whether passed in value is a pointer schema.
  *
  * @example
  * isPtrFn(d.ptrFn(d.f32)) // true

packages/typegpu/src/std/array.ts

@@ -1,9 +1,8 @@
-import { snip } from '../data/snippet.ts';
+import { createDualImpl } from '../core/function/dualImpl.ts';
 import { abstractInt, u32 } from '../data/numeric.ts';
 import { ptrFn } from '../data/ptr.ts';
-import type { AnyWgslData } from '../data/wgslTypes.ts';
-import { isPtr, isWgslArray } from '../data/wgslTypes.ts';
-import { createDualImpl } from '../core/function/dualImpl.ts';
+import { snip } from '../data/snippet.ts';
+import { isPtr, isWgslArray, type StorableData } from '../data/wgslTypes.ts';
 
 export const arrayLength = createDualImpl(
   // CPU implementation
@@ -19,5 +18,5 @@ export const arrayLength = createDualImpl(
     return snip(`arrayLength(${a.value})`, u32);
   },
   'arrayLength',
-  (a) => [ptrFn(a.dataType as AnyWgslData)],
+  (a) => [ptrFn(a.dataType as StorableData)],
 );

packages/typegpu/tests/data/ptr.test.ts

@@ -18,6 +18,20 @@ describe('d.ptrFn', () => {
       tgpu.resolve({ externals: { ptrToU32 }, template: 'ptrToU32' }),
     ).toMatchInlineSnapshot(`"ptr<function, u32>"`);
   });
+
+  it('modifies reference types in JS', () => {
+    const modifyVec = tgpu.fn([d.ptrFn(d.vec2f)])((ptr) => {
+      ptr.x += 1;
+    });
+
+    const testFn = tgpu.fn([], d.vec2f)(() => {
+      const vec = d.vec2f(1, 2);
+      modifyVec(vec);
+      return vec;
+    });
+
+    expect(testFn()).toStrictEqual(d.vec2f(2, 2));
+  });
 });
 
 describe('d.ptrPrivate', () => {