fix: Rework how schemas propagate usage validity (#1547)

* Rework how schemas propagate usage validity

* Fix types errors

* Lenient for generic cases

* Fix tests

* Verify buffer shorthands as well

* Back to square one

* Fix index validity checking

* Fix invalid inference in buffer shorthands

* Made Unstruct wider, just like WgslStruct

* Not using deprecated types internally

* Update buffer shorthand defs

* Cleanup

* Unused import

* Fix merge issue

* No reason to prefix in case of decorated

* Review fixes

Author: iwoplaza

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

@@ -3,14 +3,25 @@ import { getCompiledWriterForSchema } from '../../data/compiledIO.ts';
 import { readData, writeData } from '../../data/dataIO.ts';
 import { getWriteInstructions } from '../../data/partialIO.ts';
 import { sizeOf } from '../../data/sizeOf.ts';
-import type { BaseData, WgslTypeLiteral } from '../../data/wgslTypes.ts';
+import type { BaseData } from '../../data/wgslTypes.ts';
 import { isWgslData } from '../../data/wgslTypes.ts';
 import type { StorageFlag } from '../../extension.ts';
 import type { TgpuNamable } from '../../shared/meta.ts';
 import { getName, setName } from '../../shared/meta.ts';
-import type { Infer, InferPartial, MemIdentity } from '../../shared/repr.ts';
+import type {
+  Infer,
+  InferPartial,
+  IsValidIndexSchema,
+  IsValidStorageSchema,
+  IsValidUniformSchema,
+  IsValidVertexSchema,
+  MemIdentity,
+} from '../../shared/repr.ts';
 import { $internal } from '../../shared/symbols.ts';
-import type { UnionToIntersection } from '../../shared/utilityTypes.ts';
+import type {
+  Prettify,
+  UnionToIntersection,
+} from '../../shared/utilityTypes.ts';
 import { isGPUBuffer } from '../../types.ts';
 import type { ExperimentalTgpuRoot } from '../root/rootTypes.ts';
 import {
@@ -23,7 +34,6 @@ import {
   type TgpuFixedBufferUsage,
 } from './bufferUsage.ts';
 import type { AnyData } from '../../data/dataTypes.ts';
-import type { Undecorate } from '../../data/decorateUtils.ts';
 
 // ----------
 // Public API
@@ -75,15 +85,18 @@ const usageToUsageConstructor = {
   readonly: asReadonly,
 };
 
-type IsIndexCompatible<TData extends BaseData> = Undecorate<TData> extends {
-  readonly type: 'array';
-  readonly elementType: infer TElement;
-}
-  ? TElement extends BaseData
-    ? Undecorate<TElement> extends { readonly type: 'u32' | 'u16' } ? true
-    : false
-  : false
-  : false;
+/**
+ * Done as an object to later Prettify it
+ */
+type InnerValidUsagesFor<T> = {
+  usage:
+    | (IsValidStorageSchema<T> extends true ? 'storage' : never)
+    | (IsValidUniformSchema<T> extends true ? 'uniform' : never)
+    | (IsValidVertexSchema<T> extends true ? 'vertex' : never)
+    | (IsValidIndexSchema<T> extends true ? 'index' : never);
+};
+
+export type ValidUsagesFor<T> = InnerValidUsagesFor<T>['usage'];
 
 export interface TgpuBuffer<TData extends BaseData> extends TgpuNamable {
   readonly [$internal]: true;
@@ -99,7 +112,12 @@ export interface TgpuBuffer<TData extends BaseData> extends TgpuNamable {
   usableAsVertex: boolean;
   usableAsIndex: boolean;
 
-  $usage<T extends RestrictUsages<TData>>(
+  $usage<
+    T extends [
+      Prettify<InnerValidUsagesFor<TData>>['usage'],
+      ...Prettify<InnerValidUsagesFor<TData>>['usage'][],
+    ],
+  >(
     ...usages: T
   ): this & UnionToIntersection<LiteralToUsageType<T[number]>>;
   $addFlags(flags: GPUBufferUsageFlags): this;
@@ -152,30 +170,6 @@ export function isUsableAsIndex<T extends TgpuBuffer<AnyData>>(
 // --------------
 const endianness = getSystemEndianness();
 
-type IsArrayOfU32<TData extends BaseData> = Undecorate<TData> extends {
-  readonly type: 'array';
-  readonly elementType: infer TElement;
-}
-  ? TElement extends BaseData
-    ? Undecorate<TElement> extends { readonly type: 'u32' } ? true
-    : false
-  : false
-  : false;
-
-type IsWgslLiteral<TData extends BaseData> = TData extends {
-  readonly type: WgslTypeLiteral;
-} ? true
-  : false;
-
-type RestrictUsages<TData extends BaseData> = string extends TData['type']
-  ? ('uniform' | 'storage' | 'vertex' | 'index')[]
-  : IsIndexCompatible<TData> extends true
-    ? IsArrayOfU32<TData> extends true
-      ? ('uniform' | 'storage' | 'vertex' | 'index')[]
-    : ['index']
-  : IsWgslLiteral<TData> extends true ? ('uniform' | 'storage' | 'vertex')[]
-  : ['vertex'];
-
 class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
   public readonly [$internal] = true;
   public readonly resourceType = 'buffer';

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

@@ -235,7 +235,7 @@ export class TgpuLaidOutBufferImpl<
 
     ctx.addDeclaration(
       `@group(${group}) @binding(${this._membership.idx}) var<${usage}> ${id}: ${
-        ctx.resolve(this.dataType as AnyWgslData)
+        ctx.resolve(this.dataType)
       };`,
     );
 

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

@@ -4,10 +4,10 @@ import { undecorate } from '../../data/decorateUtils.ts';
 import type { AnyData } from '../../data/dataTypes.ts';
 import { snip } from '../../data/snippet.ts';
 import {
-  type AnyWgslStruct,
   isWgslData,
   isWgslStruct,
   Void,
+  type WgslStruct,
 } from '../../data/wgslTypes.ts';
 import { MissingLinksError } from '../../errors.ts';
 import { getMetaData, getName, setName } from '../../shared/meta.ts';
@@ -182,7 +182,7 @@ export function createFnCore(
                   alias,
                   snip(
                     `_arg_${i}.${name}`,
-                    (argTypes[i] as AnyWgslStruct)
+                    (argTypes[i] as WgslStruct)
                       .propTypes[name],
                   ),
                 ])

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

@@ -129,7 +129,7 @@ function resolveStruct(ctx: ResolutionCtx, struct: WgslStruct) {
   ctx.addDeclaration(`\
 struct ${id} {
 ${
-    Object.entries(struct.propTypes)
+    Object.entries(struct.propTypes as Record<string, BaseData>)
       .map((prop) => resolveStructProperty(ctx, prop))
       .join('')
   }\
@@ -159,7 +159,7 @@ function resolveUnstruct(ctx: ResolutionCtx, unstruct: Unstruct) {
   ctx.addDeclaration(`\
 struct ${id} {
 ${
-    Object.entries(unstruct.propTypes)
+    Object.entries(unstruct.propTypes as Record<string, BaseData>)
       .map((prop) =>
         isAttribute(prop[1])
           ? resolveStructProperty(ctx, [

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

@@ -1,10 +1,16 @@
 import type { AnyComputeBuiltin, OmitBuiltins } from '../../builtin.ts';
 import type { TgpuQuerySet } from '../../core/querySet/querySet.ts';
+import type { AnyData, Disarray } from '../../data/dataTypes.ts';
 import type { UndecorateRecord } from '../../data/decorateUtils.ts';
-import type { AnyData, Disarray, HasNestedType } from '../../data/dataTypes.ts';
 import type { AnyWgslData, U16, U32, WgslArray } from '../../data/wgslTypes.ts';
 import type { NameRegistry } from '../../nameRegistry.ts';
-import type { Infer } from '../../shared/repr.ts';
+import type {
+  ExtractInvalidSchemaError,
+  Infer,
+  IsValidBufferSchema,
+  IsValidStorageSchema,
+  IsValidUniformSchema,
+} from '../../shared/repr.ts';
 import type {
   Mutable,
   OmitProps,
@@ -387,16 +393,20 @@ export interface RenderPass {
   ): undefined;
 }
 
-type ValidateSchema<TData extends AnyData> = HasNestedType<
-  [TData],
-  'bool'
-> extends true ? 'Error: Bool is not host-shareable, use U32 or I32 instead'
-  : HasNestedType<[TData], 'u16'> extends true ? TData extends {
-      type: 'array';
-      elementType: { type: 'u16' };
-    } ? TData
-    : 'Error: U16 is only usable inside arrays for index buffers'
-  : TData;
+export type ValidateBufferSchema<TData extends AnyData> =
+  IsValidBufferSchema<TData> extends false
+    ? ExtractInvalidSchemaError<TData, '(Error) '>
+    : TData;
+
+export type ValidateStorageSchema<TData extends AnyData> =
+  IsValidStorageSchema<TData> extends false
+    ? ExtractInvalidSchemaError<TData, '(Error) '>
+    : TData;
+
+export type ValidateUniformSchema<TData extends AnyData> =
+  IsValidUniformSchema<TData> extends false
+    ? ExtractInvalidSchemaError<TData, '(Error) '>
+    : TData;
 
 export interface TgpuRoot extends Unwrapper {
   /**
@@ -414,8 +424,9 @@ export interface TgpuRoot extends Unwrapper {
    * @param initial The initial value of the buffer. (optional)
    */
   createBuffer<TData extends AnyData>(
-    typeSchema: ValidateSchema<TData>,
-    initial?: Infer<TData> | undefined,
+    typeSchema: ValidateBufferSchema<TData>,
+    // NoInfer is there to infer the schema type just based on the first parameter
+    initial?: Infer<NoInfer<TData>> | undefined,
   ): TgpuBuffer<TData>;
 
   /**
@@ -428,7 +439,7 @@ export interface TgpuRoot extends Unwrapper {
    * @param gpuBuffer A vanilla WebGPU buffer.
    */
   createBuffer<TData extends AnyData>(
-    typeSchema: ValidateSchema<TData>,
+    typeSchema: ValidateBufferSchema<TData>,
     gpuBuffer: GPUBuffer,
   ): TgpuBuffer<TData>;
 
@@ -441,8 +452,9 @@ export interface TgpuRoot extends Unwrapper {
    * @param initial The initial value of the buffer. (optional)
    */
   createUniform<TData extends AnyWgslData>(
-    typeSchema: TData,
-    initial?: Infer<TData>,
+    typeSchema: ValidateUniformSchema<TData>,
+    // NoInfer is there to infer the schema type just based on the first parameter
+    initial?: Infer<NoInfer<TData>>,
   ): TgpuUniform<TData>;
 
   /**
@@ -454,7 +466,7 @@ export interface TgpuRoot extends Unwrapper {
    * @param gpuBuffer A vanilla WebGPU buffer.
    */
   createUniform<TData extends AnyWgslData>(
-    typeSchema: TData,
+    typeSchema: ValidateUniformSchema<TData>,
     gpuBuffer: GPUBuffer,
   ): TgpuUniform<TData>;
 
@@ -467,8 +479,9 @@ export interface TgpuRoot extends Unwrapper {
    * @param initial The initial value of the buffer. (optional)
    */
   createMutable<TData extends AnyWgslData>(
-    typeSchema: TData,
-    initial?: Infer<TData>,
+    typeSchema: ValidateStorageSchema<TData>,
+    // NoInfer is there to infer the schema type just based on the first parameter
+    initial?: Infer<NoInfer<TData>>,
   ): TgpuMutable<TData>;
 
   /**
@@ -480,7 +493,7 @@ export interface TgpuRoot extends Unwrapper {
    * @param gpuBuffer A vanilla WebGPU buffer.
    */
   createMutable<TData extends AnyWgslData>(
-    typeSchema: TData,
+    typeSchema: ValidateStorageSchema<TData>,
     gpuBuffer: GPUBuffer,
   ): TgpuMutable<TData>;
 
@@ -493,8 +506,9 @@ export interface TgpuRoot extends Unwrapper {
    * @param initial The initial value of the buffer. (optional)
    */
   createReadonly<TData extends AnyWgslData>(
-    typeSchema: TData,
-    initial?: Infer<TData>,
+    typeSchema: ValidateStorageSchema<TData>,
+    // NoInfer is there to infer the schema type just based on the first parameter
+    initial?: Infer<NoInfer<TData>>,
   ): TgpuReadonly<TData>;
 
   /**
@@ -506,7 +520,7 @@ export interface TgpuRoot extends Unwrapper {
    * @param gpuBuffer A vanilla WebGPU buffer.
    */
   createReadonly<TData extends AnyWgslData>(
-    typeSchema: TData,
+    typeSchema: ValidateStorageSchema<TData>,
     gpuBuffer: GPUBuffer,
   ): TgpuReadonly<TData>;
 

packages/typegpu/src/core/vertexLayout/vertexAttribute.ts

@@ -1,13 +1,9 @@
 import type {
-  AnyUnstruct,
   Disarray,
   LooseDecorated,
+  Unstruct,
 } from '../../data/dataTypes.ts';
-import type {
-  AnyWgslStruct,
-  Decorated,
-  WgslArray,
-} from '../../data/wgslTypes.ts';
+import type { Decorated, WgslArray, WgslStruct } from '../../data/wgslTypes.ts';
 import type {
   KindToAcceptedAttribMap,
   KindToDefaultFormatMap,
@@ -23,8 +19,7 @@ import type {
  * - WgslStruct<{ a: Vec3f, b: unorm8x2 }>
  * - WgslStruct<{ nested: WgslStruct<{ a: Vec3f }> }>
  */
-export type DataToContainedAttribs<T> = T extends AnyWgslStruct | AnyUnstruct
-  ? {
+export type DataToContainedAttribs<T> = T extends WgslStruct | Unstruct ? {
     [Key in keyof T['propTypes']]: DataToContainedAttribs<
       T['propTypes'][Key]
     >;

packages/typegpu/src/core/vertexLayout/vertexLayout.ts

@@ -89,8 +89,9 @@ function dataToContainedAttribs<
   if (isWgslStruct(data)) {
     let memberOffset = offset;
 
+    const propTypes = data.propTypes as Record<string, BaseData>;
     return Object.fromEntries(
-      Object.entries(data.propTypes).map(([key, value]) => {
+      Object.entries(propTypes).map(([key, value]) => {
         memberOffset = roundUp(memberOffset, alignmentOf(value));
         const attrib = [
           key,
@@ -111,8 +112,9 @@ function dataToContainedAttribs<
   if (isUnstruct(data)) {
     let memberOffset = offset;
 
+    const propTypes = data.propTypes as Record<string, BaseData>;
     return Object.fromEntries(
-      Object.entries(data.propTypes).map(([key, value]) => {
+      Object.entries(propTypes).map(([key, value]) => {
         memberOffset = roundUp(memberOffset, customAlignmentOf(value));
         const attrib = [
           key,

packages/typegpu/src/data/alignmentOf.ts

@@ -45,7 +45,7 @@ function computeAlignment(data: object): number {
   }
 
   if (isWgslStruct(data)) {
-    return Object.values(data.propTypes)
+    return Object.values(data.propTypes as Record<string, BaseData>)
       .map(alignmentOf)
       .reduce((a, b) => (a > b ? a : b));
   }
@@ -56,7 +56,8 @@ function computeAlignment(data: object): number {
 
   if (isUnstruct(data)) {
     // A loose struct is aligned to its first property.
-    const firstProp = Object.values(data.propTypes)[0];
+    const firstProp =
+      Object.values(data.propTypes as Record<string, BaseData>)[0];
     return firstProp ? (getCustomAlignment(firstProp) ?? 1) : 1;
   }
 
@@ -80,7 +81,8 @@ function computeAlignment(data: object): number {
 function computeCustomAlignment(data: BaseData): number {
   if (isUnstruct(data)) {
     // A loose struct is aligned to its first property.
-    const firstProp = Object.values(data.propTypes)[0];
+    const firstProp =
+      Object.values(data.propTypes as Record<string, BaseData>)[0];
     return firstProp ? customAlignmentOf(firstProp) : 1;
   }
 

packages/typegpu/src/data/atomic.ts

@@ -1,6 +1,13 @@
 import type { Infer, MemIdentity } from '../shared/repr.ts';
 import { $internal } from '../shared/symbols.ts';
-import type { $gpuRepr, $memIdent, $repr } from '../shared/symbols.ts';
+import type {
+  $gpuRepr,
+  $memIdent,
+  $repr,
+  $validStorageSchema,
+  $validUniformSchema,
+  $validVertexSchema,
+} from '../shared/symbols.ts';
 import type { Atomic, atomicI32, atomicU32, I32, U32 } from './wgslTypes.ts';
 
 // ----------
@@ -34,6 +41,9 @@ class AtomicImpl<TSchema extends U32 | I32> implements Atomic<TSchema> {
   declare readonly [$repr]: Infer<TSchema>;
   declare readonly [$memIdent]: MemIdentity<TSchema>;
   declare readonly [$gpuRepr]: TSchema extends U32 ? atomicU32 : atomicI32;
+  declare readonly [$validStorageSchema]: true;
+  declare readonly [$validUniformSchema]: true;
+  declare readonly [$validVertexSchema]: true;
   // ---
 
   constructor(public readonly inner: TSchema) {}