fix: Let a decorator replace an existing one of the same kind

packages/typegpu/src/data/attributes.ts

@@ -101,25 +101,42 @@ export type ExtractAttributes<T> = T extends {
   ? T['attribs']
   : [];
 
+/**
+ * `TAttribs` without the attributes of type `TType`. An array that is not a tuple is kept as it is.
+ */
+type WithoutAttribute<TAttribs extends unknown[], TType> = number extends TAttribs['length']
+  ? TAttribs
+  : TAttribs extends [infer Head, ...infer Rest]
+    ? Head extends { readonly type: TType }
+      ? WithoutAttribute<Rest, TType>
+      : [Head, ...WithoutAttribute<Rest, TType>]
+    : [];
+
 /**
  * Decorates a data-type `TData` with an attribute `TAttrib`.
  *
  * - if `TData` is loose
  *   - if `TData` is already `LooseDecorated`
- *     - Prepend `TAttrib` to the existing attribute tuple.
+ *     - Prepend `TAttrib` to the existing attribute tuple, replacing an attribute of the same type.
  *   - else
  *     - Wrap `TData` with `LooseDecorated` and a single attribute `[TAttrib]`
  * - else
  *   - if `TData` is already `Decorated`
- *     - Prepend `TAttrib` to the existing attribute tuple.
+ *     - Prepend `TAttrib` to the existing attribute tuple, replacing an attribute of the same type.
  *   - else
  *     - Wrap `TData` with `Decorated` and a single attribute `[TAttrib]`
  */
 export type Decorate<TData extends BaseData, TAttrib extends AnyAttribute> =
   IsWgslData<TData> extends true
-    ? Decorated<Undecorate<TData>, [TAttrib, ...ExtractAttributes<TData>]>
+    ? Decorated<
+        Undecorate<TData>,
+        [TAttrib, ...WithoutAttribute<ExtractAttributes<TData>, TAttrib['type']>]
+      >
     : IsLooseData<TData> extends true
-      ? LooseDecorated<Undecorate<TData>, [TAttrib, ...ExtractAttributes<TData>]>
+      ? LooseDecorated<
+          Undecorate<TData>,
+          [TAttrib, ...WithoutAttribute<ExtractAttributes<TData>, TAttrib['type']>]
+        >
       : never;
 
 export type IsBuiltin<T> = ExtractAttributes<T>[number] extends []
@@ -134,13 +151,23 @@ export type HasCustomLocation<T> = ExtractAttributes<T>[number] extends []
     ? true
     : false;
 
+/**
+ * WGSL allows one attribute of each kind, so a new decorator replaces an existing one of the same kind.
+ */
+function withAttribute(attribs: unknown[], attrib: AnyAttribute): AnyAttribute[] {
+  return [
+    attrib,
+    ...(attribs as AnyAttribute[]).filter((existing) => existing.type !== attrib.type),
+  ];
+}
+
 export function attribute(data: BaseData, attrib: AnyAttribute): Decorated | LooseDecorated {
   if (isDecorated(data)) {
-    return new DecoratedImpl(data.inner, [attrib, ...data.attribs]);
+    return new DecoratedImpl(data.inner, withAttribute(data.attribs, attrib));
   }
 
   if (isLooseDecorated(data)) {
-    return new LooseDecoratedImpl(data.inner, [attrib, ...data.attribs]);
+    return new LooseDecoratedImpl(data.inner, withAttribute(data.attribs, attrib));
   }
 
   if (isLooseData(data)) {

packages/typegpu/src/data/schemasEqual.ts

@@ -123,16 +123,16 @@ function schemasEqual(a: AnyData, b: AnyData, structural: boolean): boolean {
       return false;
     }
     // Attribute order carries no meaning, so we compare attributes by type.
-    // Only the first (outermost) attribute of each type is effective.
-    const attribsA = effectiveAttribs(a.attribs as AnyAttribute[]);
-    const attribsB = effectiveAttribs(b.attribs as AnyAttribute[]);
+    // A decorator replaces an existing one of the same type, so each type appears once.
+    const attribsA = a.attribs as AnyAttribute[];
+    const attribsB = b.attribs as AnyAttribute[];
 
-    if (attribsA.size !== attribsB.size) {
+    if (attribsA.length !== attribsB.length) {
       return false;
     }
 
-    for (const [type, attrA] of attribsA) {
-      const attrB = attribsB.get(type);
+    for (const attrA of attribsA) {
+      const attrB = attribsB.find((attr) => attr.type === attrA.type);
       if (!attrB) {
         return false;
       }
@@ -151,13 +151,3 @@ function schemasEqual(a: AnyData, b: AnyData, structural: boolean): boolean {
   // All other types have been checked for equality at the start
   return true;
 }
-
-function effectiveAttribs(attribs: AnyAttribute[]): Map<string, AnyAttribute> {
-  const result = new Map<string, AnyAttribute>();
-  for (const attrib of attribs) {
-    if (!result.has(attrib.type)) {
-      result.set(attrib.type, attrib);
-    }
-  }
-  return result;
-}

packages/typegpu/tests/attributes.test.ts

@@ -19,4 +19,22 @@ describe('attributes', () => {
       }>
     >();
   });
+
+  it('replaces an attribute of the same kind instead of repeating it', () => {
+    const s1 = d.struct({
+      a: d.align(16, d.align(8, d.u32)),
+      b: d.size(16, d.align(8, d.size(8, d.u32))),
+    });
+
+    const wgsl = tgpu.resolve([s1]);
+    expect(wgsl).toContain('@align(16) a: u32,');
+    expect(wgsl).toContain('@size(16) @align(8) b: u32,');
+
+    expectTypeOf(s1).toEqualTypeOf<
+      d.WgslStruct<{
+        a: d.Decorated<d.U32, [d.Align<16>]>;
+        b: d.Decorated<d.U32, [d.Size<16>, d.Align<8>]>;
+      }>
+    >();
+  });
 });