diff --git a/packages/typegpu/src/data/attributes.ts b/packages/typegpu/src/data/attributes.ts index 9c6b289435..216d35170b 100644 --- a/packages/typegpu/src/data/attributes.ts +++ b/packages/typegpu/src/data/attributes.ts @@ -101,25 +101,42 @@ export type ExtractAttributes = 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 = number extends TAttribs['length'] + ? TAttribs + : TAttribs extends [infer Head, ...infer Rest] + ? Head extends { readonly type: TType } + ? WithoutAttribute + : [Head, ...WithoutAttribute] + : []; + /** * 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 = IsWgslData extends true - ? Decorated, [TAttrib, ...ExtractAttributes]> + ? Decorated< + Undecorate, + [TAttrib, ...WithoutAttribute, TAttrib['type']>] + > : IsLooseData extends true - ? LooseDecorated, [TAttrib, ...ExtractAttributes]> + ? LooseDecorated< + Undecorate, + [TAttrib, ...WithoutAttribute, TAttrib['type']>] + > : never; export type IsBuiltin = ExtractAttributes[number] extends [] @@ -134,13 +151,23 @@ export type HasCustomLocation = ExtractAttributes[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)) { diff --git a/packages/typegpu/src/data/schemasEqual.ts b/packages/typegpu/src/data/schemasEqual.ts index 3123550527..37d358d879 100644 --- a/packages/typegpu/src/data/schemasEqual.ts +++ b/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 { - const result = new Map(); - for (const attrib of attribs) { - if (!result.has(attrib.type)) { - result.set(attrib.type, attrib); - } - } - return result; -} diff --git a/packages/typegpu/tests/attributes.test.ts b/packages/typegpu/tests/attributes.test.ts index 90ebda9d1b..62b9b7ec36 100644 --- a/packages/typegpu/tests/attributes.test.ts +++ b/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]>; + b: d.Decorated, d.Align<8>]>; + }> + >(); + }); });