impr: Infix dispatch handling

packages/typegpu/src/data/dataTypes.ts

@@ -30,7 +30,6 @@ import type { Snippet } from './snippet.ts';
 import type { PackedData } from './vertexFormatData.ts';
 import * as wgsl from './wgslTypes.ts';
 import type { WgslComparisonSampler, WgslSampler } from './sampler.ts';
-import type { ResolutionCtx } from '../types.ts';
 import type { BaseData } from './wgslTypes.ts';
 
 /**
@@ -237,22 +236,6 @@ export type AnyConcreteData = Exclude<
 export const UnknownData = Symbol('UNKNOWN');
 export type UnknownData = typeof UnknownData;
 
-export class InfixDispatch {
-  readonly name: string;
-  readonly lhs: Snippet;
-  readonly operator: (ctx: ResolutionCtx, args: [lhs: Snippet, rhs: Snippet]) => Snippet;
-
-  constructor(
-    name: string,
-    lhs: Snippet,
-    operator: (ctx: ResolutionCtx, args: [lhs: Snippet, rhs: Snippet]) => Snippet,
-  ) {
-    this.name = name;
-    this.lhs = lhs;
-    this.operator = operator;
-  }
-}
-
 export class MatrixColumnsAccess {
   readonly matrix: Snippet;
 

packages/typegpu/src/data/index.ts

@@ -5,48 +5,39 @@
 // NOTE: This is a barrel file, internal files should not import things from this file
 
 import { Operator } from 'tsover-runtime';
-import { type InfixOperator, infixOperators } from '../tgsl/accessProp.ts';
+import { type InfixOperatorName, infixOperators } from '../tgsl/accessProp.ts';
 import { MatBase } from './matrix.ts';
 import { VecBase } from './vectorImpl.ts';
+import { infixDispatch } from '../tgsl/infixDispatch.ts';
 
 function assignInfixOperator<T extends typeof VecBase | typeof MatBase>(
-  object: T,
-  operator: InfixOperator,
+  base: T,
+  operator: InfixOperatorName,
   operatorSymbol: symbol,
 ) {
-  // oxlint-disable-next-line typescript/no-explicit-any -- anything is possible
-  const proto = object.prototype as any;
-  const opImpl = infixOperators[operator] as (lhs: unknown, rhs: unknown) => unknown;
+  const opImpl = infixOperators[operator];
 
-  proto[operator] = function (this: unknown, other: unknown): unknown {
-    return opImpl(this, other);
-  };
+  Object.defineProperty(base.prototype, operatorSymbol, {
+    value: opImpl,
+  });
 
-  proto[operatorSymbol] = (lhs: unknown, rhs: unknown): unknown => {
-    return opImpl(lhs, rhs);
-  };
+  Object.defineProperty(base.prototype, operator, {
+    get() {
+      return infixDispatch(this, opImpl);
+    },
+  });
 }
 
 assignInfixOperator(VecBase, 'add', Operator.plus);
+assignInfixOperator(MatBase, 'add', Operator.plus);
 assignInfixOperator(VecBase, 'sub', Operator.minus);
+assignInfixOperator(MatBase, 'sub', Operator.minus);
 assignInfixOperator(VecBase, 'mul', Operator.star);
+assignInfixOperator(MatBase, 'mul', Operator.star);
 assignInfixOperator(VecBase, 'div', Operator.slash);
 assignInfixOperator(VecBase, 'mod', Operator.percent);
-assignInfixOperator(MatBase, 'add', Operator.plus);
-assignInfixOperator(MatBase, 'sub', Operator.minus);
-assignInfixOperator(MatBase, 'mul', Operator.star);
-
-// bitShift does not yet have tsover operator symbol
-{
-  // oxlint-disable-next-line typescript/no-explicit-any -- anything is possible
-  const proto = VecBase.prototype as any;
-  proto.bitShiftLeft = function (this: unknown, other: unknown) {
-    return (infixOperators.bitShiftLeft as (a: unknown, b: unknown) => unknown)(this, other);
-  };
-  proto.bitShiftRight = function (this: unknown, other: unknown) {
-    return (infixOperators.bitShiftRight as (a: unknown, b: unknown) => unknown)(this, other);
-  };
-}
+assignInfixOperator(VecBase, 'bitShiftLeft', Symbol()); // bitShift does not yet have tsover operator symbol
+assignInfixOperator(VecBase, 'bitShiftRight', Symbol()); // bitShift does not yet have tsover operator symbol
 
 export { bool, f16, f32, i32, u16, u32 } from './numeric.ts';
 export {

packages/typegpu/src/tgsl/accessProp.ts

@@ -1,13 +1,7 @@
 import { stitch } from '../core/resolve/stitch.ts';
 import { AutoStruct } from '../data/autoStruct.ts';
 import { EntryInputRouter } from '../core/function/entryInputRouter.ts';
-import {
-  InfixDispatch,
-  isUnstruct,
-  MatrixColumnsAccess,
-  undecorate,
-  UnknownData,
-} from '../data/dataTypes.ts';
+import { isUnstruct, MatrixColumnsAccess, undecorate, UnknownData } from '../data/dataTypes.ts';
 import { abstractInt, bool, f16, f32, i32, u32 } from '../data/numeric.ts';
 import { derefSnippet } from '../data/ref.ts';
 import { isEphemeralSnippet, isSnippet, snip, type Snippet } from '../data/snippet.ts';
@@ -37,10 +31,10 @@ import {
   isWgslArray,
   isWgslStruct,
 } from '../data/wgslTypes.ts';
-import { $gpuCallable } from '../shared/symbols.ts';
 import { add, bitShiftLeft, bitShiftRight, div, mod, mul, sub } from '../std/operators.ts';
 import { isKnownAtComptime } from '../types.ts';
 import { coerceToSnippet } from './generationHelpers.ts';
+import { infixDispatch } from './infixDispatch.ts';
 
 const infixKinds = [
   'vec2f',
@@ -70,7 +64,8 @@ export const infixOperators = {
   bitShiftRight,
 } as const;
 
-export type InfixOperator = keyof typeof infixOperators;
+export type InfixOperatorName = keyof typeof infixOperators;
+export type InfixOperator = (typeof infixOperators)[InfixOperatorName];
 
 type SwizzleableType = 'f' | 'h' | 'i' | 'u' | 'b';
 type SwizzleLength = 1 | 2 | 3 | 4;
@@ -110,12 +105,8 @@ const swizzleLenToType: Record<SwizzleableType, Record<SwizzleLength, BaseData>>
 
 export function accessProp(target: Snippet, propName: string): Snippet | undefined {
   if (infixKinds.includes((target.dataType as BaseData).type) && propName in infixOperators) {
-    const operator = infixOperators[propName as InfixOperator];
-    return snip(
-      new InfixDispatch(propName, target, operator[$gpuCallable].call.bind(operator)),
-      UnknownData,
-      /* origin */ target.origin,
-    );
+    const operator = infixOperators[propName as InfixOperatorName];
+    return snip(infixDispatch(target, operator), UnknownData, /* origin */ target.origin);
   }
 
   if (isWgslArray(target.dataType) && propName === 'length') {
@@ -191,15 +182,12 @@ export function accessProp(target: Snippet, propName: string): Snippet | undefin
   }
 
   const propLength = propName.length;
-  if (isVec(target.dataType) && propLength >= 1 && propLength <= 4) {
-    const isXYZW = /^[xyzw]+$/.test(propName);
-    const isRGBA = /^[rgba]+$/.test(propName);
-
-    if (!isXYZW && !isRGBA) {
-      // Not a valid swizzle
-      return undefined;
-    }
-
+  if (
+    isVec(target.dataType) &&
+    propLength >= 1 &&
+    propLength <= 4 &&
+    /^[xyzw]+$|^[rgba]+$/.test(propName)
+  ) {
     const swizzleTypeChar = target.dataType.type.includes('bool')
       ? 'b'
       : (target.dataType.type[4] as SwizzleableType);

packages/typegpu/src/tgsl/infixDispatch.ts

@@ -0,0 +1,52 @@
+import { isSnippet, type Snippet } from '../data/snippet.ts';
+import type { AnyMatInstance, AnyNumericVecInstance } from '../data/wgslTypes.ts';
+import { $internal, isMarkedInternal } from '../shared/symbols.ts';
+import type { InfixOperator } from './accessProp.ts';
+
+type Numeric = number | AnyNumericVecInstance | AnyMatInstance;
+
+/**
+ * In wgslGenerator, the lhs may either be Numeric or Snippet,
+ * and InfixDispatch is recognized by the $internal symbol.
+ * InfixDispatch is not called in wgslGenerator.
+ * @example
+ * const dispatch = d.vec2u(1).mul;
+ * const fn = () => {
+ *    'use gpu';
+ *    dispatch(2); // lhs is Numeric
+ *    d.vec2u(1).mul(2) // lhs is a snippet
+ * }
+ *
+ * In JS, the lhs is always numeric, and InfixDispatch is callable.
+ * @example
+ * const dispatch = d.vec2u(1).mul;
+ * dispatch(2);
+ */
+export interface InfixDispatch {
+  [$internal]: true;
+  type: 'infix-dispatch';
+  readonly lhs: Snippet | Numeric;
+  readonly operator: InfixOperator;
+  (other: Numeric): Numeric;
+}
+
+export function infixDispatch(lhs: Snippet | Numeric, operator: InfixOperator): InfixDispatch {
+  const callable = (other: Numeric | Snippet) => {
+    if (isSnippet(lhs)) {
+      throw new Error('Unexpected snippet lhs in JS infix operator.');
+    }
+    // operator will perform all necessary type checks
+    return operator(lhs as never, other as never);
+  };
+  const infix = Object.assign(callable, {
+    [$internal]: true as const,
+    type: 'infix-dispatch' as const,
+    lhs,
+    operator,
+  });
+  return infix;
+}
+
+export function isInfixDispatch(o: unknown): o is InfixDispatch {
+  return isMarkedInternal(o) && (o as InfixDispatch)?.type === 'infix-dispatch';
+}

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -1,7 +1,7 @@
 import * as tinyest from 'tinyest';
 import { stitch } from '../core/resolve/stitch.ts';
 import { arrayOf } from '../data/array.ts';
-import { type AnyData, InfixDispatch, UnknownData, unptr } from '../data/dataTypes.ts';
+import { type AnyData, UnknownData, unptr } from '../data/dataTypes.ts';
 import { bool, i32, u32 } from '../data/numeric.ts';
 import { vec2u, vec3u, vec4u } from '../data/vector.ts';
 import {
@@ -47,6 +47,7 @@ import { isTgpuRange } from '../std/range.ts';
 import { stringifyNode } from '../shared/tseynit.ts';
 import type { FunctionDefinitionOptions } from './shaderGenerator_members.ts';
 import { getAttributesString } from '../data/attributes.ts';
+import { isInfixDispatch } from './infixDispatch.ts';
 
 const { NodeTypeCatalog: NODE } = tinyest;
 
@@ -646,15 +647,16 @@ ${this.ctx.pre}}`;
         );
       }
 
-      if (callee.value instanceof InfixDispatch) {
-        // Infix operator dispatch.
+      if (isInfixDispatch(callee.value)) {
         if (!argNodes[0]) {
           throw new WgslTypeError(
-            `An infix operator '${callee.value.name}' was called without any arguments`,
+            `An infix operator '${getName(callee.value.operator)}' was called without any arguments`,
           );
         }
+        const lhs = coerceToSnippet(callee.value.lhs);
         const rhs = this._expression(argNodes[0]);
-        return callee.value.operator(this.ctx, [callee.value.lhs, rhs]);
+        const callable = callee.value.operator[$gpuCallable];
+        return callable.call(this.ctx, [lhs, rhs]);
       }
 
       if (isGPUCallable(callee.value)) {

packages/typegpu/tests/swizzleMixedValidation.test.ts

@@ -71,7 +71,7 @@ describe('Mixed swizzle validation', () => {
         return mixed;
       };
 
-      // The resolution should fail because accessProp returns undefined for mixed swizzles
+      // The resolution should fail because accessProp won't match any prop and will return undefined
       expect(() => tgpu.resolve([main])).toThrowErrorMatchingInlineSnapshot(`
         [Error: Resolution of the following tree failed:
         - <root>