From 7b2cbaa2ce7e5b368a776e34bae23062bf336e61 Mon Sep 17 00:00:00 2001
From: Iwo Plaza <iwoplaza@gmail.com>
Date: Wed, 29 Jan 2025 13:54:56 +0000
Subject: [PATCH] =?UTF-8?q?feat(=F0=9F=AB=B5):=20Add=20ptr<function,=20...?=
 =?UTF-8?q?>=20schema=20(#806)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../typegpu/src/core/resolve/resolveData.ts   |  4 ++++
 packages/typegpu/src/data/dataIO.ts           |  8 ++++++++
 packages/typegpu/src/data/index.ts            |  3 +++
 packages/typegpu/src/data/ptr.ts              | 10 ++++++++++
 packages/typegpu/src/data/wgslTypes.ts        | 20 +++++++++++++++++++
 packages/typegpu/tests/data/ptr.test.ts       | 19 ++++++++++++++++++
 6 files changed, 64 insertions(+)
 create mode 100644 packages/typegpu/src/data/ptr.ts
 create mode 100644 packages/typegpu/tests/data/ptr.test.ts

diff --git a/packages/typegpu/src/core/resolve/resolveData.ts b/packages/typegpu/src/core/resolve/resolveData.ts
index 86d942bd7..a3e206cd4 100644
--- a/packages/typegpu/src/core/resolve/resolveData.ts
+++ b/packages/typegpu/src/core/resolve/resolveData.ts
@@ -165,5 +165,9 @@ export function resolveData(ctx: ResolutionCtx, data: AnyWgslData): string {
     return ctx.resolve(data.inner as AnyWgslData);
   }
 
+  if (data.type === 'ptrFn') {
+    return `ptr<function, ${ctx.resolve(data.inner)}>`;
+  }
+
   assertExhaustive(data, 'resolveData');
 }
diff --git a/packages/typegpu/src/data/dataIO.ts b/packages/typegpu/src/data/dataIO.ts
index 1363e6111..b3e1cfbad 100644
--- a/packages/typegpu/src/data/dataIO.ts
+++ b/packages/typegpu/src/data/dataIO.ts
@@ -186,6 +186,10 @@ const dataWriters = {
     output.seekTo(beginning + sizeOf(schema));
   },
 
+  ptrFn() {
+    throw new Error('Pointers are not host-shareable');
+  },
+
   atomic(output, schema: wgsl.Atomic, value: number) {
     dataWriters[schema.inner.type]?.(output, schema, value);
   },
@@ -615,6 +619,10 @@ const dataReaders = {
     return elements as never[];
   },
 
+  ptrFn() {
+    throw new Error('Pointers are not host-shareable');
+  },
+
   atomic(input, schema: wgsl.Atomic): number {
     return readData(input, schema.inner);
   },
diff --git a/packages/typegpu/src/data/index.ts b/packages/typegpu/src/data/index.ts
index 92657dfa4..e37ade6d0 100644
--- a/packages/typegpu/src/data/index.ts
+++ b/packages/typegpu/src/data/index.ts
@@ -7,6 +7,7 @@ export {
   isWgslData,
   isWgslArray,
   isWgslStruct,
+  isPtrFn,
   isAtomic,
   isDecorated,
   isAlignAttrib,
@@ -39,6 +40,7 @@ export type {
   Mat4x4f,
   WgslStruct,
   WgslArray,
+  PtrFn,
   Atomic,
   Decorated,
   Size,
@@ -68,6 +70,7 @@ export {
   TgpuArray,
   arrayOf,
 } from './array';
+export { ptrFn } from './ptr';
 export type {
   Disarray,
   Unstruct,
diff --git a/packages/typegpu/src/data/ptr.ts b/packages/typegpu/src/data/ptr.ts
new file mode 100644
index 000000000..be506c3c3
--- /dev/null
+++ b/packages/typegpu/src/data/ptr.ts
@@ -0,0 +1,10 @@
+import type { AnyData } from './dataTypes';
+import type { Exotic } from './exotic';
+import type { PtrFn } from './wgslTypes';
+
+export function ptrFn<T extends AnyData>(inner: T): PtrFn<Exotic<T>> {
+  return {
+    type: 'ptrFn',
+    inner: inner as Exotic<T>,
+  } as PtrFn<Exotic<T>>;
+}
diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts
index 865d17e93..f63324e74 100644
--- a/packages/typegpu/src/data/wgslTypes.ts
+++ b/packages/typegpu/src/data/wgslTypes.ts
@@ -694,6 +694,13 @@ export interface WgslArray<TElement = BaseWgslData> {
   readonly '~repr': Infer<TElement>[];
 }
 
+export interface PtrFn<TInner = BaseWgslData> {
+  readonly type: 'ptrFn';
+  readonly inner: TInner;
+  /** Type-token, not available at runtime */
+  readonly '~repr': Infer<TInner>;
+}
+
 /**
  * Schema representing the `atomic<...>` WGSL data type.
  */
@@ -770,6 +777,7 @@ export const wgslTypeLiterals = [
   'mat4x4f',
   'struct',
   'array',
+  'ptrFn',
   'atomic',
   'decorated',
 ] as const;
@@ -820,6 +828,7 @@ export type AnyWgslData =
   | Mat4x4f
   | AnyWgslStruct
   | WgslArray
+  | PtrFn
   | Atomic
   | Decorated;
 
@@ -865,6 +874,17 @@ export function isWgslStruct<T extends WgslStruct>(
   return (schema as T)?.type === 'struct';
 }
 
+/**
+ * Checks whether passed in value is a pointer ('function' scope) schema.
+ *
+ * @example
+ * isPtrFn(d.ptrFn(d.f32)) // true
+ * isPtrFn(d.f32) // false
+ */
+export function isPtrFn<T extends PtrFn>(schema: T | unknown): schema is T {
+  return (schema as T)?.type === 'ptrFn';
+}
+
 /**
  * Checks whether the passed in value is an atomic schema.
  *
diff --git a/packages/typegpu/tests/data/ptr.test.ts b/packages/typegpu/tests/data/ptr.test.ts
new file mode 100644
index 000000000..2ba44a3b9
--- /dev/null
+++ b/packages/typegpu/tests/data/ptr.test.ts
@@ -0,0 +1,19 @@
+import { describe, expect, expectTypeOf, it } from 'vitest';
+import tgpu from '../../src';
+import * as d from '../../src/data';
+
+describe('d.ptrFn', () => {
+  it('wraps a schema and infers type properly', () => {
+    const ptrToU32 = d.ptrFn(d.u32);
+
+    expectTypeOf(ptrToU32).toEqualTypeOf<d.PtrFn<d.U32>>();
+  });
+
+  it('resolves to matching WGSL', () => {
+    const ptrToU32 = d.ptrFn(d.u32);
+
+    expect(
+      tgpu.resolve({ externals: { ptrToU32 }, template: 'ptrToU32' }),
+    ).toMatchInlineSnapshot(`"ptr<function, u32>"`);
+  });
+});