Preliminary babel safe transform impl

src/index.lsc

@@ -14,7 +14,7 @@ import { getShadowingIdentifiers, checkFalsePositiveReference } from "./variable
 import { ensureConstructor, bindMethodsInConstructor, bindMethods } from "./classes";
 import * as matching from "./match";
 import * as comprehension from "./comprehension";
-import { replaceWithSafeCall, transformSafeMemberExpression, transformExistentialExpression } from "./safe";
+import { replaceWithSafeCall, transformSafeMemberExpression, transformExistentialExpression, maybeTransformSafe } from "./safe";
 import { replaceWithInlinedOperator } from "./inlinedOperators";
 import { transformForInArrayStatement, transformForInObjectStatement, lintForInArrayStatement, lintForInObjectStatement } from "./for";
 import { resetHelpers } from "./helpers";
@@ -90,27 +90,21 @@ Lightscript(babel) ->
           comprehension.transformPlainObjectComprehension(path)
 
       TildeCallExpression: {
-        // run on exit instead of enter so that SafeMemberExpression
-        // can process differently from a wrapping CallExpression
-        // eg; `a?.b~c()` -> `a == null ? null : c(a.b)`
-        exit(path): void ->
+        enter(path): void ->
           { node } = path
           args = [ node.left, ...node.arguments ]
           callExpr = t.callExpression(node.right, args)~atNode(node)
-
-          if path~replaceWithInlinedOperator(node.right, args): return
-
-          if (node.optional):
-            path~replaceWithSafeCall(callExpr)
-          else:
-            path.replaceWith(callExpr)
+          if node.optional: callExpr.optional = true
+          path.replaceWith(callExpr)
       }
 
       CallExpression: {
+        enter(path): void ->
+          maybeTransformSafe(path)
+
         exit(path): void ->
           { node } = path
           if replaceWithInlinedOperator(path, node.callee, node.arguments): return
-          if node.optional: replaceWithSafeCall(path, node)
       }
 
       NamedArrowFunction(path): void ->
@@ -239,7 +233,10 @@ Lightscript(babel) ->
         path.replaceWith(awaitExpr)
 
       MemberExpression(path): void ->
-        if path.node.optional: transformSafeMemberExpression(path)
+        path~maybeTransformSafe()
+
+      NewExpression(path): void ->
+        path~maybeTransformSafe()
 
       AwaitExpression(path): void ->
         if (path.get("argument").isArrayExpression() || path.node.argument.type === "ArrayComprehension") {

src/safe.lsc

@@ -6,6 +6,94 @@ import { ensureBlockArrowFunctionExpression } from './functions'
 
 import { getLoc, placeTreeAtLocWhenUnplaced as allAtLoc } from 'ast-loc-utils'
 
+// Optional-replacement algorithm based on babel-plugin-transform-optional-chaining
+// See https://github.com/babel/babel/blob/master/packages/babel-plugin-transform-optional-chaining/src/index.js
+
+findReplacementPath(path) ->
+  path.find(path ->
+    { key, parentPath } = path
+    if key == "left" and parentPath~isa("AssignmentExpression"): false
+    elif key == "object" and parentPath~isa("MemberExpression"): false
+    elif key == "callee" and (parentPath~isa("CallExpression") or parentPath~isa("NewExpression")): false
+    elif key == "argument" and parentPath~isa("UpdateExpression"): false
+    elif key == "argument" and (parentPath~isa("UnaryExpression") and parentPath.node.operator == "delete"): false
+    else: true
+  )
+
+replaceOptionals(path, replacementPath, loose = true): void ->
+  { scope } = path
+  optionals = []
+  nil = scope.buildUndefinedNode()
+
+  // Collect all optional nodes within the local cluster of nodes
+  let objectPath = path
+  while objectPath~isa("MemberExpression") or objectPath~isa("CallExpression") or objectPath~isa("NewExpression"):
+    { node } = objectPath;
+    if node.optional: optionals.push(node)
+
+    if objectPath~isa("MemberExpression"):
+      now objectPath = objectPath.get("object")
+    else:
+      now objectPath = objectPath.get("callee")
+
+  // Traverse optionals from innermost to outermost
+  for let i = optionals.length - 1; i >= 0; i--:
+    node = optionals[i]
+    node.optional = false
+
+    isCall = node~isa("CallExpression");
+    replaceKey = if isCall or node~isa("NewExpression"): "callee" else: "object"
+    chain = node[replaceKey]
+
+    let ref, check
+    if loose and isCall:
+      // If we are using a loose transform (avoiding a Function#call) and we are at the call,
+      // we can avoid a needless memoize.
+      now check = now ref = chain;
+    else:
+      now ref = scope.maybeGenerateMemoised(chain);
+      if (ref) {
+        now check = t.assignmentExpression("=", ref, chain)
+        node[replaceKey] = ref;
+      } else {
+        now check = now ref = chain;
+      }
+
+    // Ensure call expressions have the proper `this`
+    // `foo.bar()` has context `foo`.
+    if isCall && chain~isa("MemberExpression"):
+      if loose:
+        // To avoid a Function#call, we can instead re-grab the property from the context object.
+        // `a.?b.?()` translates roughly to `_a.b != null && _a.b()`
+        node.callee = chain
+      else:
+        // Otherwise, we need to memoize the context object, and change the call into a Function#call.
+        // `a.?b.?()` translates roughly to `(_b = _a.b) != null && _b.call(_a)`
+        { object } = chain
+        let context = scope.maybeGenerateMemoised(object);
+        if (context) {
+          chain.object = t.assignmentExpression("=", context, object);
+        } else {
+          now context = object
+        }
+
+        node.arguments.unshift(context)
+        node.callee = t.memberExpression(node.callee, t.identifier("call"))
+
+    replacementPath.replaceWith(
+      t.conditionalExpression(
+        t.binaryExpression("==", check, t.nullLiteral()),
+        nil,
+        replacementPath.node,
+      ),
+    )
+
+    now replacementPath = replacementPath.get("alternate")
+
+export maybeTransformSafe(path): void ->
+  if path.node.optional:
+    replaceOptionals(path, path~findReplacementPath(), true)
+
 export transformExistentialExpression(path) ->
   path.replaceWith(
     t.binaryExpression(