feat: Route binary operations through the shader generator

Author: iwoplaza

packages/typegpu-gl/src/glslGenerator.ts

@@ -1,6 +1,6 @@
 import { NodeTypeCatalog as NODE } from 'tinyest';
 import type { Expression, Return } from 'tinyest';
-import { tgpu, d, type ShaderStage } from 'typegpu';
+import { tgpu, d, type ShaderStage, std } from 'typegpu';
 import { abstractInt, getName, snip, UnknownData, WgslGenerator } from 'typegpu/~internal';
 import type {
   ResolutionCtx,
@@ -10,6 +10,7 @@ import type {
   Origin,
   Snippet,
   ResolvedSnippet,
+  BinaryOperator,
 } from 'typegpu/~internal';
 
 // ----------
@@ -175,6 +176,28 @@ export function getCrossShaderStageState(ctx: ResolutionCtx) {
   return state;
 }
 
+function isF32VecfSchema(
+  schema: d.BaseData | UnknownData,
+): schema is d.F32 | d.Vec2f | d.Vec3f | d.Vec4f {
+  return (
+    typeof schema !== 'symbol' &&
+    (schema.type === 'f32' ||
+      schema.type === 'vec2f' ||
+      schema.type === 'vec3f' ||
+      schema.type === 'vec4f')
+  );
+}
+
+const HELPERS = {
+  // TODO(#2821): Make signature more accurate when std.sign and std.abs
+  // accept a wider union
+  remainder: (x: number, y: number): number => {
+    'use gpu';
+    const truncDiv = std.sign(x / y) * std.floor(std.abs(x / y));
+    return x - y * truncDiv;
+  },
+};
+
 /**
  * A GLSL ES 3.0 shader generator that extends WgslGenerator.
  * Overrides `dataType` to emit GLSL type names instead of WGSL ones,
@@ -247,7 +270,7 @@ export class GlslGenerator extends WgslGenerator {
     return snip(options.id, options.dataType, options.scope);
   }
 
-  override typeAnnotation(data: d.BaseData): string {
+  override emitTypeAnnotation(data: d.BaseData): string {
     if (!d.isLooseData(data)) {
       const glslName = WGSL_TO_GLSL_TYPE[data.type];
       if (glslName !== undefined) {
@@ -257,17 +280,17 @@ export class GlslGenerator extends WgslGenerator {
 
     if (d.isWgslArray(data)) {
       // The array size suffix is handled elsewhere
-      return this.typeAnnotation(data.elementType);
+      return this.emitTypeAnnotation(data.elementType);
     }
 
     if (d.isWgslStruct(data)) {
       return resolveStruct(this.ctx, data);
     }
 
-    return super.typeAnnotation(data);
+    return super.emitTypeAnnotation(data);
   }
 
-  override call(
+  override emitCall(
     name: string,
     templateParams: readonly Snippet[],
     args: readonly Snippet[],
@@ -291,16 +314,16 @@ export class GlslGenerator extends WgslGenerator {
         : sourceSchema;
 
       if (sourcePrimitive.type === 'u32' && targetPrimitive.type === 'f32') {
-        return super.call('uintBitsToFloat', [], [source]);
+        return super.emitCall('uintBitsToFloat', [], [source]);
       }
       if (sourcePrimitive.type === 'i32' && targetPrimitive.type === 'f32') {
-        return super.call('intBitsToFloat', [], [source]);
+        return super.emitCall('intBitsToFloat', [], [source]);
       }
       if (sourcePrimitive.type === 'f32' && targetPrimitive.type === 'u32') {
-        return super.call('floatBitsToUint', [], [source]);
+        return super.emitCall('floatBitsToUint', [], [source]);
       }
       if (sourcePrimitive.type === 'f32' && targetPrimitive.type === 'i32') {
-        return super.call('floatBitsToInt', [], [source]);
+        return super.emitCall('floatBitsToInt', [], [source]);
       }
       if (sourceSchema.type === targetSchema.type) {
         return this.ctx.resolveSnippet(source).value;
@@ -317,9 +340,9 @@ export class GlslGenerator extends WgslGenerator {
 
       if (falsy.dataType !== UnknownData && falsy.dataType.type.startsWith('vec')) {
         if (cond.dataType !== UnknownData && cond.dataType.type.startsWith('vec')) {
-          return super.call('mix', templateParams, args);
+          return super.emitCall('mix', templateParams, args);
         }
-        return super.call('mix', templateParams, [
+        return super.emitCall('mix', templateParams, [
           falsy,
           truthy,
           this.typeInstantiation(correspondingBooleanVectorSchema(falsy.dataType), [cond]),
@@ -334,10 +357,14 @@ export class GlslGenerator extends WgslGenerator {
       if (!arg) {
         throw new Error(`Invalid number of arguments for 'saturate'`);
       }
-      return super.call('clamp', [], [arg, snip(0, d.f32, 'constant'), snip(1, d.f32, 'constant')]);
+      return super.emitCall(
+        'clamp',
+        [],
+        [arg, snip(0, d.f32, 'constant'), snip(1, d.f32, 'constant')],
+      );
     }
 
-    return super.call(name, templateParams, args);
+    return super.emitCall(name, templateParams, args);
   }
 
   override typeInstantiation(schema: d.BaseData, args: Snippet[]): ResolvedSnippet {
@@ -393,6 +420,22 @@ export class GlslGenerator extends WgslGenerator {
     return `${this.ctx.pre}${glslTypeName} ${name}${resolveArraySizeSuffix(this.ctx, dataType)} = ${rhsStr};`;
   }
 
+  override emitBinaryOp(lhs: Snippet, op: BinaryOperator, rhs: Snippet): string {
+    if (op === '%' && (isF32VecfSchema(lhs.dataType) || isF32VecfSchema(rhs.dataType))) {
+      const result = this._callShellless(HELPERS.remainder, [lhs, rhs]);
+      if (!result) {
+        const lhsStr = this.ctx.resolveSnippet(lhs).value;
+        const rhsStr = this.ctx.resolveSnippet(rhs).value;
+        throw new Error(
+          `[@typegpu/gl] Invalid use of '%', incompatible with the GLSL generator: ${lhsStr} (type: ${String(lhs.dataType)}) ${op} ${rhsStr} (type: ${String(rhs.dataType)})`,
+        );
+      }
+      return result.value;
+    }
+
+    return super.emitBinaryOp(lhs, op, rhs);
+  }
+
   /**
    * GLSL has no pointers, so `const x = <alias>;` cannot be turned into an implicit
    * pointer definition like it is in WGSL. Instead:

packages/typegpu-gl/tests/glslGenerator.test.ts

@@ -202,6 +202,30 @@ describe('GlslGenerator - standard function calls', () => {
   });
 });
 
+describe('GlslGenerator - operator', () => {
+  it('translates % with floating-point arguments to a call to the `remainder` helper function', () => {
+    function foo() {
+      'use gpu';
+      const value = 2;
+      const rem = value % 5;
+      return (1 + rem) % 0.5;
+    }
+
+    expect(tgpu.resolve([foo], glOptions())).toMatchInlineSnapshot(`
+      "float remainder(float x, float y) {
+        float truncDiv = (sign((x / y)) * floor(abs((x / y))));
+        return (x - (y * truncDiv));
+      }
+
+      float foo() {
+        int value = 2;
+        int rem = (value % 5);
+        return remainder(float((1 + rem)), 0.5);
+      }"
+    `);
+  });
+});
+
 describe('GlslGenerator - function definitions', () => {
   it('generates proper function signatures', () => {
     function add(a: number, b: number) {

packages/typegpu/src/internal.ts

@@ -15,6 +15,7 @@ export type { Snippet, ResolvedSnippet, Origin } from './data/snippet.ts';
 
 export type {
   ShaderGenerator,
+  BinaryOperator,
   ShaderGeneratorClass,
   FunctionDefinitionOptions,
   ConstantDefinitionOptions,

packages/typegpu/src/resolutionCtx.ts

@@ -947,7 +947,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
       let result: ResolvedSnippet;
       if (isData(item)) {
         // Ref is arbitrary, as we're resolving a schema
-        result = snip(this.gen.typeAnnotation(item), Void, /* origin */ 'runtime');
+        result = snip(this.gen.emitTypeAnnotation(item), Void, /* origin */ 'runtime');
       } else if (isLazy(item) || isSlot(item)) {
         result = this.resolve(this.unwrap(item));
       } else if (isSelfResolvable(item)) {

packages/typegpu/src/std/array.ts

@@ -19,7 +19,7 @@ export const arrayLength = dualImpl({
   normalImpl: (a: unknown[] | ref<unknown[]>) => (isRef(a) ? a.$.length : a.length),
   codegenImpl(ctx, [a]) {
     const length = sizeOfPointedToArray(a.dataType);
-    return length > 0 ? `${length}` : ctx.gen.call('arrayLength', [], [a]);
+    return length > 0 ? `${length}` : ctx.gen.emitCall('arrayLength', [], [a]);
   },
   sideEffects: false,
 });

packages/typegpu/src/std/bitcast.ts

@@ -65,7 +65,7 @@ export const bitcastU32toF32 = dualImpl({
     return VectorOps.bitcastU32toF32[value.kind](value);
   }) as BitcastU32toF32Overload,
   codegenImpl: (ctx, [n], returnType) => {
-    return ctx.gen.call('bitcast', [coerceToSnippet(returnType)], [n]);
+    return ctx.gen.emitCall('bitcast', [coerceToSnippet(returnType)], [n]);
   },
   signature: (...arg) => {
     const uargs = unifyStrict(arg, u32AllowedSchemas);
@@ -102,7 +102,7 @@ export const bitcastU32toI32 = dualImpl({
     return VectorOps.bitcastU32toI32[value.kind](value);
   }) as BitcastU32toI32Overload,
   codegenImpl: (ctx, [n], returnType) => {
-    return ctx.gen.call('bitcast', [coerceToSnippet(returnType)], [n]);
+    return ctx.gen.emitCall('bitcast', [coerceToSnippet(returnType)], [n]);
   },
   signature: (...arg) => {
     const uargs = unifyStrict(arg, u32AllowedSchemas);
@@ -141,7 +141,7 @@ export const bitcastF32toU32 = dualImpl({
     return VectorOps.bitcastF32toU32[value.kind](value);
   }) as BitcastF32toU32Overload,
   codegenImpl: (ctx, [n], returnType) => {
-    return ctx.gen.call('bitcast', [coerceToSnippet(returnType)], [n]);
+    return ctx.gen.emitCall('bitcast', [coerceToSnippet(returnType)], [n]);
   },
   signature: (...arg) => {
     const uargs = unifyStrict(arg, f32AllowedSchemas);
@@ -277,7 +277,7 @@ function bitcastFor<In extends BitcastAllowedTypes, Out extends BitcastAllowedTy
   return dualImpl({
     name: 'bitcast',
     normalImpl: getCpuBitcast<In, Out>(inType, outType),
-    codegenImpl: (ctx, [n]) => ctx.gen.call('bitcast', [coerceToSnippet(outType)], [n]),
+    codegenImpl: (ctx, [n]) => ctx.gen.emitCall('bitcast', [coerceToSnippet(outType)], [n]),
     signature: (arg) => {
       const uarg = unifyStrict([arg], [inType]);
       if (!uarg) {

packages/typegpu/src/std/boolean.ts

@@ -443,7 +443,7 @@ export const select = dualImpl({
   },
   normalImpl: cpuSelect,
   codegenImpl: (ctx, [f, t, cond]) => {
-    const result = ctx.gen.call('select', [], [f, t, cond]);
+    const result = ctx.gen.emitCall('select', [], [f, t, cond]);
     if (
       !validSelectBranchTypes.includes(f.dataType as AnyWgslData) ||
       !validSelectBranchTypes.includes(t.dataType as AnyWgslData)

packages/typegpu/src/std/numeric.ts

@@ -1072,7 +1072,7 @@ export const saturate = dualImpl({
   name: 'saturate',
   signature: unifyRestrictedSignature(anyFloat),
   normalImpl: cpuSaturate,
-  codegenImpl: (ctx, [value]) => ctx.gen.call('saturate', [], [value]),
+  codegenImpl: (ctx, [value]) => ctx.gen.emitCall('saturate', [], [value]),
   sideEffects: false,
 });
 

packages/typegpu/src/std/operators.ts

@@ -116,7 +116,7 @@ export const add = dualImpl({
   name: 'add',
   signature: binaryArithmeticSignature,
   normalImpl: cpuAdd,
-  codegenImpl: (_ctx, [lhs, rhs]) => stitch`(${lhs} + ${rhs})`,
+  codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '+', rhs),
   sideEffects: false,
 });
 
@@ -144,7 +144,7 @@ export const sub = dualImpl({
   name: 'sub',
   signature: binaryArithmeticSignature,
   normalImpl: cpuSub,
-  codegenImpl: (_ctx, [lhs, rhs]) => stitch`(${lhs} - ${rhs})`,
+  codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '-', rhs),
   sideEffects: false,
 });
 
@@ -196,7 +196,7 @@ export const mul = dualImpl({
   name: 'mul',
   signature: binaryMulSignature,
   normalImpl: cpuMul,
-  codegenImpl: (_ctx, [lhs, rhs]) => stitch`(${lhs} * ${rhs})`,
+  codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '*', rhs),
   sideEffects: false,
 });
 
@@ -226,7 +226,7 @@ export const div = dualImpl({
   name: 'div',
   signature: binaryDivSignature,
   normalImpl: cpuDiv,
-  codegenImpl: (_ctx, [lhs, rhs]) => stitch`(${lhs} / ${rhs})`,
+  codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '/', rhs),
   ignoreImplicitCastWarning: true,
   sideEffects: false,
 });
@@ -266,7 +266,7 @@ export const mod = dualImpl({
     }
     throw new Error('Mod called with invalid arguments, expected types: number or vector.');
   }) as ModOverload,
-  codegenImpl: (_ctx, [lhs, rhs]) => stitch`(${lhs} % ${rhs})`,
+  codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '%', rhs),
   sideEffects: false,
 });
 
@@ -352,13 +352,13 @@ export const bitShiftLeft = dualImpl({
   name: 'bitShiftLeft',
   signature: bitShiftSignature,
   normalImpl: cpuBitShiftLeft,
-  codegenImpl: (_ctx, [lhs, rhs]) => {
+  codegenImpl: (ctx, [lhs, rhs]) => {
     if (isVec(lhs.dataType) && !isVec(rhs.dataType)) {
       const cc = lhs.dataType.componentCount;
-      const schema = cc === 2 ? 'vec2u' : cc === 3 ? 'vec3u' : 'vec4u';
-      return stitch`(${lhs} << ${schema}(${rhs}))`;
+      const schema = cc === 2 ? vec2u : cc === 3 ? vec3u : vec4u;
+      return ctx.gen.emitBinaryOp(lhs, '<<', ctx.gen.typeInstantiation(schema, [rhs]));
     }
-    return stitch`(${lhs} << ${rhs})`;
+    return ctx.gen.emitBinaryOp(lhs, '<<', rhs);
   },
   sideEffects: false,
 });
@@ -389,13 +389,13 @@ export const bitShiftRight = dualImpl({
   name: 'bitShiftRight',
   signature: bitShiftSignature,
   normalImpl: cpuBitShiftRight,
-  codegenImpl: (_ctx, [lhs, rhs]) => {
+  codegenImpl: (ctx, [lhs, rhs]) => {
     if (isVec(lhs.dataType) && !isVec(rhs.dataType)) {
       const cc = lhs.dataType.componentCount;
-      const schema = cc === 2 ? 'vec2u' : cc === 3 ? 'vec3u' : 'vec4u';
-      return stitch`(${lhs} >> ${schema}(${rhs}))`;
+      const schema = cc === 2 ? vec2u : cc === 3 ? vec3u : vec4u;
+      return ctx.gen.emitBinaryOp(lhs, '>>', ctx.gen.typeInstantiation(schema, [rhs]));
     }
-    return stitch`(${lhs} >> ${rhs})`;
+    return ctx.gen.emitBinaryOp(lhs, '>>', rhs);
   },
   sideEffects: false,
 });

packages/typegpu/src/tgsl/shaderGenerator.ts

@@ -43,6 +43,25 @@ export interface ShaderGeneratorClass<T extends ShaderGenerator = ShaderGenerato
   new (): T;
 }
 
+export type BinaryOperator =
+  | '='
+  | '^'
+  | '|'
+  | '&'
+  | '*'
+  | '/'
+  | '%'
+  | '+'
+  | '-'
+  | '<<'
+  | '>>'
+  | '<'
+  | '>'
+  | '<='
+  | '>='
+  | '=='
+  | '!=';
+
 /**
  * Represents generators that, once instantiated, will generate `wgsl` (as opposed to e.g. `glsl`)
  */
@@ -65,6 +84,8 @@ export interface ShaderGenerator {
 
   typeInstantiation(schema: BaseData, args: readonly Snippet[]): ResolvedSnippet;
   numericLiteral(value: number, schema: BaseData): ResolvedSnippet;
-  typeAnnotation(schema: BaseData): string;
-  call(name: string, templateParams: readonly Snippet[], args: readonly Snippet[]): string;
+
+  emitTypeAnnotation(schema: BaseData): string;
+  emitCall(name: string, templateParams: readonly Snippet[], args: readonly Snippet[]): string;
+  emitBinaryOp(lhs: Snippet, op: BinaryOperator, rhs: Snippet): string;
 }