Fix for SafeMemberExpressions in ArrowFunctionExpressions

See wcjohnson/lightscript#23

.gitignore

@@ -4,3 +4,4 @@ lib
 /stdlib.js
 npm-debug.log
 lerna-debug.log
+/test/fixtures-debug

src/functions.lsc

@@ -71,6 +71,18 @@ export toArrowFunction(node) ->
 export replaceWithArrowFunction(path) ->
   path.replaceWith(toArrowFunction(path.node))
 
+export toBlockArrowFunctionExpression(node) ->
+  let { params, body, async } = node
+  fn = t.arrowFunctionExpression(params, body~toBlockStatement(), async)~atNode(node)
+  if node.returnType: fn.returnType = node.returnType
+  if node.typeParameters: fn.typeParameters = node.typeParameters
+
+  fn
+
+export ensureBlockArrowFunctionExpression(path) ->
+  if path.node~isa("ArrowFunctionExpression"):
+    path.replaceWith(path.node~toBlockArrowFunctionExpression())
+
 // Replace a functional node with a new node binding the function to `this`.
 export replaceWithBoundFunction(path) ->
   { node } = path

src/safe.lsc

@@ -1,7 +1,8 @@
 import t from './types'
 import atNode from 'ast-loc-utils/lib/placeAtNode'
 import { hoistRef } from './ref'
-import is from './is'
+import is, { isa } from './is'
+import { ensureBlockArrowFunctionExpression } from './functions'
 
 import { getLoc, placeTreeAtLocWhenUnplaced as allAtLoc } from 'ast-loc-utils'
 
@@ -46,19 +47,29 @@ export replaceWithSafeCall(path, callExpr) ->
 export transformSafeMemberExpression(path) ->
   // x?.y -> x == null ? x : x.y
   // x?[y] -> x == null ? x : x[y]
-  { node } = path;
-  { object } = node;
+  { node } = path
+  { object } = node
 
+  // Demote to an ordinary member expression
+  node.type = "MemberExpression"
+
+  // Generate null check, hoisting ref if necessary.
   left = if object.type === "Identifier" or object.type === "SafeMemberExpression":
     object
   else:
+    // "Lazy man's" fix for the `?.` fat arrow bug.
+    // The ref hoist below would cause a bodiless ArrowExpression to gain a body with a declared ref at the top.
+    // This invalidates the `path`, which now points to a bodiless arrow function that no longer exists in the ast.
+    // Instead let's handle that case early.
+    if path.parent~isa("ArrowFunctionExpression"):
+      path.parentPath~ensureBlockArrowFunctionExpression()
+      now path = path.parentPath.get("body.body.0.expression")
+
     ref = path.scope.generateDeclaredUidIdentifier("ref")~atNode(object)
     node.object = ref
     t.assignmentExpression("=", ref, object)~atNode(object)
 
   nullCheck = t.binaryExpression("==", left, t.nullLiteral()~atNode(object))~atNode(object)
-  node.type = "MemberExpression"
-  path.replaceWith(node)
 
   // Gather trailing subscripts/calls, which are parent nodes:
   // eg; in `o?.x.y()`, group trailing `.x.y()` into the ternary

test/fixtures/safe-member-expression/fat-arrow-expression-body/actual.js

@@ -0,0 +1,5 @@
+=> _?.a
+
+=> _.a?.a
+
+=> f()?.x

test/fixtures/safe-member-expression/fat-arrow-expression-body/expected.js

@@ -0,0 +1,11 @@
+() => _ == null ? null : _.a;
+
+() => {
+  var _ref;
+
+  return (_ref = _.a) == null ? null : _ref.a;
+};() => {
+  var _ref2;
+
+  return (_ref2 = f()) == null ? null : _ref2.x;
+};