feat: Expand property access in externals

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

@@ -10,12 +10,17 @@
   declaredNames: string[];
 };
 
+interface Externals {
+  [key: string]: Externals | string;
+}
+
 type Context = {
   /** Holds a set of all identifiers that were used in code, but were not declared in code. */
-  externalNames: Set<string>;
+  externalNames: Externals;
   /** Used to signal to identifiers that they should not treat their resolution as possible external uses. */
   ignoreExternalDepth: number;
   stack: Scope[];
+  ancestorChain: JsNode[];
 };
 
 type JsNode = babel.Node | acorn.AnyNode;
@@ -28,6 +33,64 @@
   return transpile(ctx, node.expression);
 };
 
+function addExternal(
+  ctx: Context,
+  node: babel.ThisExpression | acorn.ThisExpression | babel.Identifier | acorn.Identifier,
+) {
+  // TODO: clean up this mess
+  const chain: string[] = [];
+  for (let i = ctx.ancestorChain.length - 1; i >= 0; i--) {
+    const current = ctx.ancestorChain[i];
+
+    if (!current) {
+      break;
+    }
+
+    if (current.type === 'Identifier') {
+      chain.push(current.name);
+    } else if (current.type === 'ThisExpression') {
+      chain.push('this');
+    } else if (current.type === 'MemberExpression' && !current.computed) {
+      chain.push(`${(current.property as { name: string }).name}`); // TODO: better handling of other nodes
+    } else {
+      break;
+    }
+  }
+
+  let currentExternals = ctx.externalNames;
+  if (typeof currentExternals !== 'object') {
+    throw new Error('??');
+  }
+  for (const elem of chain) {
+    let nextExternals = currentExternals[elem];
+    if (nextExternals) {
+      if (typeof nextExternals !== 'string') {
+        currentExternals = nextExternals;
+      } else {
+        // we already need this in externals, so we break
+        break;
+      }
+    } else {
+      const newExt = Object.create(null);
+      currentExternals[elem] = newExt;
+      currentExternals = newExt;
+    }
+  }
+
+  const lastKey = chain[chain.length - 1];
+  if (lastKey !== undefined) {
+    let parent = ctx.externalNames;
+    for (const key of chain.slice(0, -1)) {
+      const next = parent[key];
+      if (!next || typeof next === 'string') break;
+      parent = next;
+    }
+    if (typeof parent[lastKey] === 'object') {
+      parent[lastKey] = chain.join('.');
+    }
+  }
+}
+
 const Transpilers: Partial<{
   [Type in JsNode['type']]: (
     ctx: Context,
@@ -70,14 +133,13 @@
 
   Identifier(ctx, node) {
     if (ctx.ignoreExternalDepth === 0 && !isDeclared(ctx, node.name)) {
-      ctx.externalNames.add(node.name);
+      addExternal(ctx, node);
     }
-
     return node.name;
   },
 
-  ThisExpression(ctx) {
-    ctx.externalNames.add('this');
+  ThisExpression(ctx, node) {
+    addExternal(ctx, node);
     return 'this';
   },
 
@@ -308,8 +370,11 @@
     throw new Error(`Unsupported JS functionality: ${node.type}`);
   }
 
+  ctx.ancestorChain.push(node);
   // @ts-expect-error <too much for typescript, it seems :/ >
-  return transpiler(ctx, node);
+  const result = transpiler(ctx, node);
+  ctx.ancestorChain.pop();
+  return result;
 }
 
 export type TranspilationResult = {
@@ -319,7 +384,7 @@
    * All identifiers found in the function code that are not declared in the
    * function itself, or in the block that is accessing that identifier.
    */
-  externalNames: string[];
+  externalNames: Externals;
 };
 
 export function extractFunctionParts(rootNode: JsNode): {
@@ -424,7 +489,7 @@
   const { params, body } = extractFunctionParts(rootNode);
 
   const ctx: Context = {
-    externalNames: new Set(),
+    externalNames: Object.create(null),
     ignoreExternalDepth: 0,
     stack: [
       {
@@ -435,35 +500,36 @@
         ),
       },
     ],
+    ancestorChain: [],
   };
 
   const tinyestBody = transpile(ctx, body);
-  const externalNames = [...ctx.externalNames];
 
   if (body.type === 'BlockStatement') {
     return {
       params,
       body: tinyestBody as tinyest.Block,
-      externalNames,
+      externalNames: ctx.externalNames,
     };
   }
 
   return {
     params,
     body: [NODE.block, [[NODE.return, tinyestBody as tinyest.Expression]]],
-    externalNames,
+    externalNames: ctx.externalNames,
   };
 }
 
 export function transpileNode(node: JsNode): tinyest.AnyNode {
   const ctx: Context = {
-    externalNames: new Set(),
+    externalNames: Object.create(null),
     ignoreExternalDepth: 0,
     stack: [
       {
         declaredNames: [],
       },
     ],
+    ancestorChain: [],
   };
 
   return transpile(ctx, node);

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

@@ -30,7 +30,7 @@ describe('transpileFn', () => {
 
       expect(params).toStrictEqual([]);
       expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[]]"`);
-      expect(externalNames).toStrictEqual([]);
+      expect(externalNames).toMatchInlineSnapshot(`{}`);
     }),
   );
 
@@ -41,7 +41,7 @@ describe('transpileFn', () => {
 
       expect(params).toStrictEqual([]);
       expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[]]"`);
-      expect(externalNames).toStrictEqual([]);
+      expect(externalNames).toMatchInlineSnapshot(`{}`);
     }),
   );
 
@@ -57,7 +57,11 @@ describe('transpileFn', () => {
       expect(JSON.stringify(body)).toMatchInlineSnapshot(
         `"[0,[[10,[1,[1,"a","+","b"],"-","c"]]]]"`,
       );
-      expect(externalNames).toStrictEqual(['c']);
+      expect(externalNames).toMatchInlineSnapshot(`
+        {
+          "c": "c",
+        }
+      `);
     }),
   );
 
@@ -76,7 +80,11 @@ describe('transpileFn', () => {
         `"[0,[[13,"a",[5,"0"]],[2,"c","=",[1,"a","+",[5,"2"]]]]]"`,
       );
       // Only 'c' is external, as 'a' is declared in the same scope.
-      expect(externalNames).toStrictEqual(['c']);
+      expect(externalNames).toMatchInlineSnapshot(`
+        {
+          "c": "c",
+        }
+      `);
     }),
   );
 
@@ -97,7 +105,11 @@ describe('transpileFn', () => {
         `"[0,[[13,"a",[5,"0"]],[0,[[2,"c","=",[1,"a","+",[5,"2"]]]]]]]"`,
       );
       // Only 'c' is external, as 'a' is declared in the outer scope.
-      expect(externalNames).toStrictEqual(['c']);
+      expect(externalNames).toMatchInlineSnapshot(`
+        {
+          "c": "c",
+        }
+      `);
     }),
   );
 
@@ -111,7 +123,15 @@ describe('transpileFn', () => {
         `"[0,[[10,[7,[7,"external","outside"],"prop"]]]]"`,
       );
       // Only 'external' is external.
-      expect(externalNames).toStrictEqual(['external']);
+      expect(externalNames).toMatchInlineSnapshot(`
+        {
+          "external": {
+            "outside": {
+              "prop": "external.outside.prop",
+            },
+          },
+        }
+      `);
     }),
   );
 
@@ -140,7 +160,7 @@ describe('transpileFn', () => {
         },
       ]);
 
-      expect(externalNames).toStrictEqual([]);
+      expect(externalNames).toMatchInlineSnapshot(`{}`);
     }),
   );
 
@@ -186,7 +206,7 @@ describe('transpileFn', () => {
         },
       ]);
 
-      expect(externalNames).toStrictEqual([]);
+      expect(externalNames).toMatchInlineSnapshot(`{}`);
     }),
   );
 
@@ -195,4 +215,44 @@ describe('transpileFn', () => {
 
     expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[[10,[7,"x","y"]]]]"`);
   });
+
+  it(
+    'handles complex external trees',
+    dualTest((p) => {
+      const { externalNames } = transpileFn(
+        p(`() => {
+          const a = ext.p;
+          const b = ext.q.a;
+          const c = ext.q.b;
+
+          const d = ext.r.a;
+          const e = ext.r;
+
+          const f = ext.s;
+          const g = ext.s.a;
+
+          const h = ext.t.fn().x;
+          const i = ext.t.comp['computed'].x;
+        }`),
+      );
+
+      expect(externalNames).toMatchInlineSnapshot(`
+        {
+          "ext": {
+            "p": "ext.p",
+            "q": {
+              "a": "ext.q.a",
+              "b": "ext.q.b",
+            },
+            "r": "ext.r",
+            "s": "ext.s",
+            "t": {
+              "comp": "ext.t.comp",
+              "fn": "ext.t.fn",
+            },
+          },
+        }
+      `);
+    }),
+  );
 });

packages/typegpu/src/core/function/fnCore.ts

@@ -3,7 +3,7 @@ import { undecorate } from '../../data/dataTypes.ts';
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
 import { type BaseData, isWgslData, isWgslStruct, Void } from '../../data/wgslTypes.ts';
 import { MissingLinksError } from '../../errors.ts';
-import { getMetaData, getName } from '../../shared/meta.ts';
+import { getMetaData, getName, type Externals2 } from '../../shared/meta.ts';
 import { $getNameForward } from '../../shared/symbols.ts';
 import type { ResolutionCtx } from '../../types.ts';
 import { applyExternals, type ExternalMap, replaceExternalsInWgsl } from '../resolve/externals.ts';
@@ -148,11 +148,14 @@ export function createFnCore(implementation: Implementation, fnAttribute = ''):
       const pluginData = getMetaData(implementation);
 
       // Passing a record happens prior to version 0.9.0
+      // Passing a function happens prior to version 0.12.0
       // TODO: Support for this can be removed down the line
-      const pluginExternals =
-        typeof pluginData?.externals === 'function'
-          ? pluginData.externals()
-          : pluginData?.externals;
+      let pluginExternals: ExternalMap | Record<string, unknown> | undefined =
+        pluginData?.externals2
+          ? externals2ToExternalMap(pluginData.externals2)
+          : typeof pluginData?.externals === 'function'
+            ? pluginData.externals()
+            : pluginData?.externals;
 
       if (pluginExternals) {
         const missing = Object.fromEntries(
@@ -170,7 +173,9 @@ export function createFnCore(implementation: Implementation, fnAttribute = ''):
       }
 
       // verify all required externals are present
-      const missingExternals = ast.externalNames.filter((name) => !(name in externalMap));
+      const missingExternals = Object.keys(ast.externalNames).filter(
+        (name) => !(name in externalMap),
+      );
       if (missingExternals.length > 0) {
         throw new MissingLinksError(getName(this), missingExternals);
       }
@@ -218,6 +223,22 @@ export function createFnCore(implementation: Implementation, fnAttribute = ''):
   return core;
 }
 
+// TODO: deslopify, document, make sure it works as intended
+function externals2ToExternalMap(ext2: Externals2): ExternalMap {
+  const result: ExternalMap = {};
+  for (const [key, value] of Object.entries(ext2)) {
+    if (typeof value === 'function') {
+      Object.defineProperty(result, key, {
+        get: value,
+        enumerable: true,
+      });
+    } else {
+      result[key] = externals2ToExternalMap(value);
+    }
+  }
+  return result;
+}
+
 function isArgUsedInBody(argName: string, body: string): boolean {
   return new RegExp(`\\b${argName}\\b`).test(body);
 }

packages/typegpu/src/core/function/fnTypes.ts

@@ -1,4 +1,3 @@
-import type * as tinyest from 'tinyest';
 import type { BuiltinClipDistances } from '../../builtin.ts';
 import type { AnyAttribute } from '../../data/attributes.ts';
 import type {
@@ -27,19 +26,6 @@ import type { InferGPU } from '../../shared/repr.ts';
 
 export type AnyFn = (...args: never[]) => unknown;
 
-/**
- * Information extracted from transpiling a JS function.
- */
-export type TranspilationResult = {
-  params: tinyest.FuncParameter[];
-  body: tinyest.Block;
-  /**
-   * All identifiers found in the function code that are not declared in the
-   * function itself, or in the block that is accessing that identifier.
-   */
-  externalNames: string[];
-};
-
 export type InferArgs<T extends unknown[]> = {
   [Idx in keyof T]: InferGPU<T[Idx]>;
 };