Skip to content

Commit

Permalink
feat(🫵): Add ptr<function, ...> schema (#806)
Browse files Browse the repository at this point in the history
  • Loading branch information
iwoplaza authored Jan 29, 2025
1 parent e5c4db8 commit 7b2cbaa
Show file tree
Hide file tree
Showing 6 changed files with 64 additions and 0 deletions.
4 changes: 4 additions & 0 deletions packages/typegpu/src/core/resolve/resolveData.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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');
}
8 changes: 8 additions & 0 deletions packages/typegpu/src/data/dataIO.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
},
Expand Down Expand Up @@ -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);
},
Expand Down
3 changes: 3 additions & 0 deletions packages/typegpu/src/data/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ export {
isWgslData,
isWgslArray,
isWgslStruct,
isPtrFn,
isAtomic,
isDecorated,
isAlignAttrib,
Expand Down Expand Up @@ -39,6 +40,7 @@ export type {
Mat4x4f,
WgslStruct,
WgslArray,
PtrFn,
Atomic,
Decorated,
Size,
Expand Down Expand Up @@ -68,6 +70,7 @@ export {
TgpuArray,
arrayOf,
} from './array';
export { ptrFn } from './ptr';
export type {
Disarray,
Unstruct,
Expand Down
10 changes: 10 additions & 0 deletions packages/typegpu/src/data/ptr.ts
Original file line number Diff line number Diff line change
@@ -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>>;
}
20 changes: 20 additions & 0 deletions packages/typegpu/src/data/wgslTypes.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*/
Expand Down Expand Up @@ -770,6 +777,7 @@ export const wgslTypeLiterals = [
'mat4x4f',
'struct',
'array',
'ptrFn',
'atomic',
'decorated',
] as const;
Expand Down Expand Up @@ -820,6 +828,7 @@ export type AnyWgslData =
| Mat4x4f
| AnyWgslStruct
| WgslArray
| PtrFn
| Atomic
| Decorated;

Expand Down Expand Up @@ -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.
*
Expand Down
19 changes: 19 additions & 0 deletions packages/typegpu/tests/data/ptr.test.ts
Original file line number Diff line number Diff line change
@@ -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>"`);
});
});

0 comments on commit 7b2cbaa

Please sign in to comment.