fix: Unify operand types in arithmetic operators (#1603)

reczkok · web-flow · commit ec1a771a7827 · 2025-08-13T09:16:53.000+02:00
* Unify operand types in arithmetic operators

* Remove rogue console logs

* Tweak concretizeTypes logic in convertToCommonType helper

* Make convertToCommonType generic over input tuple type

* chefs kiss
diff --git a/packages/typegpu/src/std/operators.ts b/packages/typegpu/src/std/operators.ts
@@ -1,12 +1,12 @@
 import { stitch } from '../core/resolve/stitch.ts';
 import { createDualImpl } from '../core/function/dualImpl.ts';
-import { f16, f32 } from '../data/numeric.ts';
 import { isSnippetNumeric, snip, type Snippet } from '../data/snippet.ts';
 import { vecTypeToConstructor } from '../data/vector.ts';
 import { VectorOps } from '../data/vectorOps.ts';
 import {
   type AnyMatInstance,
   type AnyNumericVecInstance,
+  type AnyWgslData,
   isFloat32VecInstance,
   isMatInstance,
   isVecInstance,
@@ -17,6 +17,24 @@ import { convertToCommonType } from '../tgsl/generationHelpers.ts';
 import { getResolutionCtx } from '../execMode.ts';
 import type { ResolutionCtx } from '../types.ts';
 import { $internal } from '../shared/symbols.ts';
+import { f16, f32 } from '../data/numeric.ts';
+
+function tryUnify<T extends Snippet[]>(
+  values: T,
+  restrictTo?: AnyWgslData[],
+  concretizeTypes = false,
+  verbose = true,
+): T {
+  const ctx = getResolutionCtx() as ResolutionCtx;
+  const converted = convertToCommonType({
+    ctx,
+    values,
+    restrictTo,
+    concretizeTypes,
+    verbose,
+  });
+  return converted ?? values;
+}
 
 type NumVec = AnyNumericVecInstance;
 type Mat = AnyMatInstance;
@@ -58,7 +76,10 @@ export const add = createDualImpl(
   cpuAdd,
   // CODEGEN implementation
   (lhs, rhs) => {
-    const resultType = isSnippetNumeric(lhs) ? rhs.dataType : lhs.dataType;
+    const [convLhs, convRhs] = tryUnify([lhs, rhs]);
+    const resultType = isSnippetNumeric(convLhs)
+      ? convRhs.dataType
+      : convLhs.dataType;
 
     if (
       (typeof lhs.value === 'number' ||
@@ -72,7 +93,7 @@ export const add = createDualImpl(
       return snip(cpuAdd(lhs.value as never, rhs.value as never), resultType);
     }
 
-    return snip(stitch`(${lhs} + ${rhs})`, resultType);
+    return snip(stitch`(${convLhs} + ${convRhs})`, resultType);
   },
   'add',
 );
@@ -99,7 +120,10 @@ export const sub = createDualImpl(
   cpuSub,
   // CODEGEN implementation
   (lhs, rhs) => {
-    const resultType = isSnippetNumeric(lhs) ? rhs.dataType : lhs.dataType;
+    const [convLhs, convRhs] = tryUnify([lhs, rhs]);
+    const resultType = isSnippetNumeric(convLhs)
+      ? convRhs.dataType
+      : convLhs.dataType;
 
     if (
       (typeof lhs.value === 'number' ||
@@ -113,7 +137,7 @@ export const sub = createDualImpl(
       return snip(cpuSub(lhs.value as never, rhs.value as never), resultType);
     }
 
-    return snip(stitch`(${lhs} - ${rhs})`, resultType);
+    return snip(stitch`(${convLhs} - ${convRhs})`, resultType);
   },
   'sub',
 );
@@ -164,20 +188,21 @@ export const mul = createDualImpl(
   cpuMul,
   // GPU implementation
   (lhs, rhs) => {
-    const returnType = isSnippetNumeric(lhs)
+    const [convLhs, convRhs] = tryUnify([lhs, rhs]);
+    const returnType = isSnippetNumeric(convLhs)
       // Scalar * Scalar/Vector/Matrix
-      ? rhs.dataType
-      : isSnippetNumeric(rhs)
+      ? convRhs.dataType
+      : isSnippetNumeric(convRhs)
       // Vector/Matrix * Scalar
-      ? lhs.dataType
-      : lhs.dataType.type.startsWith('vec')
+      ? convLhs.dataType
+      : convLhs.dataType.type.startsWith('vec')
       // Vector * Vector/Matrix
-      ? lhs.dataType
-      : rhs.dataType.type.startsWith('vec')
+      ? convLhs.dataType
+      : convRhs.dataType.type.startsWith('vec')
       // Matrix * Vector
-      ? rhs.dataType
+      ? convRhs.dataType
       // Matrix * Matrix
-      : lhs.dataType;
+      : convLhs.dataType;
 
     if (
       (typeof lhs.value === 'number' ||
@@ -191,7 +216,7 @@ export const mul = createDualImpl(
       return snip(cpuMul(lhs.value as never, rhs.value as never), returnType);
     }
 
-    return snip(stitch`(${lhs} * ${rhs})`, returnType);
+    return snip(stitch`(${convLhs} * ${convRhs})`, returnType);
   },
   'mul',
 );
@@ -223,35 +248,20 @@ export const div = createDualImpl(
   cpuDiv,
   // CODEGEN implementation
   (lhs, rhs) => {
-    let conv: [Snippet, Snippet] = [lhs, rhs];
-
-    if (isSnippetNumeric(lhs) && isSnippetNumeric(rhs)) {
-      const ctx = getResolutionCtx() as ResolutionCtx;
-      const converted = convertToCommonType({
-        ctx,
-        values: [lhs, rhs],
-        restrictTo: [f32, f16],
-        concretizeTypes: true,
-      }) as
-        | [Snippet, Snippet]
-        | undefined;
-      if (converted) {
-        conv = converted;
-      }
-    }
-
-    const lhsVal = conv[0].value;
-    const rhsVal = conv[1].value;
+    const [convLhs, convRhs] = tryUnify([lhs, rhs], [f32, f16], true, false);
 
     if (
-      (typeof lhsVal === 'number' || isVecInstance(lhsVal)) &&
-      (typeof rhsVal === 'number' || isVecInstance(rhsVal))
+      (typeof lhs.value === 'number' || isVecInstance(lhs.value)) &&
+      (typeof rhs.value === 'number' || isVecInstance(rhs.value))
     ) {
       // Precomputing
-      return snip(cpuDiv(lhsVal as never, rhsVal as never), conv[0].dataType);
+      return snip(
+        cpuDiv(lhs.value as never, rhs.value as never),
+        convLhs.dataType,
+      );
     }
 
-    return snip(stitch`(${conv[0]} / ${conv[1]})`, conv[0].dataType);
+    return snip(stitch`(${convLhs} / ${convRhs})`, convLhs.dataType);
   },
   'div',
 );
@@ -294,8 +304,9 @@ export const mod: ModOverload = createDualImpl(
   },
   // GPU implementation
   (a, b) => {
-    const type = isSnippetNumeric(a) ? b.dataType : a.dataType;
-    return snip(stitch`(${a} % ${b})`, type);
+    const [convA, convB] = tryUnify([a, b]);
+    const type = isSnippetNumeric(convA) ? convB.dataType : convA.dataType;
+    return snip(stitch`(${convA} % ${convB})`, type);
   },
   'mod',
 );
diff --git a/packages/typegpu/src/tgsl/generationHelpers.ts b/packages/typegpu/src/tgsl/generationHelpers.ts
@@ -506,23 +506,31 @@ function applyActionToSnippet(
   }
 }
 
-export type ConvertToCommonTypeOptions = {
+export type ConvertToCommonTypeOptions<T extends Snippet[]> = {
   ctx: ResolutionCtx;
-  values: Snippet[];
+  values: T;
   restrictTo?: AnyData[] | undefined;
   concretizeTypes?: boolean | undefined;
   verbose?: boolean | undefined;
 };
 
-export function convertToCommonType({
+export function convertToCommonType<T extends Snippet[]>({
   ctx,
   values,
   restrictTo,
   concretizeTypes = false,
   verbose = true,
-}: ConvertToCommonTypeOptions): Snippet[] | undefined {
+}: ConvertToCommonTypeOptions<T>): T | undefined {
+  const needsConcretization = concretizeTypes &&
+    // If we have any concrete type among the values, we don't need to concretize
+    !values.some((value) =>
+      concretize(value.dataType as AnyWgslData) === value.dataType
+    );
+
   const types = values.map((value) =>
-    concretizeTypes ? concretize(value.dataType as AnyWgslData) : value.dataType
+    needsConcretization
+      ? concretize(value.dataType as AnyWgslData)
+      : value.dataType
   );
 
   if (types.some((type) => type === UnknownData)) {
@@ -557,7 +565,7 @@ Consider using explicit conversions instead.`,
     const action = conversion.actions[index];
     invariant(action, 'Action should not be undefined');
     return applyActionToSnippet(ctx, value, action, conversion.targetType);
-  });
+  }) as T;
 }
 
 export function tryConvertSnippet(
diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts
@@ -215,11 +215,9 @@ export function generateExpression(
 
     const converted = convertToCommonType({
       ctx,
-      values: [lhsExpr, rhsExpr],
+      values: [lhsExpr, rhsExpr] as const,
       restrictTo: forcedType,
-    }) as
-      | [Snippet, Snippet]
-      | undefined;
+    });
     const [convLhs, convRhs] = converted || [lhsExpr, rhsExpr];
 
     const lhsStr = ctx.resolve(convLhs.value);
diff --git a/packages/typegpu/tests/compiledIO.test.ts b/packages/typegpu/tests/compiledIO.test.ts
@@ -6,7 +6,6 @@ import {
 import * as d from '../src/data/index.ts';
 import { sizeOf } from '../src/data/sizeOf.ts';
 import { it } from './utils/extendedIt.ts';
-import tgpu from '../src/index.ts';
 
 describe('buildWriter', () => {
   it('should compile a writer for a struct', () => {
@@ -293,8 +292,6 @@ describe('createCompileInstructions', () => {
       b: d.align(64, d.arrayOf(d.f32, 2)),
     });
 
-    console.log(tgpu.resolve({ externals: { schema } }));
-
     const builtWriter = buildWriter(schema, 'offset', 'value');
     expect(builtWriter).toMatchInlineSnapshot(`
       "output.setFloat32((offset + 0), value.a, littleEndian);
diff --git a/packages/typegpu/tests/struct.test.ts b/packages/typegpu/tests/struct.test.ts
@@ -388,8 +388,6 @@ describe('abstruct', () => {
       return result.exp;
     });
 
-    console.log(tgpu.resolve({ externals: { testFn } }));
-
     expect(parseResolved({ testFn })).toBe(
       parse(`
         fn testFn(x: f32) -> f32 {
diff --git a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts
@@ -997,7 +997,7 @@ describe('wgslGenerator division operator', () => {
     });
     expect(div()).toBe(0.1);
     expect(parseResolved({ divide1: div })).toMatchInlineSnapshot(
-      `"fn div ( ) -> f32 { return ( f32 ( f16 ( ( f32 ( 1 ) / f32 ( 2 ) ) ) ) / f32 ( 5 ) ) ; }"`,
+      `"fn div ( ) -> f32 { return ( f32 ( f16 ( 0.5 ) ) / f32 ( 5 ) ) ; }"`,
     );
   });
 
