Merge pull request #658 from OpenFn/lazy-state-expressions

Lazy state expressions

.changeset/afraid-seals-share.md

@@ -0,0 +1,5 @@
+---
+'@openfn/ws-worker': minor
+---
+
+Initial release of "lazy state" operators ($)

.changeset/five-ligers-suffer.md

@@ -0,0 +1,5 @@
+---
+'@openfn/cli': minor
+---
+
+Initial release of "lazy state" operators ($)

.changeset/nasty-frogs-join.md

@@ -0,0 +1,5 @@
+---
+'@openfn/compiler': patch
+---
+
+Allow lazy state functions to be hoisted further up the tree (lazy state expressions)

packages/compiler/src/transform.ts

@@ -27,10 +27,9 @@ export type Transformer = {
   id: TransformerName;
   types: string[];
   visitor: TransformFunction;
+  order?: number;
 };
 
-type TransformerIndex = Partial<Record<keyof Visitor, Transformer[]>>;
-
 export type TransformOptions = {
   logger?: Logger; //  TODO maybe in the wrong place?
 
@@ -49,61 +48,44 @@ export default function transform(
   options: TransformOptions = {}
 ) {
   if (!transformers) {
-    transformers = [lazyState, ensureExports, topLevelOps, addImports] as Transformer[];
+    transformers = [
+      lazyState,
+      ensureExports,
+      topLevelOps,
+      addImports,
+    ] as Transformer[];
   }
   const logger = options.logger || defaultLogger;
-  const transformerIndex = indexTransformers(transformers, options);
-
-  const v = buildVisitors(transformerIndex, logger, options);
-  // @ts-ignore generic disagree on Visitor, so disabling type checking for now
-  visit(ast, v);
 
-  return ast;
-}
-
-// Build a map of AST node types against an array of transform functions
-export function indexTransformers(
-  transformers: Transformer[],
-  options: TransformOptions = {}
-): TransformerIndex {
-  const index: TransformerIndex = {};
-  for (const t of transformers) {
-    const { types, id } = t;
-    if (options[id] !== false) {
+  transformers
+    // Ignore transformers which are explicitly disabled
+    .filter(({ id }) => options[id] ?? true)
+    // Set default orders
+    .map((t) => ({ ...t, order: t.order ?? 1 }))
+    // Sort by order
+    .sort((a, b) => {
+      if (a.order > b.order) return 1;
+      if (a.order < b.order) return -1;
+      return 0;
+    })
+    // Run each transformer
+    .forEach(({ id, types, visitor }) => {
+      const astTypes: Visitor = {};
       for (const type of types) {
         const name = `visit${type}` as keyof Visitor;
-        if (!index[name]) {
-          index[name] = [];
-        }
-        index[name]!.push(t);
+        astTypes[name] = function (path: NodePath) {
+          const opts = options[id] || {};
+          const abort = visitor!(path, logger, opts);
+          if (abort) {
+            return false;
+          }
+          this.traverse(path);
+        };
       }
-    }
-  }
-  return index;
-}
 
-// Build an index of AST visitors, where each node type is mapped to a visitor function which
-// calls out to the correct transformer, passing a logger and options
-export function buildVisitors(
-  transformerIndex: TransformerIndex,
-  logger: Logger,
-  options: TransformOptions = {}
-) {
-  const visitors: Visitor = {};
+      // @ts-ignore
+      visit(ast, astTypes);
+    });
 
-  for (const k in transformerIndex) {
-    const astType = k as keyof Visitor;
-    visitors[astType] = function (path: NodePath) {
-      const transformers = transformerIndex[astType]!;
-      for (const { id, visitor } of transformers) {
-        const opts = options[id] || {};
-        const abort = visitor!(path, logger, opts);
-        if (abort) {
-          return false;
-        }
-      }
-      this.traverse(path);
-    };
-  }
-  return visitors;
+  return ast;
 }

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

@@ -15,8 +15,6 @@ import type { Transformer } from '../transform';
 import type { Logger } from '@openfn/logger';
 
 const globals = [
-  '\\$', // TMP hack to fix a problem with lazy-state (needs double escaping to work)
-
   'AggregateError',
   'Array',
   'ArrayBuffer',

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

@@ -1,27 +1,71 @@
 /*
  * Convert $.a.b.c references into (state) => state.a.b.c
- * 
+ *
  * Converts all $.a.b chains unless:
- * - $ was assigned previously in that scope 
+ * - $ was assigned previously in that scope
  * 
- * TODO (maybe):
- *  - only convert $-expressions which are arguments to operations (needs type defs)
- *  - warn if converting a non-top-level $-expression
- *  - if not top level, convert to state.a.b.c (ie don't wrap the function)
+ *
  */
-import { builders as b, namedTypes } from 'ast-types';
+import { builders as b, namedTypes as n} from 'ast-types';
 import type { NodePath } from 'ast-types/lib/node-path';
 import type { Transformer } from '../transform';
 
-function visitor(path: NodePath<namedTypes.MemberExpression>) {
+// Walk up the AST and work out where the parent arrow function should go
+const ensureParentArrow = (path: NodePath<n.MemberExpression>) => {
+  let root = path;
+  let last;
+
+  // find the parenting call expression
+  // Ie, the operation we're passing this arrow into
+  while(root && !n.CallExpression.check(root.node)) {
+    last = root;
+    root = root.parent;
+  }
+
+  if (root) {
+    const arg = last as NodePath;
+
+    if (!isOpenFunction(arg)) {
+      const params = b.identifier('state');
+      const arrow = b.arrowFunctionExpression([params], arg.node);
+      arg.replace(arrow);
+    }
+  }
+}
+
+// Checks whether the passed node is an open function, ie, (state) => {...}
+const isOpenFunction = (path: NodePath) => {
+  // is it a function?
+  if (n.ArrowFunctionExpression.check(path.node))  {
+    const arrow = path.node as n.ArrowFunctionExpression;
+    // does it have one param?
+    if(arrow.params.length == 1) {
+      const name = (arrow.params[0] as n.Identifier).name
+      // is the param called state?
+      if (name === "state") {
+        // We already have a valid open function here
+        return true;
+      }
+      throw new Error(`invalid lazy state: parameter "${name}" should be called "state"`)
+    }
+    throw new Error('invalid lazy state: parent has wrong arity')
+  }
+
+  // if we get here, then path is:
+  // a) a Javascript Expression (and not an arrow)
+  // b) appropriate for being wrapped in an arrow
+  return false;
+};
+
+function visitor(path: NodePath<n.MemberExpression>) {
   let first = path.node.object;
-  while(first.hasOwnProperty('object')) {
-    first = (first as namedTypes.MemberExpression).object;
+  while (first.hasOwnProperty('object')) {
+    first = (first as n.MemberExpression).object;
   }
 
-  let firstIdentifer = first as namedTypes.Identifier;
-  
-  if (first && firstIdentifer.name === "$") {
+  let firstIdentifer = first as n.Identifier;
+
+  if (first && firstIdentifer.name === '$') {
     // But if a $ declared a parent scope, ignore it
     let scope = path.scope;
     while (scope) {
@@ -32,15 +76,10 @@ function visitor(path: NodePath<namedTypes.MemberExpression>) {
     }
 
     // rename $ to state
-    firstIdentifer.name = "state";
-
-    // Now nest the whole thing in an arrow
-    const params = b.identifier('state')
-    const arrow = b.arrowFunctionExpression(
-      [params],
-      path.node
-    )
-    path.replace(arrow)
+    firstIdentifer.name = 'state';
+
+    // from the parenting member expression, ensure the parent arrow is nicely wrapped
+    ensureParentArrow(path);
   }
 
   // Stop parsing this member expression
@@ -51,4 +90,6 @@ export default {
   id: 'lazy-state',
   types: ['MemberExpression'],
   visitor,
+  // It's important that $ symbols are escaped before any other transformations can run
+  order: 0,
 } as Transformer;