feat: optimize top level operations

packages/compiler/src/transform.ts

@@ -82,7 +82,7 @@ export default function transform(
     // Run each transformer
     .forEach(({ id, types, visitor }) => {
       const t = `transformer ${id}`;
-      console.time(t);
+      console.time?.(t);
       const astTypes: Visitor = {};
       for (const type of types) {
         const name = `visit${type}` as keyof Visitor;
@@ -102,7 +102,7 @@ export default function transform(
       if (options.trace) {
         console.log();
         printHeap(`${t}`);
-        console.timeEnd(t);
+        console.timeEnd?.(t);
       }
     });
   if (options.trace) {

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/test/transforms/top-level-operations.test.ts

@@ -217,7 +217,7 @@ test('should only take the top of a nested operation call (and preserve its argu
 
 test('appends an operations map to simple operation', (t) => {
   // We have to parse source here rather than building an AST so that we get positional information
-  const { program } = parse(`fn();`);
+  const { program } = parse(`fn(); export default [];`);
 
   transform(program, [visitors]);