feat: Identifier and Bool nodes in tinyest

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

@@ -16,7 +16,7 @@ import { extractFunctionParts } from './functionParts.ts';
 
 const { NodeTypeCatalog: NODE } = tinyest;
 
-function createContext(params: tinyest.FuncParameter[]): Context {
+function createContext(params: tinyest.FuncParameter[], opts: TranspilationOptions): Context {
   return {
     externalNames: new Map(),
     ignoreExternalDepth: 0,
@@ -30,6 +30,7 @@ function createContext(params: tinyest.FuncParameter[]): Context {
         ),
       },
     ],
+    opts,
   };
 }
 
@@ -51,6 +52,9 @@ function createParser(ast: AstKind) {
       const externalChain = tryFindExternalChain(ctx, node);
       if (externalChain) {
         ctx.externalNames.set(externalChain, externalChain);
+        if (ctx.opts.verboseNodes) {
+          return [NODE.identifier, externalChain];
+        }
         return externalChain;
       }
     }
@@ -60,9 +64,9 @@ function createParser(ast: AstKind) {
   };
 
   return {
-    transpileFn(rootNode: JsNode): TranspilationResult {
+    transpileFn(rootNode: JsNode, options: TranspilationOptions): TranspilationResult {
       const { params, body } = extractFunctionParts(rootNode);
-      const ctx = createContext(params);
+      const ctx = createContext(params, options);
 
       const tinyestBody = transpile(ctx, body);
 
@@ -81,8 +85,8 @@ function createParser(ast: AstKind) {
       };
     },
 
-    transpileNode(node: JsNode): tinyest.AnyNode {
-      return transpile(createContext([]), node);
+    transpileNode(node: JsNode, options: TranspilationOptions): tinyest.AnyNode {
+      return transpile(createContext([], options), node);
     },
   };
 }
@@ -100,8 +104,8 @@ export function transpileFn(
   rootNode: babel.Node,
   options: TranspilationOptions<'babel'>,
 ): TranspilationResult;
-export function transpileFn(rootNode: JsNode, { ast }: TranspilationOptions): TranspilationResult {
-  return parsers[ast].transpileFn(rootNode);
+export function transpileFn(rootNode: JsNode, opts: TranspilationOptions): TranspilationResult {
+  return parsers[opts.ast].transpileFn(rootNode, opts);
 }
 
 export function transpileNode(
@@ -112,6 +116,6 @@ export function transpileNode(
   rootNode: babel.Node,
   options: TranspilationOptions<'babel'>,
 ): tinyest.AnyNode;
-export function transpileNode(rootNode: JsNode, { ast }: TranspilationOptions): tinyest.AnyNode {
-  return parsers[ast].transpileNode(rootNode);
+export function transpileNode(rootNode: JsNode, opts: TranspilationOptions): tinyest.AnyNode {
+  return parsers[opts.ast].transpileNode(rootNode, opts);
 }

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

@@ -45,7 +45,10 @@ export const baseTranspilers = {
       : [NODE.return];
   },
 
-  Identifier(_ctx, node) {
+  Identifier(ctx, node) {
+    if (ctx.opts.verboseNodes) {
+      return [NODE.identifier, node.name];
+    }
     return node.name;
   },
 
@@ -92,13 +95,9 @@ export const baseTranspilers = {
 
     // If the property is not computed, we don't want to register identifiers as external.
     ctx.ignoreExternalDepth++;
-    const property = transpile(ctx, node.property) as tinyest.Expression;
+    const property = transpile(ctx, node.property) as tinyest.Identifier;
     ctx.ignoreExternalDepth--;
 
-    if (typeof property !== 'string') {
-      throw new Error('Expected identifier as property access key.');
-    }
-
     return [NODE.memberAccess, object, property];
   },
 
@@ -147,14 +146,10 @@ export const baseTranspilers = {
 
     const decl = node.declarations[0];
     ctx.ignoreExternalDepth++;
-    const id = transpile(ctx, decl.id);
+    const id = transpile(ctx, decl.id) as tinyest.Identifier;
     ctx.ignoreExternalDepth--;
 
-    if (typeof id !== 'string') {
-      throw new Error('Invalid variable declaration, expected identifier.');
-    }
-
-    ctx.stack[ctx.stack.length - 1]?.declaredNames.push(id);
+    ctx.stack[ctx.stack.length - 1]?.declaredNames.push(extractId(id));
 
     const init = decl.init ? (transpile(ctx, decl.init) as tinyest.Expression) : undefined;
 
@@ -221,14 +216,17 @@ export const baseTranspilers = {
 } satisfies Pick<Transpilers<JsNode>, SharedTranspilers>;
 
 const acornSpecificTranspilers = {
-  Literal(_ctx, node) {
+  Literal(ctx, node) {
     if (node.regex) {
       throw new Error('Regular expression literals are not representable in WGSL.');
     }
     if (node.raw === 'null') {
       return [NODE.nullLiteral];
     }
     if (typeof node.value === 'boolean') {
+      if (ctx.opts.verboseNodes) {
+        return [NODE.booleanLiteral, node.value];
+      }
       return node.value;
     }
     if (typeof node.value === 'string') {
@@ -298,7 +296,10 @@ const babelSpecificTranspilers = {
     return [NODE.numericLiteral, String(Number(node.value))];
   },
 
-  BooleanLiteral(_ctx, node) {
+  BooleanLiteral(ctx, node) {
+    if (ctx.opts.verboseNodes) {
+      return [NODE.booleanLiteral, node.value];
+    }
     return node.value;
   },
 
@@ -362,3 +363,10 @@ export const babelTranspilers = {
   ...(baseTranspilers as Pick<Transpilers<babel.Node>, SharedTranspilers>),
   ...babelSpecificTranspilers,
 } satisfies Transpilers<babel.Node>;
+
+function extractId(ident: tinyest.Identifier): string {
+  if (typeof ident === 'string') {
+    return ident;
+  }
+  return ident[1];
+}

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

@@ -22,6 +22,7 @@ export type Context = {
    */
   visitedNodes: Set<babel.MemberExpression | acorn.MemberExpression>;
   stack: Scope[];
+  opts: TranspilationOptions;
 };
 
 export type TranspilationResult = {
@@ -50,4 +51,11 @@ export type AstKind = 'acorn' | 'babel';
 
 export type TranspilationOptions<TAst extends AstKind = AstKind> = {
   ast: TAst;
+  /**
+   * With this option enabled, identifiers and boolean literals will be wrapped
+   * in dedicated nodes, instead of being transpiled as string/boolean.
+   *
+   * @default false
+   */
+  verboseNodes?: boolean;
 };

packages/tinyest-for-wgsl/tests/helpers.ts

@@ -1,7 +1,7 @@
 import babel from '@babel/parser';
 import type { Node } from '@babel/types';
 import * as acorn from 'acorn';
-import { transpileFn, type TranspilationResult } from 'tinyest-for-wgsl';
+import { transpileFn, type TranspilationOptions, type TranspilationResult } from 'tinyest-for-wgsl';
 
 export const parseRollup = (code: string) => acorn.parse(code, { ecmaVersion: 'latest' });
 export const parseBabel = (code: string) =>
@@ -10,11 +10,15 @@ export const parseBabel = (code: string) =>
 export function dualTest(
   test: <TNode extends Node | acorn.AnyNode>(
     p: (code: string) => TNode,
-    transpileFn: (node: TNode) => TranspilationResult,
+    transpileFn: (node: TNode, options?: Partial<TranspilationOptions>) => TranspilationResult,
   ) => void,
 ) {
   return () => {
-    test<Node>(parseBabel, (node) => transpileFn(node, { ast: 'babel' }));
-    test<acorn.AnyNode>(parseRollup, (node) => transpileFn(node, { ast: 'acorn' }));
+    test<Node>(parseBabel, (node, options) =>
+      transpileFn(node, { ast: 'babel', ...options } as TranspilationOptions<'babel'>),
+    );
+    test<acorn.AnyNode>(parseRollup, (node, options) =>
+      transpileFn(node, { ast: 'acorn', ...options } as TranspilationOptions<'acorn'>),
+    );
   };
 }

packages/tinyest-for-wgsl/tests/verboseNodes.test.ts

@@ -0,0 +1,152 @@
+import type { ClassDeclaration, ClassProperty, Expression, Node } from '@babel/types';
+import * as acorn from 'acorn';
+import { describe, expect, it } from 'vitest';
+import { transpileFn } from '../src/parsers.ts';
+import { dualTest, parseBabel } from './helpers.ts';
+
+describe('verbose nodes', () => {
+  it(
+    'uses nodes for identifiers',
+    dualTest((p, transpileFn) => {
+      const { params, body, externalNames } = transpileFn(
+        p(`(a, b, c) => {
+          return a + b + c;
+        }`),
+        { verboseNodes: true },
+      );
+
+      expect(params).toMatchInlineSnapshot(`
+        [
+          {
+            "name": "a",
+            "type": "i",
+          },
+          {
+            "name": "b",
+            "type": "i",
+          },
+          {
+            "name": "c",
+            "type": "i",
+          },
+        ]
+      `);
+      expect(JSON.stringify(body)).toMatchInlineSnapshot(
+        `"[0,[[10,[1,[1,[9,"a"],"+",[9,"b"]],"+",[9,"c"]]]]]"`,
+      );
+      expect(externalNames).toMatchInlineSnapshot(`Map {}`);
+    }),
+  );
+
+  it(
+    'uses nodes for boolean literals',
+    dualTest((p, transpileFn) => {
+      const { params, body, externalNames } = transpileFn(
+        p(`() => {
+          return true && false;
+        }`),
+        { verboseNodes: true },
+      );
+
+      expect(params).toMatchInlineSnapshot(`[]`);
+      expect(JSON.stringify(body)).toMatchInlineSnapshot(
+        `"[0,[[10,[3,[107,true],"&&",[107,false]]]]]"`,
+      );
+      expect(externalNames).toMatchInlineSnapshot(`Map {}`);
+    }),
+  );
+
+  it(
+    'uses nodes for const declarations',
+    dualTest((p, transpileFn) => {
+      const { params, body, externalNames } = transpileFn(
+        p(`() => {
+          const a = 1;
+        }`),
+        { verboseNodes: true },
+      );
+
+      expect(params).toMatchInlineSnapshot(`[]`);
+      expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[[13,[9,"a"],[5,"1"]]]]"`);
+      expect(externalNames).toMatchInlineSnapshot(`Map {}`);
+    }),
+  );
+
+  it(
+    'uses nodes for let declarations',
+    dualTest((p, transpileFn) => {
+      const { params, body, externalNames } = transpileFn(
+        p(`() => {
+          let a = 1;
+        }`),
+        { verboseNodes: true },
+      );
+
+      expect(params).toMatchInlineSnapshot(`[]`);
+      expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[[12,[9,"a"],[5,"1"]]]]"`);
+      expect(externalNames).toMatchInlineSnapshot(`Map {}`);
+    }),
+  );
+
+  it(
+    'uses nodes for member expressions',
+    dualTest((p, transpileFn) => {
+      const { params, body, externalNames } = transpileFn(
+        p(`() => {
+          const o = {};
+          return o.prop;
+        }`),
+        { verboseNodes: true },
+      );
+
+      expect(params).toMatchInlineSnapshot(`[]`);
+      expect(JSON.stringify(body)).toMatchInlineSnapshot(
+        `"[0,[[13,[9,"o"],[104,{}]],[10,[7,[9,"o"],[9,"prop"]]]]]"`,
+      );
+      expect(externalNames).toMatchInlineSnapshot(`Map {}`);
+    }),
+  );
+
+  it(
+    'uses nodes for externals',
+    dualTest((p, transpileFn) => {
+      const { params, body, externalNames } = transpileFn(
+        p(`() => {
+          return ext + ext.prop;
+        }`),
+        { verboseNodes: true },
+      );
+
+      expect(params).toMatchInlineSnapshot(`[]`);
+      expect(JSON.stringify(body)).toMatchInlineSnapshot(
+        `"[0,[[10,[1,[9,"ext"],"+",[9,"ext.prop"]]]]]"`,
+      );
+      expect(externalNames).toMatchInlineSnapshot(`
+        Map {
+          "ext" => "ext",
+          "ext.prop" => "ext.prop",
+        }
+      `);
+    }),
+  );
+
+  it(
+    'does not use nodes for object expressions',
+    dualTest((p, transpileFn) => {
+      const { params, body, externalNames } = transpileFn(
+        p(`() => {
+          return { p: ext };
+        }`),
+        { verboseNodes: true },
+      );
+
+      expect(params).toMatchInlineSnapshot(`[]`);
+      expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[[10,[104,{"p":[9,"ext"]}]]]]"`);
+      expect(externalNames).toMatchInlineSnapshot(`
+        Map {
+          "ext" => "ext",
+        }
+      `);
+    }),
+  );
+});