Avoid duplicate nested function body(Part1)

Summary:
Release Note: Avoid duplicate nested function body(Part1)
Avoids duplicating the nested function body with other residual functions. Currently the implementation has the following limitations:
1. Did not deal with the situation that nested function duplicates with a residual function did not use factory function.
2. Did not deal with function declaration

I will address these limitations in later PR.
Closes https://github.com/facebook/prepack/pull/1081

Differential Revision: D6078328

Pulled By: yinghuitan

fbshipit-source-id: 3a6b6dd0b99a5cdb8cffeeba85d257682c3db40e

flow-typed/npm/babel-traverse_vx.x.x.js

@@ -24,6 +24,7 @@ declare module 'babel-traverse' {
     scope: BabelTraverseScope;
     node: BabelNode;
     parent: BabelNode;
+    parentPath: BabelTraversePath;
   }
 
   declare export class BabelTraverseScope {

src/methods/function.js

@@ -10,7 +10,7 @@
 /* @flow */
 
 import type { LexicalEnvironment } from "../environment.js";
-import type { PropertyKeyValue } from "../types.js";
+import type { PropertyKeyValue, FunctionBodyAstNode } from "../types.js";
 import type { Realm } from "../realm.js";
 import type { ECMAScriptFunctionValue } from "../values/index.js";
 import {
@@ -584,6 +584,7 @@ export function FunctionInitialize(
   F.$FormalParameters = ParameterList;
 
   // 7. Set the [[ECMAScriptCode]] internal slot of F to Body.
+  ((Body: any): FunctionBodyAstNode).uniqueTag = realm.functionBodyUniqueTagSeed++;
   F.$ECMAScriptCode = Body;
 
   // 8. Set the [[ScriptOrModule]] internal slot of F to GetActiveScriptOrModule().

src/realm.js

@@ -254,6 +254,10 @@ export class Realm {
   errorHandler: ?ErrorHandler;
   objectCount = 0;
   symbolCount = 867501803871088;
+  // Unique tag for identifying function body ast node. It is neeeded
+  // instead of ast node itself because we may perform ast tree deep clone
+  // during serialization which changes the ast identity.
+  functionBodyUniqueTagSeed = 1;
 
   globalSymbolRegistry: Array<{ $Key: string, $Symbol: SymbolValue }>;
 

src/serializer/ResidualFunctions.js

@@ -22,10 +22,11 @@ import type {
   BabelNodeSpreadElement,
   BabelNodeFunctionExpression,
 } from "babel-types";
-import { NameGenerator } from "../utils/generator.js";
+import type { FunctionBodyAstNode } from "../types.js";
+import type { NameGenerator } from "../utils/generator.js";
 import traverse from "babel-traverse";
 import invariant from "../invariant.js";
-import type { FunctionInfo, FunctionInstance, AdditionalFunctionInfo } from "./types.js";
+import type { FunctionInfo, FactoryFunctionInfo, FunctionInstance, AdditionalFunctionInfo } from "./types.js";
 import { BodyReference, AreSameResidualBinding, SerializerStatistics } from "./types.js";
 import { ClosureRefReplacer } from "./visitors.js";
 import { Modules } from "./modules.js";
@@ -117,6 +118,52 @@ export class ResidualFunctions {
     if (!this.firstFunctionUsages.has(val)) this.firstFunctionUsages.set(val, bodyReference);
   }
 
+  _shouldUseFactoryFunction(funcBody: BabelNodeBlockStatement, instances: Array<FunctionInstance>) {
+    function shouldInlineFunction(): boolean {
+      let shouldInline = true;
+      if (funcBody.start && funcBody.end) {
+        let bodySize = funcBody.end - funcBody.start;
+        shouldInline = bodySize <= 30;
+      }
+      return shouldInline;
+    }
+    let functionInfo = this.residualFunctionInfos.get(funcBody);
+    invariant(functionInfo);
+    let { usesArguments } = functionInfo;
+    return !shouldInlineFunction() && instances.length > 1 && !usesArguments;
+  }
+
+  // Note: this function takes linear time. Please do not call it inside loop.
+  _hasRewrittenFunctionInstance(
+    rewrittenAdditionalFunctions: Map<FunctionValue, Array<BabelNodeStatement>>,
+    instances: Array<FunctionInstance>
+  ): boolean {
+    return instances.find(instance => rewrittenAdditionalFunctions.has(instance.functionValue)) !== undefined;
+  }
+
+  _generateFactoryFunctionInfos(
+    rewrittenAdditionalFunctions: Map<FunctionValue, Array<BabelNodeStatement>>
+  ): Map<number, FactoryFunctionInfo> {
+    const factoryFunctionIds = new Map();
+    for (const [functionBody, instances] of this.functions) {
+      invariant(instances.length > 0);
+
+      if (this._shouldUseFactoryFunction(functionBody, instances)) {
+        // Rewritten function should never use factory function.
+        invariant(!this._hasRewrittenFunctionInstance(rewrittenAdditionalFunctions, instances));
+
+        const functionUniqueTag = ((functionBody: any): FunctionBodyAstNode).uniqueTag;
+        invariant(functionUniqueTag);
+        const suffix = instances[0].functionValue.__originalName || "";
+        const factoryId = t.identifier(this.factoryNameGenerator.generate(suffix));
+        const functionInfo = this.residualFunctionInfos.get(functionBody);
+        invariant(functionInfo);
+        factoryFunctionIds.set(functionUniqueTag, { factoryId, functionInfo });
+      }
+    }
+    return factoryFunctionIds;
+  }
+
   spliceFunctions(
     rewrittenAdditionalFunctions: Map<FunctionValue, Array<BabelNodeStatement>>
   ): ResidualFunctionsResult {
@@ -252,19 +299,14 @@ export class ResidualFunctions {
     }
 
     // Process normal functions
+    const factoryFunctionInfos = this._generateFactoryFunctionInfos(rewrittenAdditionalFunctions);
     for (let [funcBody, instances] of functionEntries) {
       let functionInfo = this.residualFunctionInfos.get(funcBody);
       invariant(functionInfo);
-      let { unbound, modified, usesThis, usesArguments } = functionInfo;
+      let { unbound, modified, usesThis } = functionInfo;
       let params = instances[0].functionValue.$FormalParameters;
       invariant(params !== undefined);
 
-      let shouldInline = !funcBody;
-      if (!shouldInline && funcBody.start && funcBody.end) {
-        let bodySize = funcBody.end - funcBody.start;
-        shouldInline = bodySize <= 30;
-      }
-
       // Split instances into normal or nested in an additional function
       let normalInstances = [];
       let additionalFunctionNestedInstances = [];
@@ -304,6 +346,7 @@ export class ResidualFunctions {
             requireReturns: this.requireReturns,
             requireStatistics,
             isRequire: this.modules.getIsRequire(funcParams, [functionValue]),
+            factoryFunctionInfos,
           });
 
           if (functionValue.$Strict) {
@@ -317,11 +360,14 @@ export class ResidualFunctions {
       };
 
       if (additionalFunctionNestedInstances.length > 0) naiveProcessInstances(additionalFunctionNestedInstances);
-      if (shouldInline || normalInstances.length === 1 || usesArguments) {
+      if (!this._shouldUseFactoryFunction(funcBody, normalInstances)) {
         naiveProcessInstances(normalInstances);
       } else if (normalInstances.length > 0) {
-        let suffix = normalInstances[0].functionValue.__originalName || "";
-        let factoryId = t.identifier(this.factoryNameGenerator.generate(suffix));
+        const functionUniqueTag = ((funcBody: any): FunctionBodyAstNode).uniqueTag;
+        invariant(functionUniqueTag);
+        const factoryInfo = factoryFunctionInfos.get(functionUniqueTag);
+        invariant(factoryInfo);
+        const { factoryId } = factoryInfo;
 
         // filter included variables to only include those that are different
         let factoryNames: Array<string> = [];
@@ -400,6 +446,7 @@ export class ResidualFunctions {
           requireReturns: this.requireReturns,
           requireStatistics,
           isRequire: this.modules.getIsRequire(factoryParams, normalInstances.map(instance => instance.functionValue)),
+          factoryFunctionInfos,
         });
 
         for (let instance of normalInstances) {

src/serializer/types.js

@@ -43,6 +43,8 @@ export type FunctionInfo = {
   usesThis: boolean,
 };
 
+export type FactoryFunctionInfo = { factoryId: BabelNodeIdentifier, functionInfo: FunctionInfo };
+
 export type ResidualFunctionBinding = {
   value: void | Value,
   modified: boolean,

src/serializer/visitors.js

@@ -12,10 +12,12 @@
 import { Realm } from "../realm.js";
 import { FunctionValue } from "../values/index.js";
 import * as t from "babel-types";
-import type { BabelNodeExpression, BabelNodeCallExpression } from "babel-types";
-import { BabelTraversePath } from "babel-traverse";
+import type { BabelNodeExpression, BabelNodeCallExpression, BabelNodeFunctionExpression } from "babel-types";
 import { convertExpressionToJSXIdentifier } from "../utils/jsx";
-import type { TryQuery, FunctionInfo, ResidualFunctionBinding } from "./types.js";
+import type { BabelTraversePath } from "babel-traverse";
+import type { FunctionBodyAstNode } from "../types.js";
+import type { TryQuery, FunctionInfo, FactoryFunctionInfo, ResidualFunctionBinding } from "./types.js";
+import { nullExpression } from "../utils/internalizer.js";
 
 export type ClosureRefVisitorState = {
   tryQuery: TryQuery<*>,
@@ -30,6 +32,7 @@ export type ClosureRefReplacerState = {
   requireReturns: Map<number | string, BabelNodeExpression>,
   requireStatistics: { replaced: 0, count: 0 },
   isRequire: void | ((scope: any, node: BabelNodeCallExpression) => boolean),
+  factoryFunctionInfos: Map<number, FactoryFunctionInfo>,
 };
 
 function markVisited(node, data) {
@@ -84,6 +87,23 @@ function getLiteralTruthiness(node): { known: boolean, value?: boolean } {
   return { known: false };
 }
 
+function canShareFunctionBody(duplicateFunctionInfo: FactoryFunctionInfo): boolean {
+  // Only share function when:
+  // 1. it does not access any free variables.
+  // 2. it does not use "this".
+  const { unbound, modified, usesThis } = duplicateFunctionInfo.functionInfo;
+  return unbound.size === 0 && modified.size === 0 && !usesThis;
+}
+
+// TODO: enhance for nested functions accessing read-only free variables.
+function replaceNestedFunction(functionTag: number, path: BabelTraversePath, state: ClosureRefReplacerState) {
+  const duplicateFunctionInfo = state.factoryFunctionInfos.get(functionTag);
+  if (duplicateFunctionInfo && canShareFunctionBody(duplicateFunctionInfo)) {
+    const { factoryId } = duplicateFunctionInfo;
+    path.replaceWith(t.callExpression(t.memberExpression(factoryId, t.identifier("bind")), [nullExpression]));
+  }
+}
+
 export let ClosureRefReplacer = {
   ReferencedIdentifier(path: BabelTraversePath, state: ClosureRefReplacerState) {
     if (ignorePath(path)) return;
@@ -123,6 +143,23 @@ export let ClosureRefReplacer = {
     }
   },
 
+  // TODO: handle FunctionDeclaration
+  FunctionExpression(path: BabelTraversePath, state: ClosureRefReplacerState) {
+    if (t.isProgram(path.parentPath.parentPath.node)) {
+      // Our goal is replacing duplicate nested function so skip root residual function itself.
+      // This assumes the root function is wrapped with: t.file(t.program([t.expressionStatement(rootFunction).
+      return;
+    }
+
+    const functionExpression: BabelNodeFunctionExpression = path.node;
+    const functionTag = ((functionExpression.body: any): FunctionBodyAstNode).uniqueTag;
+    if (!functionTag) {
+      // Un-interpreted nested function.
+      return;
+    }
+    replaceNestedFunction(functionTag, path, state);
+  },
+
   // A few very simple dead code elimination helpers. Eventually these should be subsumed by the partial evaluators.
   IfStatement: {
     exit: function(path: BabelTraversePath, state: ClosureRefReplacerState) {

src/types.js

@@ -94,6 +94,11 @@ export type Descriptor = {
   set?: UndefinedValue | CallableObjectValue | AbstractValue,
 };
 
+export type FunctionBodyAstNode = {
+  // Function body ast node will have uniqueTag after interpreted.
+  uniqueTag?: number,
+};
+
 export type PropertyBinding = {
   descriptor?: Descriptor,
   object: ObjectValue | AbstractObjectValue,

test/serializer/basic/NestedFunctions1.js

@@ -0,0 +1,15 @@
+// Copies of x: 2
+(function() {
+  global.f = function() {
+    return function() {
+      /* This makes the function too big to inline. */
+      var x = 10;
+      return 2 * x;
+    }
+  }
+  global.g1 = f();
+  global.g2 = f();
+  inspect = function() {
+    return g1() + g2();
+  }
+})();

test/serializer/basic/NestedFunctions2.js

@@ -0,0 +1,19 @@
+// Copies of x: 2
+(function() {
+  global.f = function() {
+    return function() {
+      return function() {
+        /* This makes the function too big to inline. */
+        var x = 10;
+        return 2 * x;
+      }
+    }
+  }
+  global.g1 = f();
+  global.g2 = f();
+  global.h1 = g1();
+  global.h2 = g1();
+  inspect = function() {
+    return g1()() + g2()() + h1() + h2();
+  }
+})();