impr: Allow numeric literals as bare numbers in the AST

packages/tinyest/src/nodes.ts

@@ -249,8 +249,11 @@ export type PreUpdate = readonly [
   argument: Expression,
 ];
 
-/** A numeric literal */
-export type Num = readonly [type: NodeTypeCatalog['numericLiteral'], string];
+/**
+ * A numeric literal can either be encoded as a plain number, or as a node holding its source text
+ * (e.g. `0x1F`, `1e-3`)
+ */
+export type Num = number | readonly [type: NodeTypeCatalog['numericLiteral'], string];
 
 /** A string literal */
 export type Str = readonly [type: NodeTypeCatalog['stringLiteral'], string];

packages/typegpu/src/shared/tseynit.ts

@@ -114,7 +114,7 @@ function stringifyExpression(node: tinyest.Expression): string {
     return node;
   }
 
-  if (typeof node === 'boolean') {
+  if (typeof node === 'boolean' || typeof node === 'number') {
     return `${node}`;
   }
 
@@ -164,7 +164,10 @@ function stringifyExpression(node: tinyest.Expression): string {
   }
 
   if (node[0] === NODE.memberAccess) {
-    if (Array.isArray(node[1]) && node[1][0] === NODE.numericLiteral) {
+    if (
+      typeof node[1] === 'number' ||
+      (Array.isArray(node[1]) && node[1][0] === NODE.numericLiteral)
+    ) {
       return `(${stringifyExpression(node[1])}).${stringifyExpression(node[2])}`;
     }
     return `${wrapIfComplex(node[1])}.${stringifyExpression(node[2])}`;
@@ -212,6 +215,7 @@ function isExpression(node: tinyest.AnyNode): node is tinyest.Expression {
   if (
     typeof node === 'string' ||
     typeof node === 'boolean' ||
+    typeof node === 'number' ||
     node[0] === NODE.identifier ||
     node[0] === NODE.booleanLiteral ||
     node[0] === NODE.numericLiteral ||
@@ -252,7 +256,7 @@ const SIMPLE_NODES: number[] = [
  */
 function wrapIfComplex(node: tinyest.Expression): string {
   const s = stringifyExpression(node);
-  if (typeof node === 'string' || typeof node === 'boolean') {
+  if (typeof node === 'string' || typeof node === 'boolean' || typeof node === 'number') {
     return s;
   }
   if (SIMPLE_NODES.includes(node[0])) {

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -513,6 +513,13 @@ export class WgslGenerator implements ShaderGenerator {
       return snip(extractBool(expression), bool, /* origin */ 'constant', false);
     }
 
+    if (isNum(expression)) {
+      const value = extractNum(expression);
+      const type = numericLiteralToSnippet(value);
+      invariant(type, `Expected ${stringifyNode(expression)} to be valid numeric literal`);
+      return type;
+    }
+
     if (expression[0] === NODE.logicalExpr) {
       const [_, lhs, op, rhs] = expression;
       const lhsExpr = this._expression(lhs);
@@ -801,16 +808,6 @@ export class WgslGenerator implements ShaderGenerator {
       return accessed;
     }
 
-    if (expression[0] === NODE.numericLiteral) {
-      // Numeric Literal
-      const type =
-        typeof expression[1] === 'string'
-          ? numericLiteralToSnippet(parseNumericString(expression[1]))
-          : numericLiteralToSnippet(expression[1]);
-      invariant(type, `Expected ${stringifyNode(expression)} to be valid numeric literal`);
-      return type;
-    }
-
     if (expression[0] === NODE.call) {
       // Function Call
       const [_, calleeNode, argNodes] = expression;
@@ -1712,7 +1709,7 @@ Try 'return ${typeStr}(${str});' instead.
       }
     }
 
-    if (isId(statement) || isBool(statement)) {
+    if (isId(statement) || isBool(statement) || isNum(statement)) {
       throw new WgslForbiddenStatementError(statement);
     }
 
@@ -2121,7 +2118,6 @@ ${stringifyNode(statement)}`);
     }
 
     if (
-      statement[0] === NODE.numericLiteral ||
       statement[0] === NODE.binaryExpr ||
       statement[0] === NODE.unaryExpr ||
       statement[0] === NODE.logicalExpr ||
@@ -2273,3 +2269,14 @@ function extractBool(ident: tinyest.Bool): boolean {
   }
   return ident[1];
 }
+
+function isNum(expr: unknown): expr is tinyest.Num {
+  return typeof expr === 'number' || (Array.isArray(expr) && expr[0] === NODE.numericLiteral);
+}
+
+function extractNum(num: tinyest.Num): number {
+  if (typeof num === 'number') {
+    return num;
+  }
+  return parseNumericString(num[1]);
+}

packages/typegpu/tests/internal/metadata.test.ts

@@ -223,6 +223,57 @@ describe('meta', () => {
     `);
   });
 
+  it('treats bare numbers and numeric literal nodes the same', () => {
+    const fn = () => {};
+    const meta: RawMetadataV2 = {
+      v: 2,
+      name: 'fn',
+      externals: {},
+      ast: {
+        params: [],
+        body: [
+          NODE.block,
+          [
+            [NODE.let, 'a', [NODE.numericLiteral, '0x10']],
+            [NODE.let, 'b', 16],
+            [NODE.let, 'c', [NODE.numericLiteral, '0.5']],
+            [NODE.let, 'd', 0.5],
+            [NODE.let, 'e', [NODE.binaryExpr, 1, '+', [NODE.numericLiteral, '2.5']]],
+          ],
+        ],
+      },
+    };
+    assignMetadata(fn, meta);
+
+    expect(tgpu.resolve([fn])).toMatchInlineSnapshot(`
+      "fn fn_1() {
+        let a = 16;
+        let b = 16;
+        let c = 0.5;
+        let d = 0.5;
+        let e = 3.5;
+      }"
+    `);
+  });
+
+  it('forbids bare number expression statements', () => {
+    const fn = () => {};
+    const meta: RawMetadataV2 = {
+      v: 2,
+      name: 'fn',
+      externals: {},
+      ast: { params: [], body: [NODE.block, [1]] },
+    };
+    assignMetadata(fn, meta);
+
+    expect(() => tgpu.resolve([fn])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:fn
+      - fn*:fn(): Expression statements like '1;' are forbidden in WGSL. Remove the statement, or use the result in code (for example, assign it to a variable).]
+    `);
+  });
+
   it('correctly handles ident node', () => {
     const fn = () => {};
     const meta: RawMetadataV2 = {

packages/typegpu/tests/internal/tseynit.test.ts

@@ -31,6 +31,13 @@ describe('ast to JS transformation', () => {
       expect(stringifyNode([N.numericLiteral, '-0.0'] satisfies tinyest.Num)).toBe('-0.0');
     });
 
+    it('handles bare numbers', () => {
+      expect(stringifyNode(42 satisfies tinyest.Num)).toBe('42');
+      expect(stringifyNode(0.5 satisfies tinyest.Num)).toBe('0.5');
+      expect(stringifyNode([N.binaryExpr, 1, '+', 2.5])).toBe('1 + 2.5');
+      expect(stringifyNode([N.memberAccess, 1, 'x'])).toBe('(1).x');
+    });
+
     it('handles string literals', () => {
       expect(stringifyNode([N.stringLiteral, 'hello'] satisfies tinyest.Str)).toMatchInlineSnapshot(
         `""hello""`,

packages/unplugin-typegpu/src/core/obfuscate.ts

@@ -101,7 +101,7 @@ const visitors = {
   unaryExpr(ctx: Context, node: tinyest.UnaryExpression) {
     return [NODE.unaryExpr, node[1], obf(ctx, node[2])];
   },
-  numericLiteral(_ctx: Context, node: tinyest.Num) {
+  numericLiteral(_ctx: Context, node: Exclude<tinyest.Num, number>) {
     return [NODE.numericLiteral, node[1]];
   },
   call(ctx: Context, node: tinyest.Call) {
@@ -229,7 +229,7 @@ function obf<T extends tinyest.AnyNode | null>(ctx: Context, node: T): T {
     return ctx.obfuscator.obfuscate(node) as T;
   }
 
-  if (typeof node === 'boolean') {
+  if (typeof node === 'boolean' || typeof node === 'number') {
     return node;
   }