fix: Properly infer layout .$ using InferGPU (#2258)

Author: iwoplaza

packages/typegpu/src/data/index.ts

@@ -57,6 +57,8 @@ export type {
   AnyWgslData,
   AnyWgslStruct,
   Atomic,
+  atomicI32,
+  atomicU32,
   BaseData,
   BaseData as BaseWgslData,
   Bool,

packages/typegpu/src/tgpuBindGroupLayout.ts

@@ -49,7 +49,7 @@ import { NotUniformError } from './errors.ts';
 import { isUsableAsStorage, NotStorageError, type StorageFlag } from './extension.ts';
 import type { TgpuNamable } from './shared/meta.ts';
 import { getName, setName } from './shared/meta.ts';
-import type { Infer, MemIdentity } from './shared/repr.ts';
+import type { InferGPU, MemIdentity } from './shared/repr.ts';
 import { safeStringify } from './shared/stringify.ts';
 import { $gpuValueOf, $internal } from './shared/symbols.ts';
 import type { Default, NullableToOptional, Prettify } from './shared/utilityTypes.ts';
@@ -392,19 +392,19 @@ export type BindLayoutEntry<T extends TgpuLayoutEntry | null> = T extends TgpuLa
               : never;
 
 export type InferLayoutEntry<T extends TgpuLayoutEntry | null> = T extends TgpuLayoutUniform
-  ? Infer<T['uniform']>
+  ? InferGPU<T['uniform']>
   : T extends TgpuLayoutStorage
-    ? Infer<UnwrapRuntimeConstructor<T['storage']>>
+    ? InferGPU<UnwrapRuntimeConstructor<T['storage']>>
     : T extends TgpuLayoutSampler
-      ? Infer<WgslSampler>
+      ? InferGPU<WgslSampler>
       : T extends TgpuLayoutComparisonSampler
-        ? Infer<WgslComparisonSampler>
+        ? InferGPU<WgslComparisonSampler>
         : T extends TgpuLayoutTexture<infer TSchema>
-          ? Infer<TSchema>
+          ? InferGPU<TSchema>
           : T extends TgpuLayoutStorageTexture<infer TSchema>
-            ? Infer<TSchema>
+            ? InferGPU<TSchema>
             : T extends TgpuLayoutExternalTexture
-              ? Infer<T['externalTexture']>
+              ? InferGPU<T['externalTexture']>
               : never;
 
 export type ExtractBindGroupInputFromLayout<T extends Record<string, TgpuLayoutEntry | null>> =

packages/typegpu/tests/bindGroupLayout.test.ts

@@ -167,6 +167,18 @@ describe('TgpuBindGroupLayout', () => {
     `);
   });
 
+  it('infers atomics as GPU atomic types in layout.$', () => {
+    const layout = tgpu.bindGroupLayout({
+      counter: { storage: d.atomic(d.u32), access: 'mutable' },
+      signedCounter: { storage: d.atomic(d.i32), access: 'mutable' },
+      atomicArray: { storage: d.arrayOf(d.atomic(d.u32)), access: 'mutable' },
+    });
+
+    expectTypeOf<typeof layout.$.counter>().toEqualTypeOf<d.atomicU32>();
+    expectTypeOf<typeof layout.$.signedCounter>().toEqualTypeOf<d.atomicI32>();
+    expectTypeOf<typeof layout.$.atomicArray>().toEqualTypeOf<d.atomicU32[]>();
+  });
+
   it('takes a pointer to layout.$... if assigned to a const variable', () => {
     const Boid = d.struct({
       pos: d.vec3f,