refactor: Double-down on BaseData (#2092)

Author: iwoplaza

packages/typegpu/src/builtin.ts

@@ -4,7 +4,6 @@ import type { LooseDecorated } from './data/dataTypes.ts';
 import { bool, f32, u32 } from './data/numeric.ts';
 import { vec3u, vec4f } from './data/vector.ts';
 import type {
-  AnyWgslData,
   BaseData,
   Bool,
   Builtin,
@@ -62,7 +61,7 @@ export type BuiltinSubgroupId = Decorated<U32, [Builtin<'subgroup_id'>]>;
 export type BuiltinNumSubgroups = Decorated<U32, [Builtin<'num_subgroups'>]>;
 
 function defineBuiltin<T extends Decorated | LooseDecorated>(
-  dataType: AnyWgslData,
+  dataType: BaseData,
   value: T['attribs'][0] extends { params: [infer TValue] } ? TValue : never,
 ): T {
   return attribute(dataType, {

packages/typegpu/src/core/buffer/buffer.ts

@@ -144,30 +144,30 @@ export function INTERNAL_createBuffer<TData extends AnyData>(
   return new TgpuBufferImpl(group, typeSchema, initialOrBuffer);
 }
 
-export function isBuffer<T extends TgpuBuffer<AnyData>>(
+export function isBuffer<T extends TgpuBuffer<BaseData>>(
   value: T | unknown,
 ): value is T {
-  return (value as TgpuBuffer<AnyData>).resourceType === 'buffer';
+  return (value as T).resourceType === 'buffer';
 }
 
-export function isUsableAsVertex<T extends TgpuBuffer<AnyData>>(
+export function isUsableAsVertex<T extends TgpuBuffer<BaseData>>(
   buffer: T,
 ): buffer is T & VertexFlag {
-  return !!(buffer as unknown as VertexFlag).usableAsVertex;
+  return !!(buffer as T).usableAsVertex;
 }
 
-export function isUsableAsIndex<T extends TgpuBuffer<AnyData>>(
+export function isUsableAsIndex<T extends TgpuBuffer<BaseData>>(
   buffer: T,
 ): buffer is T & IndexFlag {
-  return !!(buffer as unknown as IndexFlag).usableAsIndex;
+  return !!(buffer as T).usableAsIndex;
 }
 
 // --------------
 // Implementation
 // --------------
 const endianness = getSystemEndianness();
 
-class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
+class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
   public readonly [$internal] = true;
   public readonly resourceType = 'buffer';
   public flags: GPUBufferUsageFlags = GPUBufferUsage.COPY_DST |

packages/typegpu/src/core/buffer/bufferUsage.ts

@@ -1,4 +1,3 @@
-import type { AnyData } from '../../data/dataTypes.ts';
 import { schemaCallWrapper } from '../../data/schemaCallWrapper.ts';
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
 import {
@@ -80,10 +79,10 @@ export interface TgpuFixedBufferUsage<TData extends BaseData>
 export interface TgpuBufferMutable<TData extends BaseData>
   extends TgpuBufferUsage<TData, 'mutable'> {}
 
-export function isUsableAsUniform<T extends TgpuBuffer<AnyData>>(
+export function isUsableAsUniform<T extends TgpuBuffer<BaseData>>(
   buffer: T,
 ): buffer is T & UniformFlag {
-  return !!(buffer as unknown as UniformFlag).usableAsUniform;
+  return !!(buffer as T).usableAsUniform;
 }
 
 // --------------
@@ -97,7 +96,7 @@ const usageToVarTemplateMap: Record<BindableBufferUsage, string> = {
 };
 
 class TgpuFixedBufferImpl<
-  TData extends AnyWgslData,
+  TData extends BaseData,
   TUsage extends BindableBufferUsage,
 > implements
   TgpuBufferUsage<TData, TUsage>,
@@ -237,7 +236,7 @@ class TgpuFixedBufferImpl<
 }
 
 export class TgpuLaidOutBufferImpl<
-  TData extends AnyWgslData,
+  TData extends BaseData,
   TUsage extends BindableBufferUsage,
 > implements TgpuBufferUsage<TData, TUsage>, SelfResolvable {
   /** Type-token, not available at runtime */
@@ -312,8 +311,8 @@ export class TgpuLaidOutBufferImpl<
 }
 
 const mutableUsageMap = new WeakMap<
-  TgpuBuffer<AnyWgslData>,
-  TgpuFixedBufferImpl<AnyWgslData, 'mutable'>
+  TgpuBuffer<BaseData>,
+  TgpuFixedBufferImpl<BaseData, 'mutable'>
 >();
 
 export function mutable<TData extends AnyWgslData>(
@@ -336,8 +335,8 @@ export function mutable<TData extends AnyWgslData>(
 }
 
 const readonlyUsageMap = new WeakMap<
-  TgpuBuffer<AnyWgslData>,
-  TgpuFixedBufferImpl<AnyWgslData, 'readonly'>
+  TgpuBuffer<BaseData>,
+  TgpuFixedBufferImpl<BaseData, 'readonly'>
 >();
 
 export function readonly<TData extends AnyWgslData>(
@@ -360,8 +359,8 @@ export function readonly<TData extends AnyWgslData>(
 }
 
 const uniformUsageMap = new WeakMap<
-  TgpuBuffer<AnyWgslData>,
-  TgpuFixedBufferImpl<AnyWgslData, 'uniform'>
+  TgpuBuffer<BaseData>,
+  TgpuFixedBufferImpl<BaseData, 'uniform'>
 >();
 
 export function uniform<TData extends AnyWgslData>(

packages/typegpu/src/core/constant/tgpuConstant.ts

@@ -1,6 +1,6 @@
 import type { AnyData } from '../../data/dataTypes.ts';
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
-import { isNaturallyEphemeral } from '../../data/wgslTypes.ts';
+import { type BaseData, isNaturallyEphemeral } from '../../data/wgslTypes.ts';
 import { inCodegenMode } from '../../execMode.ts';
 import type { TgpuNamable } from '../../shared/meta.ts';
 import { getName, setName } from '../../shared/meta.ts';
@@ -24,7 +24,7 @@ type DeepReadonly<T> = T extends { [$internal]: unknown } ? T
     ? { readonly [K in keyof T]: DeepReadonly<T[K]> }
   : T;
 
-export interface TgpuConst<TDataType extends AnyData = AnyData>
+export interface TgpuConst<TDataType extends BaseData = BaseData>
   extends TgpuNamable {
   readonly resourceType: 'const';
   readonly [$gpuValueOf]: DeepReadonly<InferGPU<TDataType>>;
@@ -71,7 +71,7 @@ function deepFreeze<T extends object>(object: T): T {
   return Object.freeze(object);
 }
 
-class TgpuConstImpl<TDataType extends AnyData>
+class TgpuConstImpl<TDataType extends BaseData>
   implements TgpuConst<TDataType>, SelfResolvable {
   readonly [$internal] = {};
   readonly resourceType: 'const';

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

@@ -1,4 +1,3 @@
-import type { AnyData } from '../../data/dataTypes.ts';
 import { type MapValueToSnippet, snip } from '../../data/snippet.ts';
 import { setName } from '../../shared/meta.ts';
 import { $gpuCallable } from '../../shared/symbols.ts';
@@ -9,8 +8,9 @@ import {
   NormalState,
   type ResolutionCtx,
 } from '../../types.ts';
+import { type BaseData, isPtr } from '../../data/wgslTypes.ts';
 
-type MapValueToDataType<T> = { [K in keyof T]: AnyData };
+type MapValueToDataType<T> = { [K in keyof T]: BaseData };
 type AnyFn = (...args: never[]) => unknown;
 
 interface DualImplOptions<T extends AnyFn> {
@@ -21,10 +21,10 @@ interface DualImplOptions<T extends AnyFn> {
     args: MapValueToSnippet<Parameters<T>>,
   ) => string;
   readonly signature:
-    | { argTypes: AnyData[]; returnType: AnyData }
+    | { argTypes: BaseData[]; returnType: BaseData }
     | ((
       ...inArgTypes: MapValueToDataType<Parameters<T>>
-    ) => { argTypes: AnyData[]; returnType: AnyData });
+    ) => { argTypes: BaseData[]; returnType: BaseData });
   /**
    * Whether the function should skip trying to execute the "normal" implementation if
    * all arguments are known at compile time.
@@ -64,7 +64,7 @@ export function dualImpl<T extends AnyFn>(
         ? options.signature(
           ...args.map((s) => {
             // Dereference implicit pointers
-            if (s.dataType.type === 'ptr' && s.dataType.implicit) {
+            if (isPtr(s.dataType) && s.dataType.implicit) {
               return s.dataType.inner;
             }
             return s.dataType;

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

@@ -1,7 +1,12 @@
 import { getAttributesString } from '../../data/attributes.ts';
-import { type AnyData, undecorate } from '../../data/dataTypes.ts';
+import { undecorate } from '../../data/dataTypes.ts';
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
-import { isWgslData, isWgslStruct, Void } from '../../data/wgslTypes.ts';
+import {
+  type BaseData,
+  isWgslData,
+  isWgslStruct,
+  Void,
+} from '../../data/wgslTypes.ts';
 import { MissingLinksError } from '../../errors.ts';
 import { getMetaData, getName, setName } from '../../shared/meta.ts';
 import type { ResolutionCtx } from '../../types.ts';
@@ -17,12 +22,12 @@ export interface FnCore {
   applyExternals(newExternals: ExternalMap): void;
   resolve(
     ctx: ResolutionCtx,
-    argTypes: AnyData[],
+    argTypes: BaseData[],
     /**
      * The return type of the function. If undefined, the type should be inferred
      * from the implementation (relevant for shellless functions).
      */
-    returnType: AnyData | undefined,
+    returnType: BaseData | undefined,
   ): ResolvedSnippet;
 }
 
@@ -45,8 +50,8 @@ export function createFnCore(
 
     resolve(
       ctx: ResolutionCtx,
-      argTypes: AnyData[],
-      returnType: AnyData | undefined,
+      argTypes: BaseData[],
+      returnType: BaseData | undefined,
     ): ResolvedSnippet {
       const externalMap: ExternalMap = {};
 

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

@@ -1,5 +1,5 @@
-import type { AnyData } from '../../data/dataTypes.ts';
 import type { ResolvedSnippet } from '../../data/snippet.ts';
+import type { BaseData } from '../../data/wgslTypes.ts';
 import { getName } from '../../shared/meta.ts';
 import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
@@ -23,12 +23,12 @@ import { createFnCore } from './fnCore.ts';
  */
 export interface ShelllessImpl extends SelfResolvable {
   readonly resourceType: 'shellless-impl';
-  readonly argTypes: AnyData[];
+  readonly argTypes: BaseData[];
   readonly [$getNameForward]: unknown;
 }
 
 export function createShelllessImpl(
-  argTypes: AnyData[],
+  argTypes: BaseData[],
   implementation: (...args: never[]) => unknown,
 ): ShelllessImpl {
   const core = createFnCore(implementation, '');

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

@@ -1,7 +1,6 @@
-import type { AnyData } from '../../data/dataTypes.ts';
 import type { ResolvedSnippet } from '../../data/snippet.ts';
 import { schemaCallWrapper } from '../../data/schemaCallWrapper.ts';
-import { Void } from '../../data/wgslTypes.ts';
+import { type BaseData, Void } from '../../data/wgslTypes.ts';
 import { ExecutionError } from '../../errors.ts';
 import { provideInsideTgpuFn } from '../../execMode.ts';
 import type { TgpuNamable } from '../../shared/meta.ts';
@@ -51,8 +50,8 @@ import type { Withable } from '../root/rootTypes.ts';
  * Describes a function signature (its arguments and return type)
  */
 type TgpuFnShellHeader<
-  Args extends AnyData[],
-  Return extends AnyData,
+  Args extends BaseData[],
+  Return extends BaseData,
 > = {
   readonly [$internal]: true;
   readonly argTypes: Args;
@@ -66,8 +65,8 @@ type TgpuFnShellHeader<
  * and passing the implementation (as WGSL string or JS function) as the argument.
  */
 export type TgpuFnShell<
-  Args extends AnyData[],
-  Return extends AnyData,
+  Args extends BaseData[],
+  Return extends BaseData,
 > =
   & TgpuFnShellHeader<Args, Return>
   & (<T extends (...args: InferArgs<Args>) => Infer<Return>>(
@@ -89,7 +88,7 @@ interface TgpuFnBase<ImplSchema extends AnyFn>
   readonly resourceType: 'function';
   readonly shell: TgpuFnShellHeader<
     Parameters<ImplSchema>,
-    Extract<ReturnType<ImplSchema>, AnyData>
+    Extract<ReturnType<ImplSchema>, BaseData>
   >;
   readonly [$providing]?: Providing | undefined;
 
@@ -102,22 +101,22 @@ export type TgpuFn<ImplSchema extends AnyFn = (...args: any[]) => any> =
   & TgpuFnBase<ImplSchema>;
 
 export function fn<
-  Args extends AnyData[] | [],
+  Args extends BaseData[] | [],
 >(argTypes: Args, returnType?: undefined): TgpuFnShell<Args, Void>;
 
 export function fn<
-  Args extends AnyData[] | [],
-  Return extends AnyData,
+  Args extends BaseData[] | [],
+  Return extends BaseData,
 >(argTypes: Args, returnType: Return): TgpuFnShell<Args, Return>;
 
 export function fn<
-  Args extends AnyData[] | [],
-  Return extends AnyData = Void,
+  Args extends BaseData[] | [],
+  Return extends BaseData = Void,
 >(argTypes: Args, returnType?: Return | undefined): TgpuFnShell<Args, Return> {
   const shell: TgpuFnShellHeader<Args, Return> = {
     [$internal]: true,
     argTypes,
-    returnType: returnType ?? Void as Return,
+    returnType: returnType ?? Void as unknown as Return,
     isEntry: false,
   };
 
@@ -133,7 +132,7 @@ export function fn<
   return Object.assign(call, shell) as unknown as TgpuFnShell<Args, Return>;
 }
 
-export function isTgpuFn<Args extends AnyData[] | [], Return extends AnyData>(
+export function isTgpuFn<Args extends BaseData[] | [], Return extends BaseData>(
   value: unknown | TgpuFn<(...args: Args) => Return>,
 ): value is TgpuFn<(...args: Args) => Return> {
   return isMarkedInternal(value) &&
@@ -151,7 +150,7 @@ function stringifyPair([slot, value]: SlotValuePair): string {
 function createFn<ImplSchema extends AnyFn>(
   shell: TgpuFnShellHeader<
     Parameters<ImplSchema>,
-    Extract<ReturnType<ImplSchema>, AnyData>
+    Extract<ReturnType<ImplSchema>, BaseData>
   >,
   implementation: Implementation<ImplSchema>,
 ): TgpuFn<ImplSchema> {
@@ -217,7 +216,7 @@ function createFn<ImplSchema extends AnyFn>(
           }
 
           const castAndCopiedArgs = args.map((arg, index) =>
-            schemaCallWrapper(shell.argTypes[index] as unknown as AnyData, arg)
+            schemaCallWrapper(shell.argTypes[index] as unknown as BaseData, arg)
           ) as InferArgs<Parameters<ImplSchema>>;
 
           const result = implementation(...castAndCopiedArgs);