Support basic arrow function serialization (#1650)

Summary:
Release notes: adds basic ES2015 arrow function serialization

This allows Prepack to serialize back to arrow functions and also traverse them correctly when understanding `this`.
Closes https://github.com/facebook/prepack/pull/1650

Differential Revision: D8623368

Pulled By: trueadm

fbshipit-source-id: 9be7849770148b67e5f17eee88a633aac906b6e2

src/serializer/ResidualFunctionInstantiator.js

@@ -18,6 +18,7 @@ import type {
   BabelNodeFunctionExpression,
   BabelNodeExpression,
   BabelNodeClassMethod,
+  BabelNodeArrowFunctionExpression,
   BabelNodeWhileStatement,
   BabelNodeJSXIdentifier,
   BabelNodeJSXMemberExpression,
@@ -94,7 +95,9 @@ export function isPure(node: BabelNodeExpression | BabelNodeSpreadElement): bool
 // i.e. bindings to captured scopes that need to get renamed to variable ids.
 // The original nodes are never mutated; instead, nodes are cloned as needed.
 // Along the way, some trivial code optimizations are performed as well.
-export class ResidualFunctionInstantiator<T: BabelNodeClassMethod | BabelNodeFunctionExpression> {
+export class ResidualFunctionInstantiator<
+  T: BabelNodeClassMethod | BabelNodeFunctionExpression | BabelNodeArrowFunctionExpression
+> {
   factoryFunctionInfos: Map<number, FactoryFunctionInfo>;
   identifierReplacements: Map<BabelNodeIdentifier, Replacement>;
   callReplacements: Map<BabelNodeCallExpression, Replacement>;

src/serializer/ResidualFunctions.js

@@ -23,6 +23,7 @@ import type {
   BabelNodeSpreadElement,
   BabelNodeFunctionExpression,
   BabelNodeClassExpression,
+  BabelNodeArrowFunctionExpression,
 } from "babel-types";
 import type { FunctionBodyAstNode } from "../types.js";
 import type { NameGenerator } from "../utils/generator.js";
@@ -45,8 +46,12 @@ import { Referentializer } from "./Referentializer.js";
 import { getOrDefault } from "./utils.js";
 
 type ResidualFunctionsResult = {
-  unstrictFunctionBodies: Array<BabelNodeFunctionExpression | BabelNodeClassExpression>,
-  strictFunctionBodies: Array<BabelNodeFunctionExpression | BabelNodeClassExpression>,
+  unstrictFunctionBodies: Array<
+    BabelNodeFunctionExpression | BabelNodeClassExpression | BabelNodeArrowFunctionExpression
+  >,
+  strictFunctionBodies: Array<
+    BabelNodeFunctionExpression | BabelNodeClassExpression | BabelNodeArrowFunctionExpression
+  >,
 };
 
 export class ResidualFunctions {
@@ -256,10 +261,10 @@ export class ResidualFunctions {
     });
   }
 
-  _createFunctionExpression(params: Array<BabelNodeLVal>, body: BabelNodeBlockStatement) {
+  _createFunctionExpression(params: Array<BabelNodeLVal>, body: BabelNodeBlockStatement, isLexical: boolean) {
     // Additional statements might be inserted at the beginning of the body, so we clone it.
     body = ((Object.assign({}, body): any): BabelNodeBlockStatement);
-    return t.functionExpression(null, params, body);
+    return isLexical ? t.arrowFunctionExpression(params, body) : t.functionExpression(null, params, body);
   }
 
   spliceFunctions(
@@ -309,7 +314,11 @@ export class ResidualFunctions {
         funcNodes.set(functionValue, ((funcOrClassNode: any): BabelNodeFunctionExpression));
         body = getPrelude(instance);
       } else {
-        invariant(t.isCallExpression(funcOrClassNode) || t.isClassExpression(funcOrClassNode)); // .bind call
+        invariant(
+          t.isCallExpression(funcOrClassNode) ||
+            t.isClassExpression(funcOrClassNode) ||
+            t.isArrowFunctionExpression(funcOrClassNode)
+        ); // .bind call
         body = getFunctionBody(instance);
       }
       body.push(t.variableDeclaration("var", [t.variableDeclarator(funcId, funcOrClassNode)]));
@@ -353,6 +362,7 @@ export class ResidualFunctions {
       let { instance } = additionalFunctionInfo;
       let functionValue = ((funcValue: any): ECMAScriptSourceFunctionValue);
       let params = functionValue.$FormalParameters;
+      let isLexical = functionValue.$ThisMode === "lexical";
       invariant(params !== undefined);
 
       let rewrittenBody = rewrittenAdditionalFunctions.get(funcValue);
@@ -403,7 +413,9 @@ export class ResidualFunctions {
           funcOrClassNode.superClass = classSuperNode;
         }
       } else {
-        funcOrClassNode = t.functionExpression(null, params, functionBody);
+        funcOrClassNode = isLexical
+          ? t.arrowFunctionExpression(params, functionBody)
+          : t.functionExpression(null, params, functionBody);
       }
       let id = this.locationService.getLocation(funcValue);
       invariant(id !== undefined);
@@ -494,11 +506,12 @@ export class ResidualFunctions {
               funcOrClassNode.superClass = classSuperNode;
             }
           } else {
+            let isLexical = instance.functionValue.$ThisMode === "lexical";
             funcOrClassNode = new ResidualFunctionInstantiator(
               factoryFunctionInfos,
               this._getIdentifierReplacements(funcBody, residualFunctionBindings),
               this._getCallReplacements(funcBody),
-              this._createFunctionExpression(params, funcBody)
+              this._createFunctionExpression(params, funcBody, isLexical)
             ).instantiate();
 
             let scopeInitialization = [];
@@ -508,8 +521,11 @@ export class ResidualFunctions {
               );
             }
 
-            if (scopeInitialization.length > 0)
-              funcOrClassNode.body.body = scopeInitialization.concat(funcOrClassNode.body.body);
+            if (scopeInitialization.length > 0) {
+              let funcOrClassNodeBody = ((funcOrClassNode.body: any): BabelNodeBlockStatement);
+              invariant(t.isBlockStatement(funcOrClassNodeBody));
+              funcOrClassNodeBody.body = scopeInitialization.concat(funcOrClassNodeBody.body);
+            }
           }
           let id = this.locationService.getLocation(functionValue);
           invariant(id !== undefined);
@@ -591,10 +607,14 @@ export class ResidualFunctions {
           factoryFunctionInfos,
           this._getIdentifierReplacements(funcBody, sameResidualBindings),
           this._getCallReplacements(funcBody),
-          this._createFunctionExpression(factoryParams, funcBody)
+          this._createFunctionExpression(factoryParams, funcBody, false)
         ).instantiate();
 
-        if (scopeInitialization.length > 0) factoryNode.body.body = scopeInitialization.concat(factoryNode.body.body);
+        if (scopeInitialization.length > 0) {
+          let factoryNodeBody = ((factoryNode.body: any): BabelNodeBlockStatement);
+          invariant(t.isBlockStatement(factoryNodeBody));
+          factoryNodeBody.body = scopeInitialization.concat(factoryNodeBody.body);
+        }
 
         // factory functions do not depend on any nested generator scope, so they go to the prelude
         let factoryDeclaration = t.variableDeclaration("var", [t.variableDeclarator(factoryId, factoryNode)]);

src/serializer/ResidualHeapVisitor.js

@@ -546,6 +546,7 @@ export class ResidualHeapVisitor {
     if (!functionInfo) {
       functionInfo = {
         depth: 0,
+        lexicalDepth: 0,
         unbound: new Map(),
         requireCalls: new Map(),
         modified: new Set(),

src/serializer/types.js

@@ -84,6 +84,7 @@ export type FunctionInstance = {
 
 export type FunctionInfo = {
   depth: number,
+  lexicalDepth: number,
   unbound: Map<string, Array<BabelNodeIdentifier>>,
   requireCalls: Map<BabelNode, number | string>,
   modified: Set<string>,

src/serializer/visitors.js

@@ -42,7 +42,7 @@ function ignorePath(path: BabelTraversePath) {
 }
 
 export let ClosureRefVisitor = {
-  "FunctionDeclaration|FunctionExpression": {
+  "FunctionDeclaration|ArrowFunctionExpression|FunctionExpression": {
     enter(path: BabelTraversePath, state: ClosureRefVisitorState) {
       state.functionInfo.depth++;
     },
@@ -51,6 +51,17 @@ export let ClosureRefVisitor = {
     },
   },
 
+  ArrowFunctionExpression: {
+    enter(path: BabelTraversePath, state: ClosureRefVisitorState) {
+      state.functionInfo.depth++;
+      state.functionInfo.lexicalDepth++;
+    },
+    exit(path: BabelTraversePath, state: ClosureRefVisitorState) {
+      state.functionInfo.depth--;
+      state.functionInfo.lexicalDepth--;
+    },
+  },
+
   CallExpression(path: BabelTraversePath, state: ClosureRefVisitorState) {
     // Here we apply the require optimization by replacing require calls with their
     // corresponding initialized modules.
@@ -75,7 +86,7 @@ export let ClosureRefVisitor = {
   },
 
   ThisExpression(path: BabelTraversePath, state: ClosureRefVisitorState) {
-    if (state.functionInfo.depth === 1) {
+    if (state.functionInfo.depth - state.functionInfo.lexicalDepth === 1) {
       state.functionInfo.usesThis = true;
     }
   },

test/serializer/optimized-functions/ArrowFunction.js

@@ -0,0 +1,9 @@
+// does contain:=> {
+
+const foo = x => {
+  return x + 1;
+};
+
+if (global.__optimize) __optimize(foo);
+
+inspect = function() { return foo(5); }

test/serializer/optimized-functions/ArrowFunction2.js

@@ -0,0 +1,14 @@
+// does contain:=> {
+
+function foo(x) {
+  this.x = x;
+  const bar = () => {
+    return this.x;
+  };
+  if (global.__optimize) __optimize(bar);
+  this.bar = bar;
+}
+
+if (global.__optimize) __optimize(foo);
+
+inspect = function() { return foo(5).bar(); }