Update fluid with atomics example

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

@@ -7,6 +7,7 @@ import type { WgslTypeLiteral } from '../../data/wgslTypes';
 import type { Storage } from '../../extension';
 import type { TgpuNamable } from '../../namable';
 import type { Infer } from '../../shared/repr';
+import type { MemIdentity } from '../../shared/repr';
 import type { UnionToIntersection } from '../../shared/utilityTypes';
 import { isGPUBuffer } from '../../types';
 import type { ExperimentalTgpuRoot } from '../root/rootTypes';
@@ -50,7 +51,7 @@ export interface TgpuBuffer<TData extends AnyData> extends TgpuNamable {
   $addFlags(flags: GPUBufferUsageFlags): this;
 
   write(data: Infer<TData>): void;
-  copyFrom(srcBuffer: TgpuBuffer<TData>): void;
+  copyFrom(srcBuffer: TgpuBuffer<MemIdentity<TData>>): void;
   read(): Promise<Infer<TData>>;
   destroy(): void;
 }
@@ -228,7 +229,7 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
     device.queue.writeBuffer(gpuBuffer, 0, hostBuffer, 0, size);
   }
 
-  copyFrom(srcBuffer: TgpuBuffer<TData>): void {
+  copyFrom(srcBuffer: TgpuBuffer<MemIdentity<TData>>): void {
     if (this.buffer.mapState === 'mapped') {
       throw new Error('Cannot copy to a mapped buffer.');
     }

packages/typegpu/src/data/alignmentOf.ts

@@ -34,6 +34,7 @@ const knownAlignmentMap: Record<string, number> = {
   mat2x2f: 8,
   mat3x3f: 16,
   mat4x4f: 16,
+  atomic: 4,
 };
 
 function computeAlignment(data: object): number {

packages/typegpu/src/data/array.ts

@@ -1,4 +1,4 @@
-import type { Infer } from '../shared/repr';
+import type { Infer, MemIdentity } from '../shared/repr';
 import type { Exotic } from './exotic';
 import { sizeOf } from './sizeOf';
 import type { AnyWgslData, WgslArray } from './wgslTypes';
@@ -48,6 +48,8 @@ class TgpuArrayImpl<TElement extends AnyWgslData>
   public readonly '~repr'!: Infer<TElement>[];
   /** Type-token, not available at runtime */
   public readonly '~exotic'!: WgslArray<Exotic<TElement>>;
+  /** Type-token, not available at runtime */
+  public readonly '~memIdent'!: WgslArray<MemIdentity<TElement>>;
 
   constructor(
     public readonly elementType: TElement,

packages/typegpu/src/data/atomic.ts

@@ -1,4 +1,4 @@
-import type { Infer } from '../shared/repr';
+import type { Infer, MemIdentity } from '../shared/repr';
 import type { Exotic } from './exotic';
 import type { Atomic, I32, U32 } from './wgslTypes';
 
@@ -29,6 +29,8 @@ class AtomicImpl<TSchema extends U32 | I32> implements Atomic<TSchema> {
   public readonly type = 'atomic';
   /** Type-token, not available at runtime */
   public readonly '~repr'!: Infer<TSchema>;
+  /** Type-token, not available at runtime */
+  public readonly '~memIdent'!: MemIdentity<TSchema>;
 
   constructor(public readonly inner: TSchema) {}
 }

packages/typegpu/src/data/attributes.ts

@@ -1,4 +1,4 @@
-import type { Infer } from '../shared/repr';
+import type { Infer, MemIdentity } from '../shared/repr';
 import { alignmentOf } from './alignmentOf';
 import {
   type AnyData,
@@ -337,6 +337,9 @@ class DecoratedImpl<TInner extends BaseWgslData, TAttribs extends unknown[]>
   implements Decorated<TInner, TAttribs>
 {
   public readonly type = 'decorated';
+  public readonly '~memIdent'!: TAttribs extends Location<number>[]
+    ? MemIdentity<TInner> | Decorated<MemIdentity<TInner>, TAttribs>
+    : Decorated<TInner, TAttribs>;
 }
 
 class LooseDecoratedImpl<

packages/typegpu/src/data/sizeOf.ts

@@ -72,6 +72,7 @@ const knownSizesMap: Record<string, number> = {
   sint32x4: 16,
   'unorm10-10-10-2': 4,
   'unorm8x4-bgra': 4,
+  atomic: 4,
 } satisfies Partial<Record<WgslTypeLiteral | LooseTypeLiteral, number>>;
 
 function sizeOfStruct(struct: WgslStruct) {

packages/typegpu/src/data/struct.ts

@@ -1,5 +1,5 @@
 import type { TgpuNamable } from '../namable';
-import type { InferRecord } from '../shared/repr';
+import type { InferRecord, MemIdentityRecord } from '../shared/repr';
 import type { Prettify } from '../shared/utilityTypes';
 import type { ExoticRecord } from './exotic';
 import type { AnyWgslData, BaseWgslData, WgslStruct } from './wgslTypes';
@@ -51,6 +51,8 @@ class TgpuStructImpl<TProps extends Record<string, AnyWgslData>>
   public readonly '~repr'!: InferRecord<TProps>;
   /** Type-token, not available at runtime */
   public readonly '~exotic'!: WgslStruct<ExoticRecord<TProps>>;
+  /** Type-token, not available at runtime */
+  public readonly '~memIdent'!: WgslStruct<MemIdentityRecord<TProps>>;
 
   constructor(public readonly propTypes: TProps) {}
 

packages/typegpu/src/data/wgslTypes.ts

@@ -1,4 +1,14 @@
-import type { Infer, InferRecord } from '../shared/repr';
+import type {
+  Infer,
+  InferRecord,
+  MemIdentity,
+  MemIdentityRecord,
+} from '../shared/repr';
+
+type DecoratedLocation<T extends BaseWgslData> = Decorated<
+  T,
+  Location<number>[]
+>;
 
 export interface NumberArrayView {
   readonly length: number;
@@ -574,12 +584,14 @@ export interface I32 {
   readonly type: 'i32';
   /** Type-token, not available at runtime */
   readonly '~repr': number;
+  readonly '~memIdent': I32 | Atomic<I32> | DecoratedLocation<I32>;
 }
 
 export interface U32 {
   readonly type: 'u32';
   /** Type-token, not available at runtime */
   readonly '~repr': number;
+  readonly '~memIdent': U32 | Atomic<U32> | DecoratedLocation<U32>;
 }
 
 export interface Vec2f {
@@ -680,6 +692,7 @@ export interface WgslStruct<
   readonly propTypes: TProps;
   /** Type-token, not available at runtime */
   readonly '~repr': InferRecord<TProps>;
+  readonly '~memIdent': WgslStruct<MemIdentityRecord<TProps>>;
 }
 
 export interface WgslArray<TElement = BaseWgslData> {
@@ -688,6 +701,7 @@ export interface WgslArray<TElement = BaseWgslData> {
   readonly elementType: TElement;
   /** Type-token, not available at runtime */
   readonly '~repr': Infer<TElement>[];
+  readonly '~memIdent': WgslArray<MemIdentity<TElement>>;
 }
 
 /**
@@ -698,6 +712,7 @@ export interface Atomic<TInner extends U32 | I32 = U32 | I32> {
   readonly inner: TInner;
   /** Type-token, not available at runtime */
   readonly '~repr': Infer<TInner>;
+  readonly '~memIdent': MemIdentity<TInner>;
 }
 
 export interface Align<T extends number> {
@@ -741,6 +756,9 @@ export interface Decorated<
   readonly attribs: TAttribs;
   /** Type-token, not available at runtime */
   readonly '~repr': Infer<TInner>;
+  readonly '~memIdent': TAttribs extends Location<number>[]
+    ? MemIdentity<TInner> | Decorated<MemIdentity<TInner>, TAttribs>
+    : Decorated<TInner, TAttribs>;
 }
 
 export const wgslTypeLiterals = [

packages/typegpu/src/shared/repr.ts

@@ -1,3 +1,5 @@
+import type { AnyData } from '../data/dataTypes';
+
 /**
  * Extracts the inferred representation of a resource.
  * @example
@@ -9,3 +11,15 @@ export type Infer<T> = T extends { readonly '~repr': infer TRepr } ? TRepr : T;
 export type InferRecord<T extends Record<string | number | symbol, unknown>> = {
   [Key in keyof T]: Infer<T[Key]>;
 };
+
+export type MemIdentity<T> = T extends {
+  readonly '~memIdent': infer TMemIdent extends AnyData;
+}
+  ? TMemIdent
+  : T;
+
+export type MemIdentityRecord<
+  T extends Record<string | number | symbol, unknown>,
+> = {
+  [Key in keyof T]: MemIdentity<T[Key]>;
+};

packages/typegpu/tests/buffer.test.ts

@@ -181,4 +181,56 @@ describe('TgpuBuffer', () => {
         .$usage('storage');
     }).toThrow();
   });
+
+  it('should be able to copy from a buffer identical on the byte level', ({
+    root,
+  }) => {
+    const buffer = root.createBuffer(d.u32);
+    const copy = root.createBuffer(d.atomic(d.u32));
+
+    buffer.copyFrom(copy);
+
+    const buffer2 = root.createBuffer(
+      d.struct({
+        one: d.location(0, d.f32), // does nothing outside an IO struct
+        two: d.atomic(d.u32),
+        three: d.arrayOf(d.u32, 3),
+      }),
+    );
+    const copy2 = root.createBuffer(
+      d.struct({
+        one: d.f32,
+        two: d.u32,
+        three: d.arrayOf(d.atomic(d.u32), 3),
+      }),
+    );
+
+    buffer2.copyFrom(copy2);
+
+    const buffer3 = root.createBuffer(
+      d.struct({
+        one: d.size(16, d.f32),
+        two: d.atomic(d.u32),
+      }),
+    );
+
+    const copy3 = root.createBuffer(
+      d.struct({
+        one: d.f32,
+        two: d.u32,
+      }),
+    );
+
+    const copy31 = root.createBuffer(
+      d.struct({
+        one: d.location(0, d.f32),
+        two: d.u32,
+      }),
+    );
+
+    // @ts-expect-error
+    buffer3.copyFrom(copy3);
+    // @ts-expect-error
+    buffer3.copyFrom(copy31);
+  });
 });

packages/typegpu/tests/data/atomic.test.ts

@@ -23,6 +23,7 @@ describe('d.atomic', () => {
       type: 'u32' as const,
       '~repr': undefined as unknown as number,
       '~exotic': undefined as unknown as d.U32,
+      '~memIdent': undefined as unknown as d.U32,
     };
 
     const u32Atomic = d.atomic(exoticU32);
@@ -34,6 +35,7 @@ describe('d.atomic', () => {
       type: 'i32' as const,
       '~repr': undefined as unknown as number,
       '~exotic': undefined as unknown as d.I32,
+      '~memIdent': undefined as unknown as d.I32,
     };
 
     const i32Atomic = d.atomic(exoticI32);
@@ -51,11 +53,13 @@ describe('d.isAtomic', () => {
     const atomicU32 = d.atomic({
       type: 'u32',
       '~repr': undefined as unknown as number,
+      '~memIdent': undefined as unknown as d.U32,
     });
 
     const atomicI32 = d.atomic({
       type: 'i32',
       '~repr': undefined as unknown as number,
+      '~memIdent': undefined as unknown as d.I32,
     });
 
     expect(d.isAtomic(atomicU32)).toEqual(true);