Merge pull request #1053 from OpenFn/ignore-top-object-add-import

feat: Compiler optimisations

Author: josephjclark

.changeset/floppy-comics-work.md

@@ -0,0 +1,6 @@
+---
+'@openfn/compiler': patch
+'@openfn/ws-worker': patch
+---
+
+Improve compiler performance for large files

packages/compiler/src/transform-ignore.ts

@@ -0,0 +1,20 @@
+import { namedTypes as n } from 'ast-types';
+import { NodePath } from 'ast-types/lib/node-path';
+
+const IgnoreRules = (path: NodePath<n.Expression>) => {
+  return {
+    check: n.ObjectExpression.check(path.node),
+    shouldSkip: () => {
+      if (
+        n.VariableDeclarator.check(path.parentPath.node) &&
+        n.VariableDeclaration.check(path.parentPath.parentPath.node) &&
+        n.Program.check(path.parentPath.parentPath.parentPath.parentPath.node)
+      ) {
+        return true;
+      }
+      return false;
+    },
+  };
+};
+
+export default IgnoreRules;

packages/compiler/src/transform.ts

@@ -68,7 +68,6 @@ export default function transform(
     ] as Transformer[];
   }
   const logger = options.logger || defaultLogger;
-
   transformers
     // Ignore transformers which are explicitly disabled
     .filter(({ id }) => options[id] ?? true)
@@ -82,11 +81,13 @@ export default function transform(
     })
     // Run each transformer
     .forEach(({ id, types, visitor }) => {
+      const t = `transformer ${id}`;
+      console.time?.(t);
       const astTypes: Visitor = {};
       for (const type of types) {
         const name = `visit${type}` as keyof Visitor;
         astTypes[name] = function (path: NodePath) {
-          printHeap(`visit: ${path.name}`);
+          // printHeap(`visit: ${path.name}`);
           const opts = options[id] || {};
           const abort = visitor!(path, logger, opts);
           if (abort) {
@@ -98,8 +99,15 @@ export default function transform(
 
       // @ts-ignore
       visit(ast, astTypes);
+      if (options.trace) {
+        console.log();
+        printHeap(`${t}`);
+        console.timeEnd?.(t);
+      }
     });
-
+  if (options.trace) {
+    console.log();
+  }
   printHeap(`finished`);
   const duration = (Date.now() - start) / 1000;
   logger.debug(`Finished in ${duration}s`);

packages/compiler/src/transforms/add-imports.ts

@@ -107,6 +107,17 @@ export type IdentifierList = Record<string, true>;
 export function findAllDanglingIdentifiers(ast: ASTNode) {
   const result: IdentifierList = {};
   visit(ast, {
+    visitObjectExpression(path) {
+      // excaping top-level object variable declarations.
+      if (
+        n.VariableDeclarator.check(path.parentPath.node) &&
+        n.VariableDeclaration.check(path.parentPath.parentPath.node) &&
+        n.Program.check(path.parentPath.parentPath.parentPath.parentPath.node)
+      ) {
+        return false;
+      }
+      this.traverse(path);
+    },
     visitIdentifier: function (path) {
       // If this is part of an import statement, do nothing
       if (n.ImportSpecifier.check(path.parent.node)) {

packages/compiler/src/transforms/lazy-state.ts

@@ -9,6 +9,7 @@
 import { builders as b, namedTypes as n } from 'ast-types';
 import type { NodePath } from 'ast-types/lib/node-path';
 import type { Transformer } from '../transform';
+import IgnoreRules from '../transform-ignore';
 
 // Walk up the AST and work out where the parent arrow function should go
 const ensureParentArrow = (path: NodePath<n.MemberExpression>) => {
@@ -79,6 +80,11 @@ const isOpenFunction = (path: NodePath) => {
 };
 
 function visitor(path: NodePath<n.MemberExpression>) {
+  // if it was called for an ObjectExpression
+  const ignoreRule = IgnoreRules(path);
+  if (ignoreRule.check) {
+    return ignoreRule.shouldSkip();
+  }
   let first = path.node.object;
   while (first.hasOwnProperty('object')) {
     first = (first as n.MemberExpression).object;
@@ -106,7 +112,7 @@ function visitor(path: NodePath<n.MemberExpression>) {
 
 export default {
   id: 'lazy-state',
-  types: ['MemberExpression'],
+  types: ['MemberExpression', 'ObjectExpression'],
   visitor,
   // It's important that $ symbols are escaped before any other transformations can run
   order: 0,

packages/compiler/src/transforms/promises.ts

@@ -1,6 +1,7 @@
 import { namedTypes as n, builders as b } from 'ast-types';
 
 import type { NodePath } from 'ast-types/lib/node-path';
+import IgnoreRules from '../transform-ignore';
 
 const NO_DEFER_DECLARATION_ERROR = 'No defer declaration found';
 
@@ -146,6 +147,10 @@ const isTopScope = (path: NodePath<any>) => {
 };
 
 const visitor = (path: NodePath<n.CallExpression>) => {
+  const ignoreRule = IgnoreRules(path);
+  if (ignoreRule.check) {
+    return ignoreRule.shouldSkip();
+  }
   let root: NodePath<any> = path;
   while (!n.Program.check(root.node)) {
     root = root.parent;
@@ -166,7 +171,7 @@ const visitor = (path: NodePath<n.CallExpression>) => {
 
 export default {
   id: 'promises',
-  types: ['CallExpression'],
+  types: ['CallExpression', 'ObjectExpression'],
   visitor,
   // this should run before top-level operations are moved into the exports array
   order: 0,

packages/compiler/src/transforms/top-level-operations.ts

@@ -8,71 +8,53 @@ import type { Transformer } from '../transform';
 // Note that the validator should complain if it see anything other than export default []
 // What is the relationship between the validator and the compiler?
 
-export type ExtendedProgram = NodePath<namedTypes.Program> & {
-  operations: Array<{ line: number; name: string; order: number }>;
-};
+export type ExtendedProgram = NodePath<
+  namedTypes.Program & {
+    operations: Array<{ line: number; name: string; order: number }>;
+  }
+>;
 
 export type TopLevelOpsOptions = {
   // Wrap operations in a `(state) => op` wrapper
   wrap: boolean; // TODO
 };
 
-function visitor(path: NodePath<namedTypes.CallExpression>) {
-  const root = path.parent.parent.node;
-  // Check if the node is a top level Operation
+function visitor(programPath: ExtendedProgram) {
+  const operations: Array<{ line: number; name: string; order: number }> = [];
+  const children = programPath.node.body;
+  const rem = [];
+
+  const target = programPath.node.body.at(-1);
   if (
-    // Check this is a top level call expression
-    // ie, the parent must be an ExpressionStatement, and the statement's parent must be a Program
-    n.Program.check(root) &&
-    n.Statement.check(path.parent.node) &&
-    // An Operation could be an Identifier (fn()) or Member Expression (http.get())
-    (n.Identifier.check(path.node.callee) ||
-      n.MemberExpression.check(path.node.callee))
+    n.ExportDefaultDeclaration.check(target) &&
+    n.ArrayExpression.check(target.declaration)
   ) {
-    appendOperationMetadata(path, root as unknown as ExtendedProgram);
-
-    // Now Find the top level exports array
-    const target = root.body.at(-1);
-    if (
-      n.ExportDefaultDeclaration.check(target) &&
-      n.ArrayExpression.check(target.declaration)
-    ) {
-      const statement = path.parent;
-      target.declaration.elements.push(path.node);
-      // orphan the original expression statement parent
-      statement.prune();
-    } else {
-      // error! there isn't an appropriate export statement
-      // What do we do?
+    for (const child of children) {
+      if (
+        n.ExpressionStatement.check(child) &&
+        n.CallExpression.check(child.expression)
+      ) {
+        const order = operations.length + 1;
+        // @ts-ignore
+        const name = child.expression.callee.name;
+        const line = child.expression.loc?.start.line ?? -1;
+        operations.push({ name, line, order });
+        target.declaration.elements.push(child.expression as any);
+      } else rem.push(child);
     }
+    programPath.node.body = rem;
+  } else {
+    // error! there isn't an appropriate export statement
+    // What do we do?
   }
+  programPath.node.operations = operations;
 
   // if not (for now) we should cancel traversal
   return true;
 }
 
-// Add metadata to each operation for:
-// - the order (do we know? We can guess from the state of the exports array
-// - the name of the operation
-// - the original line number (do we know?)
-function appendOperationMetadata(
-  path: NodePath<namedTypes.CallExpression>,
-  root: ExtendedProgram
-) {
-  if (!root.operations) {
-    root.operations = [];
-  }
-
-  const { operations } = root;
-  const order = operations.length + 1;
-  const name = path.value.callee.name;
-  const line = path.value.loc?.start.line ?? -1;
-
-  operations.push({ line, order, name });
-}
-
 export default {
   id: 'top-level-operations',
-  types: ['CallExpression'],
+  types: ['Program'],
   visitor,
-} as Transformer;
+} as unknown as Transformer;

packages/compiler/src/util.ts

@@ -5,10 +5,10 @@ import path from 'node:path';
 import { Project, describeDts } from '@openfn/describe-package';
 import type { Logger } from '@openfn/logger';
 
-export function heap(reason: string, logger?: Logger) {
+export function heap(reason: string, _logger?: Logger) {
   const { used_heap_size } = getHeapStatistics();
   const mb = used_heap_size / 1024 / 1024;
-  logger?.debug(`[${reason}] Used heap at ${mb.toFixed(2)}mb`);
+  console.log(`[${reason}] Used heap at ${mb.toFixed(2)}mb`);
 }
 
 export const loadFile = (filePath: string) =>

packages/compiler/test/transforms/add-imports.test.ts

@@ -147,7 +147,23 @@ test('findAllDanglingIdentifiers: const x = { a: 10 };', (t) => {
   t.assert(Object.keys(result).length == 0);
 });
 
-test('findAllDanglingIdentifiers: const x = { a: b };', (t) => {
+test('findAllDanglingIdentifiers: ignore top-level objects declarations', (t) => {
+  // ignored because [bank] isn't picked as dangling but [foo] is
+  const ast = parse(`
+const bigObj = {name: "ignore-me", age: bank};
+fn(state=> {
+    const update = {
+      name: "pick-me",
+      age: foo
+    }
+    return update
+})
+    `);
+  const result = findAllDanglingIdentifiers(ast);
+  t.deepEqual(result, { fn: true, foo: true });
+});
+
+test.skip('findAllDanglingIdentifiers: const x = { a: b };', (t) => {
   const ast = parse('const x = { a: b };');
   const result = findAllDanglingIdentifiers(ast);
   t.assert(Object.keys(result).length == 1);

packages/compiler/test/transforms/lazy-state.test.ts

@@ -15,6 +15,20 @@ test('convert a simple dollar reference', (t) => {
   t.is(code, 'get(state => state.data)');
 });
 
+test("don't visit top-level object", (t) => {
+  const ast = parse(`const x = { x: $.data };
+fn($.data)`);
+  t.notThrows(() => {
+    const transformed = transform(ast, [visitors]);
+    const { code } = print(transformed);
+    t.is(
+      code,
+      `const x = { x: $.data };
+fn(state => state.data)`
+    );
+  });
+});
+
 test('convert a chained dollar reference', (t) => {
   const ast = parse('get($.a.b.c.d)');