refactor: Use internal symbol to identify TypeGPU-originating values (#1171)

* refactor: Use internal symbol to identify TypeGPU-originating values

* Fixed dual impl internal types

* Smaller builtins

Author: iwoplaza

packages/typegpu/src/builtin.ts

@@ -1,8 +1,10 @@
 import { arrayOf } from './data/array.ts';
 import { attribute } from './data/attributes.ts';
+import type { LooseDecorated } from './data/dataTypes.ts';
 import { f32, u32 } from './data/numeric.ts';
 import { vec3u, vec4f } from './data/vector.ts';
 import type {
+  AnyWgslData,
   BaseData,
   Builtin,
   Decorated,
@@ -12,6 +14,7 @@ import type {
   Vec4f,
   WgslArray,
 } from './data/wgslTypes.ts';
+import { $internal } from './shared/symbols.ts';
 
 // ----------
 // Public API
@@ -51,67 +54,49 @@ export type BuiltinSubgroupInvocationId = Decorated<
 >;
 export type BuiltinSubgroupSize = Decorated<U32, [Builtin<'subgroup_size'>]>;
 
-export const builtin = {
-  vertexIndex: attribute(u32, {
-    type: '@builtin',
-    value: 'vertex_index',
-  }) as BuiltinVertexIndex,
-  instanceIndex: attribute(u32, {
-    type: '@builtin',
-    value: 'instance_index',
-  }) as BuiltinInstanceIndex,
-  position: attribute(vec4f, {
-    type: '@builtin',
-    value: 'position',
-  }) as BuiltinPosition,
-  clipDistances: attribute(arrayOf(u32, 8), {
-    type: '@builtin',
-    value: 'clip_distances',
-  }) as BuiltinClipDistances,
-  frontFacing: attribute(f32, {
-    type: '@builtin',
-    value: 'front_facing',
-  }) as BuiltinFrontFacing,
-  fragDepth: attribute(f32, {
-    type: '@builtin',
-    value: 'frag_depth',
-  }) as BuiltinFragDepth,
-  sampleIndex: attribute(u32, {
-    type: '@builtin',
-    value: 'sample_index',
-  }) as BuiltinSampleIndex,
-  sampleMask: attribute(u32, {
-    type: '@builtin',
-    value: 'sample_mask',
-  }) as BuiltinSampleMask,
-  localInvocationId: attribute(vec3u, {
+function defineBuiltin<T extends Decorated | LooseDecorated>(
+  dataType: AnyWgslData,
+  value: T['attribs'][0] extends { value: infer TValue } ? TValue : never,
+): T {
+  return attribute(dataType, {
+    [$internal]: true,
     type: '@builtin',
-    value: 'local_invocation_id',
-  }) as BuiltinLocalInvocationId,
-  localInvocationIndex: attribute(u32, {
-    type: '@builtin',
-    value: 'local_invocation_index',
-  }) as BuiltinLocalInvocationIndex,
-  globalInvocationId: attribute(vec3u, {
-    type: '@builtin',
-    value: 'global_invocation_id',
-  }) as BuiltinGlobalInvocationId,
-  workgroupId: attribute(vec3u, {
-    type: '@builtin',
-    value: 'workgroup_id',
-  }) as BuiltinWorkgroupId,
-  numWorkgroups: attribute(vec3u, {
-    type: '@builtin',
-    value: 'num_workgroups',
-  }) as BuiltinNumWorkgroups,
-  subgroupInvocationId: attribute(u32, {
-    type: '@builtin',
-    value: 'subgroup_invocation_id',
-  }) as BuiltinSubgroupInvocationId,
-  subgroupSize: attribute(u32, {
-    type: '@builtin',
-    value: 'subgroup_size',
-  }) as BuiltinSubgroupSize,
+    // biome-ignore lint/suspicious/noExplicitAny: it's fine
+    value: value as any,
+  }) as T;
+}
+
+export const builtin = {
+  vertexIndex: defineBuiltin<BuiltinVertexIndex>(u32, 'vertex_index'),
+  instanceIndex: defineBuiltin<BuiltinInstanceIndex>(u32, 'instance_index'),
+  position: defineBuiltin<BuiltinPosition>(vec4f, 'position'),
+  clipDistances: defineBuiltin<BuiltinClipDistances>(
+    arrayOf(u32, 8),
+    'clip_distances',
+  ),
+  frontFacing: defineBuiltin<BuiltinFrontFacing>(f32, 'front_facing'),
+  fragDepth: defineBuiltin<BuiltinFragDepth>(f32, 'frag_depth'),
+  sampleIndex: defineBuiltin<BuiltinSampleIndex>(u32, 'sample_index'),
+  sampleMask: defineBuiltin<BuiltinSampleMask>(u32, 'sample_mask'),
+  localInvocationId: defineBuiltin<BuiltinLocalInvocationId>(
+    vec3u,
+    'local_invocation_id',
+  ),
+  localInvocationIndex: defineBuiltin<BuiltinLocalInvocationIndex>(
+    u32,
+    'local_invocation_index',
+  ),
+  globalInvocationId: defineBuiltin<BuiltinGlobalInvocationId>(
+    vec3u,
+    'global_invocation_id',
+  ),
+  workgroupId: defineBuiltin<BuiltinWorkgroupId>(vec3u, 'workgroup_id'),
+  numWorkgroups: defineBuiltin<BuiltinNumWorkgroups>(vec3u, 'num_workgroups'),
+  subgroupInvocationId: defineBuiltin<BuiltinSubgroupInvocationId>(
+    u32,
+    'subgroup_invocation_id',
+  ),
+  subgroupSize: defineBuiltin<BuiltinSubgroupSize>(u32, 'subgroup_size'),
 } as const;
 
 export type AnyBuiltin = (typeof builtin)[keyof typeof builtin];

packages/typegpu/src/core/function/fnCore.ts

@@ -34,7 +34,7 @@ export interface FnCore {
 
 export function createFnCore(
   shell: TgpuFnShellBase<unknown[] | Record<string, unknown>, unknown>,
-  implementation: Implementation<unknown[], unknown>,
+  implementation: Implementation,
 ): FnCore {
   /**
    * External application has to be deferred until resolution because

packages/typegpu/src/core/function/fnTypes.ts

@@ -43,10 +43,25 @@ export type InferReturn<T> = T extends undefined
     void
   : Infer<T>;
 
+export type JsImplementation<
+  Args extends unknown[] | Record<string, unknown> =
+    | unknown[]
+    | Record<string, unknown>,
+  Return = unknown,
+> = (
+  ...args: Args extends unknown[]
+    ? InferArgs<Args>
+    : Args extends Record<string, never>
+      ? []
+      : [InferIO<Args>]
+) => InferReturn<Return>;
+
 export type Implementation<
-  Args extends unknown[] = unknown[],
+  Args extends unknown[] | Record<string, unknown> =
+    | unknown[]
+    | Record<string, unknown>,
   Return = unknown,
-> = string | ((...args: Args) => Return);
+> = string | JsImplementation<Args, Return>;
 
 export type BaseIOData =
   | F32

packages/typegpu/src/core/function/ioOutputType.ts

@@ -2,7 +2,6 @@ import {
   type Decorate,
   type HasCustomLocation,
   type IsBuiltin,
-  attribute,
   location,
 } from '../../data/attributes.ts';
 import { isBuiltin } from '../../data/attributes.ts';
@@ -44,10 +43,7 @@ export function withLocations<T extends IOData>(
         return [key, member];
       }
 
-      return [
-        key,
-        attribute(member, { type: '@location', value: nextLocation++ }),
-      ];
+      return [key, location(nextLocation++, member)];
     }),
   );
 }

packages/typegpu/src/core/function/tgpuFn.ts

@@ -26,6 +26,7 @@ import type {
   InferArgs,
   InferIO,
   InferReturn,
+  JsImplementation,
 } from './fnTypes.ts';
 import { stripTemplate } from './templateUtils.ts';
 
@@ -80,6 +81,9 @@ interface TgpuFnBase<
   Return extends AnyWgslData | undefined = undefined,
 > extends TgpuNamable,
     Labelled {
+  readonly [$internal]: {
+    implementation: Implementation<Args, Return>;
+  };
   readonly resourceType: 'function';
   readonly shell: TgpuFnShellHeader<Args, Return>;
   readonly '~providing'?: Providing | undefined;
@@ -102,18 +106,7 @@ export type TgpuFn<
       : Args extends Record<string, never>
         ? []
         : [InferIO<Args>]
-  ) => InferReturn<Return>) & {
-    readonly [$internal]: {
-      implementation: Implementation<
-        Args extends AnyWgslData[]
-          ? InferArgs<Args>
-          : Args extends Record<string, never>
-            ? []
-            : [InferIO<Args>],
-        InferReturn<Return>
-      >;
-    };
-  };
+  ) => InferReturn<Return>);
 
 export function fn<
   Args extends AnyWgslData[] | Record<string, AnyWgslData> | [],
@@ -164,13 +157,16 @@ function createFn<
   Return extends AnyWgslData | undefined,
 >(
   shell: TgpuFnShellHeader<Args, Return>,
-  implementation: Implementation,
+  implementation: Implementation<Args, Return>,
 ): TgpuFn<Args, Return> {
   type This = TgpuFnBase<Args, Return> & SelfResolvable;
 
-  const core = createFnCore(shell, implementation);
+  const core = createFnCore(shell, implementation as Implementation);
 
   const fnBase: This = {
+    [$internal]: {
+      implementation,
+    },
     shell,
     resourceType: 'function' as const,
 
@@ -214,8 +210,8 @@ function createFn<
     },
   };
 
-  const call = createDualImpl(
-    (...args: unknown[]): unknown => {
+  const call = createDualImpl<JsImplementation<Args, Return>>(
+    (...args) => {
       if (typeof implementation === 'string') {
         throw new Error(
           'Cannot execute on the CPU functions constructed with raw WGSL',
@@ -229,9 +225,6 @@ function createFn<
       dataType: shell.returnType ?? UnknownData,
     }),
   );
-
-  call[$internal].implementation = implementation;
-
   const fn = Object.assign(call, fnBase as This) as unknown as TgpuFn<
     Args,
     Return
@@ -256,6 +249,9 @@ function createBoundFunction<
   type This = TgpuFnBase<Args, Return>;
 
   const fnBase: This = {
+    [$internal]: {
+      implementation: innerFn[$internal].implementation,
+    },
     resourceType: 'function',
     shell: innerFn.shell,
     '~providing': {

packages/typegpu/src/data/array.ts

@@ -5,6 +5,7 @@ import type {
   InferPartial,
   MemIdentity,
 } from '../shared/repr.ts';
+import { $internal } from '../shared/symbols.ts';
 import { sizeOf } from './sizeOf.ts';
 import type { AnyWgslData, BaseData, WgslArray } from './wgslTypes.ts';
 
@@ -35,18 +36,19 @@ export function arrayOf<TElement extends AnyWgslData>(
 // --------------
 
 class WgslArrayImpl<TElement extends BaseData> implements WgslArray<TElement> {
+  public readonly [$internal] = true;
   public readonly type = 'array';
   /** Type-token, not available at runtime */
   public declare readonly [$repr]: Infer<TElement>[];
   /** Type-token, not available at runtime */
-  public readonly '~gpuRepr'!: InferGPU<TElement>[];
+  public declare readonly '~gpuRepr': InferGPU<TElement>[];
   /** Type-token, not available at runtime */
-  public readonly '~reprPartial'!: {
+  public declare readonly '~reprPartial': {
     idx: number;
     value: InferPartial<TElement>;
   }[];
   /** Type-token, not available at runtime */
-  public readonly '~memIdent'!: WgslArray<MemIdentity<TElement>>;
+  public declare readonly '~memIdent': WgslArray<MemIdentity<TElement>>;
 
   constructor(
     public readonly elementType: TElement,

packages/typegpu/src/data/atomic.ts

@@ -1,4 +1,5 @@
 import type { $repr, Infer, MemIdentity } from '../shared/repr.ts';
+import { $internal } from '../shared/symbols.ts';
 import type { Atomic, I32, U32, atomicI32, atomicU32 } from './wgslTypes.ts';
 
 // ----------
@@ -25,6 +26,7 @@ export function atomic<TSchema extends U32 | I32>(
 // --------------
 
 class AtomicImpl<TSchema extends U32 | I32> implements Atomic<TSchema> {
+  public readonly [$internal] = true;
   public readonly type = 'atomic';
   /** Type-token, not available at runtime */
   public declare readonly [$repr]: Infer<TSchema>;