feat: Variance annotation overrides to achieve better type behavior (#2073)

Review fixes

Author: iwoplaza

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

@@ -31,7 +31,7 @@ interface TgpuBufferShorthandBase<TData extends BaseData> extends TgpuNamable {
   // ---
 }
 
-export interface TgpuMutable<TData extends BaseData>
+export interface TgpuMutable<out TData extends BaseData>
   extends TgpuBufferShorthandBase<TData> {
   readonly resourceType: 'mutable';
   readonly buffer: TgpuBuffer<TData> & StorageFlag;
@@ -45,7 +45,7 @@ export interface TgpuMutable<TData extends BaseData>
   // ---
 }
 
-export interface TgpuReadonly<TData extends BaseData>
+export interface TgpuReadonly<out TData extends BaseData>
   extends TgpuBufferShorthandBase<TData> {
   readonly resourceType: 'readonly';
   readonly buffer: TgpuBuffer<TData> & StorageFlag;
@@ -59,7 +59,7 @@ export interface TgpuReadonly<TData extends BaseData>
   // ---
 }
 
-export interface TgpuUniform<TData extends BaseData>
+export interface TgpuUniform<out TData extends BaseData>
   extends TgpuBufferShorthandBase<TData> {
   readonly resourceType: 'uniform';
   readonly buffer: TgpuBuffer<TData> & UniformFlag;

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

@@ -237,7 +237,7 @@ class TgpuFixedBufferImpl<
 }
 
 export class TgpuLaidOutBufferImpl<
-  TData extends BaseData,
+  TData extends AnyWgslData,
   TUsage extends BindableBufferUsage,
 > implements TgpuBufferUsage<TData, TUsage>, SelfResolvable {
   /** Type-token, not available at runtime */
@@ -257,21 +257,20 @@ export class TgpuLaidOutBufferImpl<
   }
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
-    const dataType = this.dataType as unknown as AnyData;
     const id = ctx.getUniqueName(this);
     const group = ctx.allocateLayoutEntry(this.#membership.layout);
     const usage = usageToVarTemplateMap[this.usage];
 
     ctx.addDeclaration(
       `@group(${group}) @binding(${this.#membership.idx}) var<${usage}> ${id}: ${
-        ctx.resolve(dataType).value
+        ctx.resolve(this.dataType).value
       };`,
     );
 
     return snip(
       id,
-      dataType,
-      isNaturallyEphemeral(dataType) ? 'runtime' : this.usage,
+      this.dataType,
+      isNaturallyEphemeral(this.dataType) ? 'runtime' : this.usage,
     );
   }
 
@@ -280,7 +279,7 @@ export class TgpuLaidOutBufferImpl<
   }
 
   get [$gpuValueOf](): InferGPU<TData> {
-    const schema = this.dataType as AnyData;
+    const schema = this.dataType;
     const usage = this.usage;
 
     return new Proxy({

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

@@ -54,7 +54,7 @@ export function dualImpl<T extends (...args: never[]) => unknown>(
 
     const argSnippets = args as MapValueToSnippet<Parameters<T>>;
     const converted = argSnippets.map((s, idx) => {
-      const argType = argTypes[idx] as AnyData | undefined;
+      const argType = argTypes[idx];
       if (!argType) {
         throw new Error('Function called with invalid arguments');
       }

packages/typegpu/src/core/resolve/resolveData.ts

@@ -16,7 +16,6 @@ import {
 import { formatToWGSLType } from '../../data/vertexFormatData.ts';
 import type {
   AnyWgslData,
-  BaseData,
   Bool,
   F16,
   F32,
@@ -119,10 +118,10 @@ function isIdentityType(data: AnyWgslData): data is IdentityType {
  */
 function resolveStructProperty(
   ctx: ResolutionCtx,
-  [key, property]: [string, BaseData],
+  [key, property]: [string, AnyData],
 ) {
   return `  ${getAttributesString(property)}${key}: ${
-    ctx.resolve(property as AnyWgslData).value
+    ctx.resolve(property).value
   },\n`;
 }
 
@@ -142,7 +141,7 @@ function resolveStruct(ctx: ResolutionCtx, struct: WgslStruct) {
   ctx.addDeclaration(`\
 struct ${id} {
 ${
-    Object.entries(struct.propTypes as Record<string, BaseData>)
+    Object.entries(struct.propTypes)
       .map((prop) => resolveStructProperty(ctx, prop))
       .join('')
   }\
@@ -172,7 +171,7 @@ function resolveUnstruct(ctx: ResolutionCtx, unstruct: Unstruct) {
   ctx.addDeclaration(`\
 struct ${id} {
 ${
-    Object.entries(unstruct.propTypes as Record<string, BaseData>)
+    Object.entries(unstruct.propTypes)
       .map((prop) =>
         isAttribute(prop[1])
           ? resolveStructProperty(ctx, [

packages/typegpu/src/core/root/init.ts

@@ -11,7 +11,6 @@ import {
 import type { AnyData, Disarray } from '../../data/dataTypes.ts';
 import type {
   AnyWgslData,
-  BaseData,
   U16,
   U32,
   v3u,
@@ -657,7 +656,7 @@ class TgpuRootImpl extends WithBindingImpl
       TgpuVertexLayout,
       {
         buffer:
-          | (TgpuBuffer<WgslArray<BaseData> | Disarray<BaseData>> & VertexFlag)
+          | (TgpuBuffer<WgslArray | Disarray> & VertexFlag)
           | GPUBuffer;
         offset?: number | undefined;
         size?: number | undefined;

packages/typegpu/src/data/attributes.ts

@@ -25,7 +25,6 @@ import {
   isLooseData,
   isLooseDecorated,
   type LooseDecorated,
-  type LooseTypeLiteral,
   type Undecorate,
 } from './dataTypes.ts';
 import { sizeOf } from './sizeOf.ts';
@@ -50,7 +49,6 @@ import {
   type PerspectiveOrLinearInterpolationType,
   type Size,
   type Vec4f,
-  type WgslTypeLiteral,
 } from './wgslTypes.ts';
 
 // ----------
@@ -113,9 +111,9 @@ export type ExtractAttributes<T> = T extends {
 export type Decorate<
   TData extends BaseData,
   TAttrib extends AnyAttribute,
-> = TData['type'] extends WgslTypeLiteral
+> = TData extends AnyWgslData
   ? Decorated<Undecorate<TData>, [TAttrib, ...ExtractAttributes<TData>]>
-  : TData['type'] extends LooseTypeLiteral
+  : TData extends AnyLooseData
     ? LooseDecorated<Undecorate<TData>, [TAttrib, ...ExtractAttributes<TData>]>
   : never;
 

packages/typegpu/src/data/dataTypes.ts

@@ -38,7 +38,7 @@ import type { WgslComparisonSampler, WgslSampler } from './sampler.ts';
  * unless they are explicitly decorated with the custom align attribute
  * via `d.align` function.
  */
-export interface Disarray<TElement extends wgsl.BaseData = wgsl.BaseData>
+export interface Disarray<out TElement extends wgsl.BaseData = wgsl.BaseData>
   extends wgsl.BaseData {
   <T extends TElement>(elements: Infer<T>[]): Infer<T>[];
   (): Infer<TElement>[];
@@ -66,8 +66,8 @@ export interface Disarray<TElement extends wgsl.BaseData = wgsl.BaseData>
  * via `d.align` function.
  */
 export interface Unstruct<
-  // biome-ignore lint/suspicious/noExplicitAny: the widest type that works with both covariance and contravariance
-  TProps extends Record<string, wgsl.BaseData> = any,
+  // @ts-expect-error: Override variance, as we want unstructs to behave like objects
+  out TProps extends Record<string, AnyData> = Record<string, AnyData>,
 > extends wgsl.BaseData, TgpuNamable {
   (props: Prettify<InferRecord<TProps>>): Prettify<InferRecord<TProps>>;
   (): Prettify<InferRecord<TProps>>;
@@ -93,8 +93,8 @@ export interface Unstruct<
 export type AnyUnstruct = Unstruct;
 
 export interface LooseDecorated<
-  TInner extends wgsl.BaseData = wgsl.BaseData,
-  TAttribs extends unknown[] = unknown[],
+  out TInner extends wgsl.BaseData = wgsl.BaseData,
+  out TAttribs extends unknown[] = unknown[],
 > extends wgsl.BaseData {
   readonly type: 'loose-decorated';
   readonly inner: TInner;

packages/typegpu/src/data/deepEqual.ts

@@ -55,7 +55,8 @@ export function deepEqual(a: AnyData, b: AnyData): boolean {
       const keyB = bKeys[i];
       if (
         keyA !== keyB || !keyA || !keyB ||
-        !deepEqual(aProps[keyA], bProps[keyB])
+        // biome-ignore lint/style/noNonNullAssertion: they exist
+        !deepEqual(aProps[keyA]!, bProps[keyB]!)
       ) {
         return false;
       }
@@ -74,12 +75,12 @@ export function deepEqual(a: AnyData, b: AnyData): boolean {
     return (
       a.addressSpace === b.addressSpace &&
       a.access === b.access &&
-      deepEqual(a.inner as AnyData, b.inner as AnyData)
+      deepEqual(a.inner, b.inner)
     );
   }
 
   if (isAtomic(a) && isAtomic(b)) {
-    return deepEqual(a.inner as AnyData, b.inner as AnyData);
+    return deepEqual(a.inner, b.inner);
   }
 
   if (

packages/typegpu/src/data/offsets.ts

@@ -4,7 +4,7 @@ import alignIO from './alignIO.ts';
 import { alignmentOf, customAlignmentOf } from './alignmentOf.ts';
 import { isUnstruct, type Unstruct } from './dataTypes.ts';
 import { sizeOf } from './sizeOf.ts';
-import type { BaseData, WgslStruct } from './wgslTypes.ts';
+import type { WgslStruct } from './wgslTypes.ts';
 
 export interface OffsetInfo {
   offset: number;
@@ -17,16 +17,18 @@ const cachedOffsets = new WeakMap<
   Record<string, OffsetInfo>
 >();
 
-export function offsetsForProps<T extends Record<string, BaseData>>(
-  struct: WgslStruct<T> | Unstruct<T>,
-): Record<keyof T, OffsetInfo> {
+export function offsetsForProps<T extends WgslStruct | Unstruct>(
+  struct: T,
+): Record<keyof T['propTypes'], OffsetInfo> {
+  type Key = keyof T['propTypes'];
+
   const cached = cachedOffsets.get(struct);
   if (cached) {
-    return cached as Record<keyof T, OffsetInfo>;
+    return cached as Record<Key, OffsetInfo>;
   }
 
   const measurer = new Measurer();
-  const offsets = {} as Record<keyof T, OffsetInfo>;
+  const offsets = {} as Record<Key, OffsetInfo>;
   let lastEntry: OffsetInfo | undefined;
 
   for (const key in struct.propTypes) {
@@ -47,7 +49,7 @@ export function offsetsForProps<T extends Record<string, BaseData>>(
     }
 
     const propSize = sizeOf(prop);
-    offsets[key] = { offset: measurer.size, size: propSize };
+    offsets[key as Key] = { offset: measurer.size, size: propSize };
     lastEntry = offsets[key];
     measurer.add(propSize);
   }
@@ -57,12 +59,6 @@ export function offsetsForProps<T extends Record<string, BaseData>>(
       measurer.size;
   }
 
-  cachedOffsets.set(
-    struct as
-      | WgslStruct<Record<string, BaseData>>
-      | Unstruct<Record<string, BaseData>>,
-    offsets,
-  );
-
+  cachedOffsets.set(struct, offsets);
   return offsets;
 }

packages/typegpu/src/data/struct.ts

@@ -1,9 +1,8 @@
 import { isValidProp } from '../nameRegistry.ts';
 import { getName, setName } from '../shared/meta.ts';
 import { $internal } from '../shared/symbols.ts';
-import type { AnyData } from './dataTypes.ts';
 import { schemaCallWrapper } from './schemaCallWrapper.ts';
-import type { AnyWgslData, BaseData, WgslStruct } from './wgslTypes.ts';
+import type { AnyWgslData, WgslStruct } from './wgslTypes.ts';
 
 // ----------
 // Public API
@@ -26,7 +25,7 @@ export function struct<TProps extends Record<string, AnyWgslData>>(
   return INTERNAL_createStruct(props, false);
 }
 
-export function abstruct<TProps extends Record<string, BaseData>>(
+export function abstruct<TProps extends Record<string, AnyWgslData>>(
   props: TProps,
 ): WgslStruct<TProps> {
   return INTERNAL_createStruct(props, true);
@@ -36,7 +35,7 @@ export function abstruct<TProps extends Record<string, BaseData>>(
 // Implementation
 // --------------
 
-function INTERNAL_createStruct<TProps extends Record<string, BaseData>>(
+function INTERNAL_createStruct<TProps extends Record<string, AnyWgslData>>(
   props: TProps,
   isAbstruct: boolean,
 ): WgslStruct<TProps> {
@@ -54,7 +53,7 @@ function INTERNAL_createStruct<TProps extends Record<string, BaseData>>(
     Object.fromEntries(
       Object.entries(props).map(([key, schema]) => [
         key,
-        schemaCallWrapper(schema as AnyData, instanceProps?.[key]),
+        schemaCallWrapper(schema, instanceProps?.[key]),
       ]),
     );