feat(🫵): Add ptr<function, ...> schema (#806)

Author: iwoplaza

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');
 }

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);
   },

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,

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>>;
+}

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.
  *

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>"`);
+  });
+});