Avoid duplicate nested function body(Part2)

Summary:
Release Note: Avoid duplicate nested function body(Part2)
Deal with the situation that nested function duplicates with a residual function did not use factory function.
This makes the sample code in issue #777 working.

TODO: Deal with function declaration.
Closes https://github.com/facebook/prepack/pull/1086

Differential Revision: D6081649

Pulled By: yinghuitan

fbshipit-source-id: 20a5a7fa3f053f1aa083345df56a089be1c10d6a

src/serializer/ResidualFunctions.js

@@ -144,24 +144,29 @@ export class ResidualFunctions {
   _generateFactoryFunctionInfos(
     rewrittenAdditionalFunctions: Map<FunctionValue, Array<BabelNodeStatement>>
   ): Map<number, FactoryFunctionInfo> {
-    const factoryFunctionIds = new Map();
+    const factoryFunctionInfos = new Map();
     for (const [functionBody, instances] of this.functions) {
       invariant(instances.length > 0);
 
+      let factoryId;
+      const suffix = instances[0].functionValue.__originalName || "";
       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 });
+        factoryId = t.identifier(this.factoryNameGenerator.generate(suffix));
+      } else {
+        // For inline function body case, use the first function as the factory function.
+        factoryId = this.locationService.getLocation(instances[0].functionValue);
       }
+
+      const functionUniqueTag = ((functionBody: any): FunctionBodyAstNode).uniqueTag;
+      invariant(functionUniqueTag);
+
+      const functionInfo = this.residualFunctionInfos.get(functionBody);
+      invariant(functionInfo);
+      factoryFunctionInfos.set(functionUniqueTag, { factoryId, functionInfo });
     }
-    return factoryFunctionIds;
+    return factoryFunctionInfos;
   }
 
   spliceFunctions(

src/serializer/visitors.js

@@ -12,8 +12,8 @@
 import { Realm } from "../realm.js";
 import { FunctionValue } from "../values/index.js";
 import * as t from "babel-types";
-import type { BabelNodeExpression, BabelNodeCallExpression, BabelNodeFunctionExpression } from "babel-types";
 import { convertExpressionToJSXIdentifier } from "../utils/jsx";
+import type { BabelNodeExpression, BabelNodeCallExpression, BabelNodeFunctionExpression } from "babel-types";
 import type { BabelTraversePath } from "babel-traverse";
 import type { FunctionBodyAstNode } from "../types.js";
 import type { TryQuery, FunctionInfo, FactoryFunctionInfo, ResidualFunctionBinding } from "./types.js";
@@ -95,15 +95,6 @@ function canShareFunctionBody(duplicateFunctionInfo: FactoryFunctionInfo): boole
   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;
@@ -143,7 +134,8 @@ export let ClosureRefReplacer = {
     }
   },
 
-  // TODO: handle FunctionDeclaration
+  // TODO: handle FunctionDeclaration.
+  // Replace "function () {}" ==> "factory_id.bind(null)".
   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.
@@ -157,7 +149,11 @@ export let ClosureRefReplacer = {
       // Un-interpreted nested function.
       return;
     }
-    replaceNestedFunction(functionTag, path, state);
+    const duplicateFunctionInfo = state.factoryFunctionInfos.get(functionTag);
+    if (duplicateFunctionInfo && canShareFunctionBody(duplicateFunctionInfo)) {
+      const { factoryId } = duplicateFunctionInfo;
+      path.replaceWith(t.callExpression(t.memberExpression(factoryId, t.identifier("bind")), [nullExpression]));
+    }
   },
 
   // A few very simple dead code elimination helpers. Eventually these should be subsumed by the partial evaluators.

test/serializer/basic/NestedFunctions3.js

@@ -0,0 +1,12 @@
+// Copies of x0: 3
+f = function() {
+  return function() {
+    var x0 = 1;
+    var x1 = x0 + x0;
+    var x2 = x1 + x1;
+    var x3 = x2 + x2;
+    var x4 = x3 + x3;
+    return x4;
+  }
+}
+g = f();

test/serializer/basic/NestedFunctions4.js

@@ -0,0 +1,16 @@
+// TODO: add copies checking after handling FunctionDeclaration
+f = function() {
+  function nested() {
+    var x0 = 1;
+    var x1 = x0 + x0;
+    var x2 = x1 + x1;
+    var x3 = x2 + x2;
+    var x4 = x3 + x3;
+    return x4;
+  }
+  return nested;
+}
+g = f();
+inspect = function() {
+  return g();
+}

test/serializer/basic/NestedFunctions5.js

@@ -0,0 +1,15 @@
+f = function() {
+  return nested; // note that declaration comes later --- this is okay!
+  function nested() {
+    var x0 = 1;
+    var x1 = x0 + x0;
+    var x2 = x1 + x1;
+    var x3 = x2 + x2;
+    var x4 = x3 + x3;
+    return x4;
+  }
+}
+g = f();
+inspect = function() {  
+  return f()();
+}