Implement smarter expression handling for lazy refs

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

@@ -15,6 +15,7 @@ import type { Transformer } from '../transform';
 import type { Logger } from '@openfn/logger';
 
 const globals = [
+  // TODO I think we can remove this now!!
   '\\$', // TMP hack to fix a problem with lazy-state (needs double escaping to work)
 
   'AggregateError',

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

@@ -3,61 +3,67 @@
  *
  * Converts all $.a.b chains unless:
  * - $ 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';
 
 // Walk up the AST and work out where the parent arrow function should go
-// This should be on the argument to the operation. So the top of whatever
-// structure we are in
-// If there's already a (state) wrapper, happily do nothing
-// if there's anything else, throw an error
-const ensureParentArrow = (path: NodePath<namedTypes.MemberExpression>) => {
-  let root;
-
-  // find the parenting call expression
-  // find the matching argument (which argument is == last?)
-  // maybe wrap the argument
-
+const ensureParentArrow = (path: NodePath<n.MemberExpression>) => {
+  let root = path;
   let last;
-  let callexpr;
-  let arg;
-  // while (!root) {
-  //   //
 
-  //   root = root.parent;
+  // 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;
 
-  // Now nest the whole thing in an arrow
-  const params = b.identifier('state');
-  const arrow = b.arrowFunctionExpression([params], path.node);
-  arg.replace(arrow);
-};
+    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<any>) => {
+const isOpenFunction = (path: NodePath) => {
   // is it a function?
-  //   -return false
-  // does it have one param?
-  //  continue else throw
-  // is the param called state?
-  //  return true else throw
+  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<namedTypes.MemberExpression>) {
+function visitor(path: NodePath<n.MemberExpression>) {
   let first = path.node.object;
   while (first.hasOwnProperty('object')) {
-    first = (first as namedTypes.MemberExpression).object;
+    first = (first as n.MemberExpression).object;
   }
 
-  let firstIdentifer = first as namedTypes.Identifier;
+  let firstIdentifer = first as n.Identifier;
 
   if (first && firstIdentifer.name === '$') {
     // But if a $ declared a parent scope, ignore it
@@ -72,7 +78,8 @@ function visitor(path: NodePath<namedTypes.MemberExpression>) {
     // rename $ to state
     firstIdentifer.name = 'state';
 
-    ensureParentArrow();
+    // from the parenting member expression, ensure the parent arrow is nicely wrapped
+    ensureParentArrow(path);
   }
 
   // Stop parsing this member expression

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

@@ -9,7 +9,6 @@ import visitors from '../../src/transforms/lazy-state';
 
 test('convert a simple dollar reference', (t) => {
   const ast = parse('get($.data)');
-
   const transformed = transform(ast, [visitors]);
   const { code } = print(transformed)
 
@@ -34,7 +33,82 @@ test('ignore a regular chain reference', (t) => {
   t.is(code, 'get(a.b.c.d)')
 })
 
-test('ignore a string', (t) => {
+test('convert a template literal', (t) => {
+  const src = 'get(`hello ${$.data}`)'
+  t.log(src)
+  const ast = parse(src);
+  const transformed = transform(ast, [visitors]);
+  const { code } = print(transformed)
+  t.log(code)
+
+  t.is(code, 'get(state => `hello ${state.data}`)')
+})
+
+test('convert a template literal with two refs', (t) => {
+  const src = 'get(`hello ${$.firstname} ${$.lastname}`)'
+  t.log(src)
+  const ast = parse(src);
+  const transformed = transform(ast, [visitors]);
+  const { code } = print(transformed)
+  t.log(code)
+
+  t.is(code, 'get(state => `hello ${state.firstname} ${state.lastname}`)')
+})
+
+test('convert a template literal with a pre-existing parent arrow', (t) => {
+  const src = 'get(state => `hello ${$.data}`)'
+  t.log(src)
+  const ast = parse(src);
+  const transformed = transform(ast, [visitors]);
+  const { code } = print(transformed)
+  t.log(code)
+
+  t.is(code, 'get(state => `hello ${state.data}`)')
+})
+
+test('throw if a $ is already inside a non-compatible arrow (state name)', (t) => {
+  const src = 'get((s) => `hello ${$.data}`)' // throw!!
+  t.log(src)
+  const ast = parse(src);
+
+  t.throws(() => transform(ast, [visitors]), {
+    message: `invalid lazy state: parameter "s" should be called "state"`
+  });
+})
+
+test('throw if a $ is already inside a non-compatible arrow (arity)', (t) => {
+  const src = 'get((state, b) => `hello ${$.data}`)' // throw!!
+  t.log(src)
+  const ast = parse(src);
+
+  t.throws(() => transform(ast, [visitors]), {
+    message: 'invalid lazy state: parent has wrong arity'
+  });
+})
+
+test('wrap a concatenation', (t) => {
+  const src = 'get($.firstname + " " + $.lastname)'
+  t.log(src)
+  const ast = parse(src);
+  const transformed = transform(ast, [visitors]);
+  const { code } = print(transformed)
+  t.log(code)
+
+  t.is(code, 'get(state => state.firstname + " " + state.lastname)')
+})
+
+test('wrap a dynamic property reference', (t) => {
+  const src = 'get(city[$.location])'
+  t.log(src)
+  const ast = parse(src);
+  const transformed = transform(ast, [visitors]);
+  const { code } = print(transformed)
+  t.log(code)
+
+  t.is(code, 'get(state => city[state.location])')
+})
+
+test('ignore a dollar ref in a string', (t) => {
   const ast = parse('get("$.a.b")');
 
   const transformed = transform(ast, [visitors]);