feat: Inspect vector type in shader function

Author: iwoplaza

packages/tinyest-for-wgsl/src/parsers.ts

@@ -24,73 +24,6 @@ function isDeclared(ctx: Context, name: string) {
   return ctx.stack.some((scope) => scope.declaredNames.includes(name));
 }
 
-const BINARY_OP_MAP = {
-  '==': '==',
-  '!=': '!=',
-  '===': '==',
-  '!==': '!=',
-  '<': '<',
-  '<=': '<=',
-  '>': '>',
-  '>=': '>=',
-  '<<': '<<',
-  '>>': '>>',
-  get '>>>'(): never {
-    throw new Error('The `>>>` operator is unsupported in TGSL.');
-  },
-  '+': '+',
-  '-': '-',
-  '*': '*',
-  '/': '/',
-  '%': '%',
-  '|': '|',
-  '^': '^',
-  '&': '&',
-  get in(): never {
-    throw new Error('The `in` operator is unsupported in TGSL.');
-  },
-  get instanceof(): never {
-    throw new Error('The `instanceof` operator is unsupported in TGSL.');
-  },
-  '**': '**',
-  get '|>'(): never {
-    throw new Error('The `|>` operator is unsupported in TGSL.');
-  },
-} as const;
-
-const LOGICAL_OP_MAP = {
-  '||': '||',
-  '&&': '&&',
-  get '??'(): never {
-    throw new Error('The `??` operator is unsupported in TGSL.');
-  },
-} as const;
-
-const ASSIGNMENT_OP_MAP = {
-  '=': '=',
-  '+=': '+=',
-  '-=': '-=',
-  '*=': '*=',
-  '/=': '/=',
-  '%=': '%=',
-  '<<=': '<<=',
-  '>>=': '>>=',
-  get '>>>='(): never {
-    throw new Error('The `>>>=` operator is unsupported in TGSL.');
-  },
-  '|=': '|=',
-  '^=': '^=',
-  '&=': '&=',
-  get '**='(): never {
-    throw new Error('The `**=` operator is unsupported in TGSL.');
-  },
-  '||=': '||=',
-  '&&=': '&&=',
-  get '??='(): never {
-    throw new Error('The `??=` operator is unsupported in TGSL.');
-  },
-} as const;
-
 const Transpilers: Partial<
   {
     [Type in JsNode['type']]: (
@@ -144,24 +77,36 @@ const Transpilers: Partial<
   },
 
   BinaryExpression(ctx, node) {
-    const wgslOp = BINARY_OP_MAP[node.operator];
     const left = transpile(ctx, node.left) as tinyest.Expression;
     const right = transpile(ctx, node.right) as tinyest.Expression;
-    return [NODE.binaryExpr, left, wgslOp, right];
+    return [
+      NODE.binaryExpr,
+      left,
+      node.operator as tinyest.BinaryOperator,
+      right,
+    ];
   },
 
   LogicalExpression(ctx, node) {
-    const wgslOp = LOGICAL_OP_MAP[node.operator];
     const left = transpile(ctx, node.left) as tinyest.Expression;
     const right = transpile(ctx, node.right) as tinyest.Expression;
-    return [NODE.logicalExpr, left, wgslOp, right];
+    return [
+      NODE.logicalExpr,
+      left,
+      node.operator as tinyest.LogicalOperator,
+      right,
+    ];
   },
 
   AssignmentExpression(ctx, node) {
-    const wgslOp = ASSIGNMENT_OP_MAP[node.operator as acorn.AssignmentOperator];
     const left = transpile(ctx, node.left) as tinyest.Expression;
     const right = transpile(ctx, node.right) as tinyest.Expression;
-    return [NODE.assignmentExpr, left, wgslOp, right];
+    return [
+      NODE.assignmentExpr,
+      left,
+      node.operator as tinyest.AssignmentOperator,
+      right,
+    ];
   },
 
   UnaryExpression(ctx, node) {

packages/tinyest/src/nodes.ts

@@ -118,12 +118,15 @@ export type Statement =
 export type BinaryOperator =
   | '=='
   | '!='
+  | '==='
+  | '!=='
   | '<'
   | '<='
   | '>'
   | '>='
   | '<<'
   | '>>'
+  | '>>>'
   | '+'
   | '-'
   | '*'
@@ -132,6 +135,8 @@ export type BinaryOperator =
   | '|'
   | '^'
   | '&'
+  | 'in'
+  | 'instanceof'
   | '**';
 
 export type BinaryExpression = readonly [
@@ -164,7 +169,7 @@ export type AssignmentExpression = readonly [
   rhs: Expression,
 ];
 
-export type LogicalOperator = '&&' | '||';
+export type LogicalOperator = '&&' | '||' | '??';
 
 export type LogicalExpression = readonly [
   type: NodeTypeCatalog['logicalExpr'],

packages/typegpu/src/tgsl/generationHelpers.ts

@@ -229,6 +229,12 @@ export function accessProp(
     return accessProp(derefed, propName);
   }
 
+  if (isVec(target.dataType)) {
+    if (propName === 'kind') {
+      return snip(target.dataType.type, UnknownData, 'constant');
+    }
+  }
+
   const propLength = propName.length;
   if (
     isVec(target.dataType) &&

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -26,7 +26,7 @@ import { safeStringify } from '../shared/stringify.ts';
 import { $internal } from '../shared/symbols.ts';
 import { pow } from '../std/numeric.ts';
 import { add, div, mul, neg, sub } from '../std/operators.ts';
-import type { FnArgsConversionHint } from '../types.ts';
+import { type FnArgsConversionHint, isKnownAtComptime } from '../types.ts';
 import {
   convertStructValues,
   convertToCommonType,
@@ -50,6 +50,8 @@ const { NodeTypeCatalog: NODE } = tinyest;
 const parenthesizedOps = [
   '==',
   '!=',
+  '===',
+  '!==',
   '<',
   '<=',
   '>',
@@ -68,7 +70,48 @@ const parenthesizedOps = [
   '||',
 ];
 
-const binaryLogicalOps = ['&&', '||', '==', '!=', '<', '<=', '>', '>='];
+const binaryLogicalOps = [
+  '&&',
+  '||',
+  '==',
+  '!=',
+  '===',
+  '!==',
+  '<',
+  '<=',
+  '>',
+  '>=',
+];
+
+const OP_MAP = {
+  //
+  // binary
+  //
+  '===': '==',
+  '!==': '!=',
+  get '>>>'(): never {
+    throw new Error('The `>>>` operator is unsupported in TypeGPU functions.');
+  },
+  get in(): never {
+    throw new Error('The `in` operator is unsupported in TypeGPU functions.');
+  },
+  get instanceof(): never {
+    throw new Error(
+      'The `instanceof` operator is unsupported in TypeGPU functions.',
+    );
+  },
+  get '|>'(): never {
+    throw new Error('The `|>` operator is unsupported in TypeGPU functions.');
+  },
+  //
+  // logical
+  //
+  '||': '||',
+  '&&': '&&',
+  get '??'(): never {
+    throw new Error('The `??` operator is unsupported in TypeGPU functions.');
+  },
+} as Record<string, string>;
 
 type Operator =
   | tinyest.BinaryOperator
@@ -248,8 +291,24 @@ ${this.ctx.pre}}`;
         );
       }
 
+      if (op === '==') {
+        throw new Error('Please use the === operator instead of ==');
+      }
+
+      if (
+        op === '===' && isKnownAtComptime(lhsExpr) && isKnownAtComptime(rhsExpr)
+      ) {
+        return snip(
+          lhsExpr.value === rhsExpr.value,
+          bool,
+          /* ref */ 'constant',
+        );
+      }
+
       if (lhsExpr.dataType.type === 'unknown') {
-        throw new WgslTypeError(`Left-hand side of '${op}' is of unknown type`);
+        throw new WgslTypeError(
+          `Left-hand side of '${op}' is of unknown type`,
+        );
       }
 
       if (rhsExpr.dataType.type === 'unknown') {
@@ -302,8 +361,8 @@ ${this.ctx.pre}}`;
 
       return snip(
         parenthesizedOps.includes(op)
-          ? `(${lhsStr} ${op} ${rhsStr})`
-          : `${lhsStr} ${op} ${rhsStr}`,
+          ? `(${lhsStr} ${OP_MAP[op] ?? op} ${rhsStr})`
+          : `${lhsStr} ${OP_MAP[op] ?? op} ${rhsStr}`,
         type,
         // Result of an operation, so not a reference to anything
         /* ref */ 'runtime',
@@ -657,7 +716,7 @@ ${this.ctx.pre}}`;
     }
 
     if (expression[0] === NODE.stringLiteral) {
-      return snip(expression[1], UnknownData, /* ref */ 'runtime'); // arbitrary ref
+      return snip(expression[1], UnknownData, /* origin */ 'constant');
     }
 
     if (expression[0] === NODE.preUpdate) {

packages/typegpu/src/types.ts

@@ -318,7 +318,8 @@ export function getOwnSnippet(value: unknown): Snippet | undefined {
 }
 
 export function isKnownAtComptime(snippet: Snippet): boolean {
-  return typeof snippet.value !== 'string' &&
+  return (typeof snippet.value !== 'string' ||
+    snippet.dataType.type === 'unknown') &&
     getOwnSnippet(snippet.value) === undefined;
 }
 

packages/typegpu/tests/examples/individual/tgsl-parsing-test.test.ts

@@ -36,10 +36,10 @@ describe('tgsl parsing test example', () => {
 
       fn logicalExpressionTests_1() -> bool {
         var s = true;
-        s = (s && (true == true));
-        s = (s && (true == true));
-        s = (s && (true == true));
-        s = (s && (false == false));
+        s = (s && true);
+        s = (s && true);
+        s = (s && true);
+        s = (s && true);
         s = (s && true);
         s = (s && !false);
         s = (s && true);

packages/typegpu/tests/vector.test.ts

@@ -4,6 +4,7 @@ import { readData, writeData } from '../src/data/dataIO.ts';
 import * as d from '../src/data/index.ts';
 import { sizeOf } from '../src/data/sizeOf.ts';
 import tgpu from '../src/index.ts';
+import * as std from '../src/std/index.ts';
 import { asWgsl } from './utils/parseResolved.ts';
 
 describe('constructors', () => {
@@ -972,3 +973,41 @@ describe('v4b', () => {
     });
   });
 });
+
+describe('type predicates', () => {
+  it('prunes branches', () => {
+    const ceil = (input: d.v3f | d.v3i): d.v3i => {
+      'use gpu';
+      if (input.kind === 'vec3f') {
+        return d.vec3i(std.ceil(input));
+      } else {
+        return input;
+      }
+    };
+
+    const main = () => {
+      'use gpu';
+      const foo = ceil(d.vec3f(1, 2, 3));
+      const bar = ceil(d.vec3i(1, 2, 3));
+    };
+
+    expect(asWgsl(main)).toMatchInlineSnapshot(`
+      "fn ceil(input: vec3f) -> vec3i {
+        {
+          return vec3i(ceil(input));
+        }
+      }
+
+      fn ceil_1(input: vec3i) -> vec3i {
+        {
+          return input;
+        }
+      }
+
+      fn main() {
+        var foo = ceil(vec3f(1, 2, 3));
+        var bar = ceil_1(vec3i(1, 2, 3));
+      }"
+    `);
+  });
+});