introduce component count (#2161)

Author: reczkok

packages/typegpu/src/data/compiledIO.ts

@@ -204,9 +204,8 @@ export function buildWriter(
     let code = '';
     const writeFunc = primitiveToWriteFunction[primitive];
     const components = ['x', 'y', 'z', 'w'];
-    const count = wgsl.isVec2(node) ? 2 : wgsl.isVec3(node) ? 3 : 4;
 
-    for (let i = 0; i < count; i++) {
+    for (let i = 0; i < node.componentCount; i++) {
       code += `output.${writeFunc}((${offsetExpr} + ${i * 4}), ${valueExpr}.${
         components[i]
       }, littleEndian);\n`;
@@ -251,13 +250,7 @@ export function buildWriter(
     ];
     const writeFunc = primitiveToWriteFunction[primitive];
     const wgslType = formatToWGSLType[formatName];
-    const componentCount = wgsl.isVec4(wgslType)
-      ? 4
-      : wgsl.isVec3(wgslType)
-      ? 3
-      : wgsl.isVec2(wgslType)
-      ? 2
-      : 1;
+    const componentCount = wgsl.isVec(wgslType) ? wgslType.componentCount : 1;
     const componentSize = primitive === 'u8' || primitive === 'i8'
       ? 1
       : primitive === 'u16' || primitive === 'i16' || primitive === 'f16'

packages/typegpu/src/data/vector.ts

@@ -275,14 +275,16 @@ export const vecTypeToConstructor = {
 
 type VecSchemaBase<TValue> = {
   readonly type: string;
+  readonly componentCount: 2 | 3 | 4;
   readonly [$repr]: TValue;
 };
 
 function makeVecSchema<TValue, S extends number | boolean>(
   VecImpl: new (...args: S[]) => VecBase<S>,
   primitive: F32 | F16 | I32 | U32 | Bool,
 ): VecSchemaBase<TValue> & ((...args: (S | AnyVecInstance)[]) => TValue) {
-  const { kind: type, length: componentCount } = new VecImpl();
+  const { kind: type, length } = new VecImpl();
+  const componentCount = length as 2 | 3 | 4;
 
   const cpuConstruct = (...args: (S | AnyVecInstance)[]): TValue => {
     const values = new Array(args.length);
@@ -331,6 +333,7 @@ function makeVecSchema<TValue, S extends number | boolean>(
       [$internal]: {},
       type,
       primitive,
+      componentCount,
       [$repr]: undefined as TValue,
     });
 

packages/typegpu/src/data/wgslTypes.ts

@@ -878,6 +878,7 @@ export interface Vec2f extends
   > {
   readonly type: 'vec2f';
   readonly primitive: F32;
+  readonly componentCount: 2;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v2f;
@@ -900,6 +901,7 @@ export interface Vec2h extends
   > {
   readonly type: 'vec2h';
   readonly primitive: F16;
+  readonly componentCount: 2;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v2h;
@@ -922,6 +924,7 @@ export interface Vec2i extends
   > {
   readonly type: 'vec2i';
   readonly primitive: I32;
+  readonly componentCount: 2;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v2i;
@@ -944,6 +947,7 @@ export interface Vec2u extends
   > {
   readonly type: 'vec2u';
   readonly primitive: U32;
+  readonly componentCount: 2;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v2u;
@@ -967,6 +971,7 @@ export interface Vec2b extends
   > {
   readonly type: 'vec2<bool>';
   readonly primitive: Bool;
+  readonly componentCount: 2;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v2b;
@@ -990,6 +995,7 @@ export interface Vec3f extends
   > {
   readonly type: 'vec3f';
   readonly primitive: F32;
+  readonly componentCount: 3;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v3f;
@@ -1014,6 +1020,7 @@ export interface Vec3h extends
   > {
   readonly type: 'vec3h';
   readonly primitive: F16;
+  readonly componentCount: 3;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v3h;
@@ -1038,6 +1045,7 @@ export interface Vec3i extends
   > {
   readonly type: 'vec3i';
   readonly primitive: I32;
+  readonly componentCount: 3;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v3i;
@@ -1062,6 +1070,7 @@ export interface Vec3u extends
   > {
   readonly type: 'vec3u';
   readonly primitive: U32;
+  readonly componentCount: 3;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v3u;
@@ -1087,6 +1096,7 @@ export interface Vec3b extends
   > {
   readonly type: 'vec3<bool>';
   readonly primitive: Bool;
+  readonly componentCount: 3;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v3b;
@@ -1114,6 +1124,7 @@ export interface Vec4f extends
   > {
   readonly type: 'vec4f';
   readonly primitive: F32;
+  readonly componentCount: 4;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v4f;
@@ -1142,6 +1153,7 @@ export interface Vec4h extends
   > {
   readonly type: 'vec4h';
   readonly primitive: F16;
+  readonly componentCount: 4;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v4h;
@@ -1170,6 +1182,7 @@ export interface Vec4i extends
   > {
   readonly type: 'vec4i';
   readonly primitive: I32;
+  readonly componentCount: 4;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v4i;
@@ -1198,6 +1211,7 @@ export interface Vec4u extends
   > {
   readonly type: 'vec4u';
   readonly primitive: U32;
+  readonly componentCount: 4;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v4u;
@@ -1227,6 +1241,7 @@ export interface Vec4b extends
   > {
   readonly type: 'vec4<bool>';
   readonly primitive: Bool;
+  readonly componentCount: 4;
 
   // Type-tokens, not available at runtime
   readonly [$repr]: v4b;