feat: Accept loose schemas in d.memoryLayoutOf

packages/typegpu/src/data/offsetUtils.ts

@@ -3,16 +3,11 @@ import { sizeOf } from './sizeOf.ts';
 import { strideOf } from './schemaMemoryLayout.ts';
 import { isContiguous } from './isContiguous.ts';
 import { getLongestContiguousPrefix } from './getLongestContiguousPrefix.ts';
-import type {
-  AnyWgslData,
-  BaseData,
-  MatData,
-  VecData,
-  WgslArray,
-  WgslStruct,
-} from './wgslTypes.ts';
+import type { BaseData, MatData, VecData, WgslArray, WgslStruct } from './wgslTypes.ts';
 import { isMat, isVec, isWgslArray, isWgslStruct } from './wgslTypes.ts';
+import type { Disarray, Unstruct } from './dataTypes.ts';
 import { isDisarray, isUnstruct, undecorate } from './dataTypes.ts';
+import { isPackedData, type PackedData } from './vertexFormatData.ts';
 import type { Infer } from '../shared/repr.ts';
 import { vec2f, vec3f, vec4f } from './vector.ts';
 
@@ -64,19 +59,9 @@ function getMarker(target: OffsetProxy, prop: PropertyKey): number | undefined {
   return undefined;
 }
 
-function throwIfLoose(schema: BaseData): void {
-  if (isUnstruct(schema) || isDisarray(schema)) {
-    throw new Error(
-      `memoryLayoutOf: loose schemas are not supported (encountered '${schema.type}').`,
-    );
-  }
-}
-
-function makeProxy(schema: AnyWgslData, baseOffset: number, contiguous: number): unknown {
+function makeProxy(schema: BaseData, baseOffset: number, contiguous: number): unknown {
   const unwrapped = undecorate(schema);
 
-  throwIfLoose(unwrapped);
-
   if (isVec(unwrapped)) {
     return makeVecProxy(unwrapped, scalarNode(baseOffset, contiguous));
   }
@@ -85,19 +70,23 @@ function makeProxy(schema: AnyWgslData, baseOffset: number, contiguous: number):
     return makeMatProxy(unwrapped, scalarNode(baseOffset, contiguous));
   }
 
-  if (isWgslStruct(unwrapped)) {
+  if (isWgslStruct(unwrapped) || isUnstruct(unwrapped)) {
     return makeStructProxy(unwrapped, scalarNode(baseOffset, contiguous));
   }
 
-  if (isWgslArray(unwrapped)) {
+  if (isWgslArray(unwrapped) || isDisarray(unwrapped)) {
     return makeArrayProxy(unwrapped, scalarNode(baseOffset, contiguous));
   }
 
+  if (isPackedData(unwrapped)) {
+    return makePackedProxy(unwrapped, scalarNode(baseOffset, contiguous));
+  }
+
   return scalarNode(baseOffset, contiguous);
 }
 
 export function createOffsetProxy<T extends BaseData>(schema: T, baseOffset = 0): unknown {
-  return makeProxy(schema as AnyWgslData, baseOffset, sizeOf(undecorate(schema)));
+  return makeProxy(schema, baseOffset, sizeOf(undecorate(schema)));
 }
 
 const vecPropToIndex = {
@@ -148,6 +137,25 @@ function makeVecProxy(vec: VecData, parent: OffsetProxy): unknown {
   });
 }
 
+function makePackedProxy(packed: PackedData, parent: OffsetProxy): unknown {
+  return new Proxy(parent, {
+    get(t, prop) {
+      const marker = getMarker(t, prop);
+      if (marker !== undefined) {
+        return marker;
+      }
+
+      if (typeof prop === 'string' && vecPropToIndex[prop] !== undefined) {
+        throw new Error(
+          `memoryLayoutOf: accessing components of packed vertex format '${packed.type}' is not supported.`,
+        );
+      }
+
+      return undefined;
+    },
+  });
+}
+
 function makeMatProxy(mat: MatData, parent: OffsetProxy): unknown {
   const [columnCount, columnSchema] =
     mat.type === 'mat2x2f' ? [2, vec2f] : mat.type === 'mat3x3f' ? [3, vec3f] : [4, vec4f];
@@ -229,9 +237,9 @@ function makeMatProxy(mat: MatData, parent: OffsetProxy): unknown {
   });
 }
 
-function makeArrayProxy(array: WgslArray, parent: OffsetProxy): unknown {
-  const elementType = array.elementType as AnyWgslData;
-  const elementSize = sizeOf(elementType);
+function makeArrayProxy(array: WgslArray | Disarray, parent: OffsetProxy): unknown {
+  const elementType = array.elementType;
+  const elementSize = sizeOf(undecorate(elementType));
   const stride = strideOf(array);
   const hasPadding = stride > elementSize;
   const elementLCP = getLongestContiguousPrefix(elementType);
@@ -287,9 +295,9 @@ type StructFieldMeta = {
   runReachesEnd: boolean;
 };
 
-function makeStructProxy(struct: WgslStruct, parent: OffsetProxy): unknown {
+function makeStructProxy(struct: WgslStruct | Unstruct, parent: OffsetProxy): unknown {
   const offsets = offsetsForProps(struct);
-  const propTypes = struct.propTypes as Record<string, AnyWgslData>;
+  const propTypes = struct.propTypes;
   const props = Object.entries(propTypes);
   const structSize = sizeOf(struct);
 
@@ -306,7 +314,7 @@ function makeStructProxy(struct: WgslStruct, parent: OffsetProxy): unknown {
   let nextPropLocalRemainingData = 0;
   let nextPropRunReachesEnd = false;
   for (let i = props.length - 1; i >= 0; i--) {
-    const [name, type] = props[i] as [string, AnyWgslData];
+    const [name, type] = props[i] as [string, BaseData];
     const info = offsets[name] as PropOffsetInfo;
     const dataSize = sizeOf(undecorate(type));
 
@@ -318,7 +326,7 @@ function makeStructProxy(struct: WgslStruct, parent: OffsetProxy): unknown {
       if (i === props.length - 1) {
         runReachesEnd = true;
       } else {
-        const [, nextType] = props[i + 1] as [string, AnyWgslData];
+        const [, nextType] = props[i + 1] as [string, BaseData];
         if (isContiguous(nextType)) {
           dataAfter = sizeOf(nextType) + nextPropLocalRemainingData;
           runReachesEnd = nextPropRunReachesEnd;
@@ -394,8 +402,6 @@ export function memoryLayoutOf<T extends BaseData>(
   schema: T,
   accessor?: (proxy: Infer<T>) => unknown,
 ): PrimitiveOffsetInfo {
-  throwIfLoose(undecorate(schema));
-
   if (!accessor) {
     return {
       offset: 0,

packages/typegpu/tests/offsetUtils.test.ts

@@ -607,26 +607,238 @@ describe('d.memoryLayoutOf (edge cases)', () => {
   });
 });
 
-describe('d.memoryLayoutOf (loose schemas)', () => {
-  it('throws for unstruct', () => {
-    const Schema = d.unstruct({ a: d.u32, b: d.vec3f });
+describe('d.memoryLayoutOf (unstruct)', () => {
+  describe('packs props without padding', () => {
+    const Schema = d.unstruct({ a: d.f32, b: d.vec3f, c: d.u32 });
 
-    expect(() => d.memoryLayoutOf(Schema)).toThrowErrorMatchingInlineSnapshot(
-      `[Error: memoryLayoutOf: loose schemas are not supported (encountered 'unstruct').]`,
-    );
-    expect(() => d.memoryLayoutOf(Schema, (s) => s.a)).toThrowErrorMatchingInlineSnapshot(
-      `[Error: memoryLayoutOf: loose schemas are not supported (encountered 'unstruct').]`,
-    );
+    it('without offset proxy', () => {
+      expect(d.memoryLayoutOf(Schema).contiguous).toBe(20);
+    });
+
+    it('with offset proxy', () => {
+      const info = d.memoryLayoutOf(Schema, (s) => s.b.y);
+
+      expect(info.offset).toBe(8);
+      expect(info.contiguous).toBe(12);
+    });
+  });
+
+  describe('clips contiguous bytes at custom alignment gap', () => {
+    const Schema = d.unstruct({ a: d.f32, b: d.align(16, d.f32) });
+
+    it('without offset proxy', () => {
+      expect(d.memoryLayoutOf(Schema).contiguous).toBe(4);
+    });
+
+    it('with offset proxy', () => {
+      const info = d.memoryLayoutOf(Schema, (s) => s.a);
+
+      expect(info.offset).toBe(0);
+      expect(info.contiguous).toBe(4);
+    });
+  });
+
+  describe('respects custom prop sizes', () => {
+    const Schema = d.unstruct({ a: d.size(8, d.f32), b: d.f32 });
+
+    it('without offset proxy', () => {
+      expect(d.memoryLayoutOf(Schema).contiguous).toBe(4);
+    });
+
+    it('with offset proxy', () => {
+      const info = d.memoryLayoutOf(Schema, (s) => s.a);
+
+      expect(info.offset).toBe(0);
+      expect(info.contiguous).toBe(4);
+    });
+  });
+
+  describe('does not count a custom size on the root as data', () => {
+    const Schema = d.size(32, d.unstruct({ a: d.f32, b: d.f32 }));
+
+    it('without offset proxy', () => {
+      expect(d.memoryLayoutOf(Schema).contiguous).toBe(8);
+    });
+
+    it('with offset proxy', () => {
+      const info = d.memoryLayoutOf(Schema, (s) => s.b);
+
+      expect(info.offset).toBe(4);
+      expect(info.contiguous).toBe(4);
+    });
+  });
+
+  it('continues from a nested unstruct into the next prop', () => {
+    const Schema = d.unstruct({ a: d.f32, n: d.unstruct({ x: d.f32, y: d.vec2f }), b: d.f32 });
+
+    const info = d.memoryLayoutOf(Schema, (s) => s.n.y.y);
+
+    expect(info.offset).toBe(12);
+    expect(info.contiguous).toBe(8);
   });
 
-  it('throws for disarray', () => {
+  it('respects matrix column padding', () => {
+    const Schema = d.unstruct({ a: d.f32, m: d.mat3x3f });
+
+    const info = d.memoryLayoutOf(Schema, (s) => s.m.columns[1].z);
+
+    expect(info.offset).toBe(28);
+    expect(info.contiguous).toBe(4);
+  });
+
+  describe('continues through a contiguous struct inside', () => {
+    const Inner = d.struct({ v: d.vec3f, w: d.f32 });
+    const Schema = d.unstruct({ a: d.u32, s: Inner, b: d.f32 });
+
+    it('without offset proxy', () => {
+      expect(d.memoryLayoutOf(Schema).contiguous).toBe(24);
+    });
+
+    it('with offset proxy', () => {
+      const info = d.memoryLayoutOf(Schema, (u) => u.s.v.y);
+
+      expect(info.offset).toBe(8);
+      expect(info.contiguous).toBe(16);
+    });
+  });
+
+  describe('limits contiguous bytes to the LCP of a padded nested struct', () => {
+    const Padded = d.struct({ a: d.f32, v: d.vec3f });
+    const Schema = d.unstruct({ x: d.f32, p: Padded, y: d.f32 });
+
+    it('without offset proxy', () => {
+      expect(d.memoryLayoutOf(Schema).contiguous).toBe(8);
+    });
+
+    it('with offset proxy', () => {
+      const info = d.memoryLayoutOf(Schema, (u) => u.p);
+
+      expect(info.offset).toBe(4);
+      expect(info.contiguous).toBe(4);
+    });
+  });
+
+  it('respects trailing padding of a struct inside', () => {
+    const Padded = d.struct({ a: d.f32, v: d.vec3f });
+    const Schema = d.unstruct({ x: d.f32, p: Padded, y: d.f32 });
+
+    const info = d.memoryLayoutOf(Schema, (u) => u.p.v);
+
+    expect(info.offset).toBe(20);
+    expect(info.contiguous).toBe(12);
+  });
+});
+
+describe('d.memoryLayoutOf (disarray)', () => {
+  describe('reports all remaining bytes for contiguous disarray', () => {
     const Schema = d.disarrayOf(d.vec3f, 4);
 
-    expect(() => d.memoryLayoutOf(Schema)).toThrowErrorMatchingInlineSnapshot(
-      `[Error: memoryLayoutOf: loose schemas are not supported (encountered 'disarray').]`,
+    it('without offset proxy', () => {
+      expect(d.memoryLayoutOf(Schema).contiguous).toBe(48);
+    });
+
+    it('with offset proxy', () => {
+      const info = d.memoryLayoutOf(Schema, (a) => a[1]?.z);
+
+      expect(info.offset).toBe(20);
+      expect(info.contiguous).toBe(28);
+    });
+  });
+
+  describe('limits contiguous bytes to element size when custom alignment adds stride padding', () => {
+    const Element = d.unstruct({ inner: d.unstruct({ x: d.align(16, d.f32) }), y: d.f32 });
+    const Schema = d.disarrayOf(Element, 2);
+
+    it('without offset proxy', () => {
+      expect(d.memoryLayoutOf(Schema).contiguous).toBe(8);
+    });
+
+    it('with offset proxy', () => {
+      const info = d.memoryLayoutOf(Schema, (a) => a[1]);
+
+      expect(info.offset).toBe(16);
+      expect(info.contiguous).toBe(8);
+    });
+  });
+
+  describe('uses custom element size as stride', () => {
+    const Schema = d.disarrayOf(d.size(32, d.unstruct({ a: d.f32, b: d.f32 })), 2);
+
+    it('without offset proxy', () => {
+      expect(d.memoryLayoutOf(Schema).contiguous).toBe(8);
+    });
+
+    it('with offset proxy', () => {
+      const info = d.memoryLayoutOf(Schema, (a) => a[1]);
+
+      expect(info.offset).toBe(32);
+      expect(info.contiguous).toBe(8);
+    });
+  });
+
+  it('continues from the last prop of an element into the next element', () => {
+    const LocallyAligned = d.unstruct({ a: d.f32, b: d.align(16, d.f32) });
+    const Schema = d.disarrayOf(LocallyAligned, 2);
+
+    const info = d.memoryLayoutOf(Schema, (a) => a[0]?.b);
+
+    expect(info.offset).toBe(16);
+    expect(info.contiguous).toBe(8);
+  });
+
+  it('respects padding before a disarray nested in an unstruct', () => {
+    const Element = d.unstruct({ x: d.align(16, d.f32) });
+    const Schema = d.unstruct({ a: d.f32, b: d.disarrayOf(Element, 2), c: d.f32 });
+
+    const info = d.memoryLayoutOf(Schema, (s) => s.b[1]?.x);
+
+    expect(info.offset).toBe(32);
+    expect(info.contiguous).toBe(4);
+  });
+});
+
+describe('d.memoryLayoutOf (packed formats)', () => {
+  describe('reports the format size for a root packed format', () => {
+    const Schema = d.unorm8x4;
+
+    it('without offset proxy', () => {
+      expect(d.memoryLayoutOf(Schema).contiguous).toBe(4);
+    });
+
+    it('with offset proxy', () => {
+      const info = d.memoryLayoutOf(Schema, (v) => v);
+
+      expect(info.offset).toBe(0);
+      expect(info.contiguous).toBe(4);
+    });
+  });
+
+  it('computes offsets of packed formats inside an unstruct', () => {
+    const Schema = d.unstruct({ pos: d.float32x3, color: d.unorm8x4, id: d.uint16 });
+
+    const info = d.memoryLayoutOf(Schema, (s) => s.color);
+
+    expect(info.offset).toBe(12);
+    expect(info.contiguous).toBe(6);
+  });
+
+  it('computes offsets of packed formats inside a disarray', () => {
+    const Schema = d.disarrayOf(d.unorm8x4, 3);
+
+    const info = d.memoryLayoutOf(Schema, (a) => a[2]);
+
+    expect(info.offset).toBe(8);
+    expect(info.contiguous).toBe(4);
+  });
+
+  it('throws when accessing components of packed formats', () => {
+    const Schema = d.unstruct({ pos: d.float32x3 });
+
+    expect(() => d.memoryLayoutOf(Schema, (s) => s.pos.x)).toThrowErrorMatchingInlineSnapshot(
+      `[Error: memoryLayoutOf: accessing components of packed vertex format 'float32x3' is not supported.]`,
     );
-    expect(() => d.memoryLayoutOf(Schema, (a) => a[1])).toThrowErrorMatchingInlineSnapshot(
-      `[Error: memoryLayoutOf: loose schemas are not supported (encountered 'disarray').]`,
+    expect(() => d.memoryLayoutOf(Schema, (s) => s.pos[1])).toThrowErrorMatchingInlineSnapshot(
+      `[Error: memoryLayoutOf: accessing components of packed vertex format 'float32x3' is not supported.]`,
     );
   });
 });