Optionally remove module factory functions (#2605)

Summary:
Re-submission of #2602, with a bug fix for speculatively executed modules.

The bad import is fixed, and guarded by #2604, which correctly flagged this PR.

```
/home/circleci/project/src/utils/modules.js
  44:1  error  '../../lib/values/ECMAScriptSourceFunctionValue.js' import is restricted from being used by a pattern  no-restricted-imports

✖ 1 problem (1 error, 0 warnings)
```
Pull Request resolved: https://github.com/facebook/prepack/pull/2605

Differential Revision: D10446661

Pulled By: sb98052

fbshipit-source-id: fcd635e2c9637211f843fd7c3098673a112399bd

scripts/test-runner.js

@@ -375,11 +375,15 @@ function runTest(name, code, options: PrepackOptions, args) {
     serialize: true,
     uniqueSuffix: "",
     arrayNestedOptimizedFunctionsEnabled: false,
+    removeModuleFactoryFunctions: false,
     modulesToInitialize,
   }): any): PrepackOptions); // Since PrepackOptions is an exact type I have to cast
   if (code.includes("// arrayNestedOptimizedFunctionsEnabled")) {
     options.arrayNestedOptimizedFunctionsEnabled = true;
   }
+  if (code.includes("// removeModuleFactoryFunctions")) {
+    options.removeModuleFactoryFunctions = true;
+  }
   if (code.includes("// throws introspection error")) {
     try {
       let realmOptions = {

src/options.js

@@ -61,6 +61,7 @@ export type RealmOptions = {
   abstractValueImpliesMax?: number,
   arrayNestedOptimizedFunctionsEnabled?: boolean,
   reactFailOnUnsupportedSideEffects?: boolean,
+  removeModuleFactoryFunctions?: boolean,
 };
 
 export type SerializerOptions = {

src/prepack-cli.js

@@ -71,6 +71,7 @@ function run(
     --statsFile              The name of the output file where statistics will be written to.
     --heapGraphFilePath      The name of the output file where heap graph will be written to.
     --dumpIRFilePath         The name of the output file where the intermediate representation will be written to.
+    --removeModuleFactoryFunctions         Forces optimized module factory functions to be removed, even if they are reachable.
     --inlineExpressions      When generating code, tells prepack to avoid naming expressions when they are only used once,
                              and instead inline them where they are used.
     --invariantLevel         0: no invariants (default); 1: checks for abstract values; 2: checks for accessed built-ins; 3: internal consistency
@@ -123,6 +124,7 @@ function run(
     debugNames: false,
     emitConcreteModel: false,
     inlineExpressions: false,
+    removeModuleFactoryFunctions: false,
     logStatistics: false,
     logModules: false,
     delayInitializations: false,

src/prepack-options.js

@@ -45,6 +45,7 @@ export type PrepackOptions = {|
   serialize?: boolean,
   check?: Array<number>,
   inlineExpressions?: boolean,
+  removeModuleFactoryFunctions?: boolean,
   sourceMaps?: boolean,
   modulesToInitialize?: Set<string | number> | "ALL",
   statsFile?: string,
@@ -90,6 +91,7 @@ export function getRealmOptions({
   debugReproArgs,
   arrayNestedOptimizedFunctionsEnabled,
   reactFailOnUnsupportedSideEffects,
+  removeModuleFactoryFunctions,
 }: PrepackOptions): RealmOptions {
   return {
     compatibility,
@@ -116,6 +118,7 @@ export function getRealmOptions({
     debugReproArgs,
     arrayNestedOptimizedFunctionsEnabled,
     reactFailOnUnsupportedSideEffects,
+    removeModuleFactoryFunctions,
   };
 }
 

src/realm.js

@@ -292,6 +292,7 @@ export class Realm {
     }
 
     this.collectedNestedOptimizedFunctionEffects = new Map();
+    this.moduleFactoryFunctionsToRemove = new Map();
     this.tracers = [];
 
     // These get initialized in construct_realm to avoid the dependency
@@ -354,6 +355,7 @@ export class Realm {
     this.optimizedFunctions = new Map();
     this.arrayNestedOptimizedFunctionsEnabled =
       opts.arrayNestedOptimizedFunctionsEnabled || opts.instantRender || false;
+    this.removeModuleFactoryFunctions = opts.removeModuleFactoryFunctions || false;
   }
 
   statistics: RealmStatistics;
@@ -467,6 +469,8 @@ export class Realm {
   simplifyAndRefineAbstractCondition: AbstractValue => Value;
 
   collectedNestedOptimizedFunctionEffects: Map<ECMAScriptSourceFunctionValue, Effects>;
+  removeModuleFactoryFunctions: boolean;
+  moduleFactoryFunctionsToRemove: Map<number, string>;
   tracers: Array<Tracer>;
 
   MOBILE_JSC_VERSION = "jsc-600-1-4-17";

src/serializer/ResidualFunctionInstantiator.js

@@ -99,17 +99,20 @@ export class ResidualFunctionInstantiator<
   T: BabelNodeClassMethod | BabelNodeFunctionExpression | BabelNodeArrowFunctionExpression
 > {
   factoryFunctionInfos: Map<number, FactoryFunctionInfo>;
+  factoryFunctionsToRemove: Map<number, string>;
   identifierReplacements: Map<BabelNodeIdentifier, Replacement>;
   callReplacements: Map<BabelNodeCallExpression, Replacement>;
   root: T;
 
   constructor(
     factoryFunctionInfos: Map<number, FactoryFunctionInfo>,
+    factoryFunctionsToRemove: Map<number, string>,
     identifierReplacements: Map<BabelNodeIdentifier, Replacement>,
     callReplacements: Map<BabelNodeCallExpression, Replacement>,
     root: T
   ) {
     this.factoryFunctionInfos = factoryFunctionInfos;
+    this.factoryFunctionsToRemove = factoryFunctionsToRemove;
     this.identifierReplacements = identifierReplacements;
     this.callReplacements = callReplacements;
     this.root = root;
@@ -203,6 +206,16 @@ export class ResidualFunctionInstantiator<
           const { factoryId } = duplicateFunctionInfo;
           return t.callExpression(t.memberExpression(factoryId, t.identifier("bind")), [nullExpression]);
         }
+
+        if (this.factoryFunctionsToRemove.has(functionTag)) {
+          let newFunctionExpression = Object.assign({}, node);
+          newFunctionExpression.body = t.blockStatement([
+            t.throwStatement(
+              t.newExpression(t.identifier("Error"), [t.stringLiteral("Function was specialized out by Prepack")])
+            ),
+          ]);
+          return newFunctionExpression;
+        }
       }
     }
 

src/serializer/ResidualFunctions.js

@@ -500,6 +500,7 @@ export class ResidualFunctions {
               let methodParams = params.slice();
               let classMethod = new ResidualFunctionInstantiator(
                 factoryFunctionInfos,
+                this.realm.moduleFactoryFunctionsToRemove,
                 this._getIdentifierReplacements(funcBody, residualFunctionBindings),
                 this._getCallReplacements(funcBody),
                 t.classMethod(
@@ -531,6 +532,7 @@ export class ResidualFunctions {
             let isLexical = instance.functionValue.$ThisMode === "lexical";
             funcOrClassNode = new ResidualFunctionInstantiator(
               factoryFunctionInfos,
+              this.realm.moduleFactoryFunctionsToRemove,
               this._getIdentifierReplacements(funcBody, residualFunctionBindings),
               this._getCallReplacements(funcBody),
               this._createFunctionExpression(params, funcBody, isLexical)
@@ -623,6 +625,7 @@ export class ResidualFunctions {
         factoryParams = factoryParams.concat(params).slice();
         let factoryNode = new ResidualFunctionInstantiator(
           factoryFunctionInfos,
+          this.realm.moduleFactoryFunctionsToRemove,
           this._getIdentifierReplacements(funcBody, sameResidualBindings),
           this._getCallReplacements(funcBody),
           this._createFunctionExpression(factoryParams, funcBody, false)

src/serializer/serializer.js

@@ -176,6 +176,11 @@ export class Serializer {
         if (this.logger.hasErrors()) return undefined;
       }
 
+      let moduleFactoryFunctionsToRemove = this.realm.moduleFactoryFunctionsToRemove;
+      for (let [functionId, moduleIdOfFunction] of this.realm.moduleFactoryFunctionsToRemove) {
+        if (!this.modules.initializedModules.has(moduleIdOfFunction)) moduleFactoryFunctionsToRemove.delete(functionId);
+      }
+
       let heapGraph;
       let ast = (() => {
         // We wrap the following in an anonymous function declaration to ensure

src/utils/modules.js

@@ -13,10 +13,13 @@ import { GlobalEnvironmentRecord, DeclarativeEnvironmentRecord } from "../enviro
 import { FatalError } from "../errors.js";
 import { Realm, Tracer } from "../realm.js";
 import type { Effects } from "../realm.js";
+import type { FunctionBodyAstNode } from "../types.js";
 import { Get } from "../methods/index.js";
 import { Environment } from "../singletons.js";
 import {
   Value,
+  BoundFunctionValue,
+  ECMAScriptSourceFunctionValue,
   FunctionValue,
   ObjectValue,
   NumberValue,
@@ -153,6 +156,7 @@ export class ModuleTracer extends Tracer {
 
       let moduleId = argumentsList[0];
       let moduleIdValue;
+
       // Do some sanity checks and request require(...) calls with bad arguments
       if (moduleId instanceof NumberValue || moduleId instanceof StringValue) moduleIdValue = moduleId.value;
       else return performCall();
@@ -175,7 +179,6 @@ export class ModuleTracer extends Tracer {
     } else if (F === this.modules.getDefine()) {
       // Here, we handle calls of the form
       //   __d(factoryFunction, moduleId, dependencyArray)
-
       let moduleId = argumentsList[1];
       if (moduleId instanceof NumberValue || moduleId instanceof StringValue) {
         let moduleIdValue = moduleId.value;
@@ -191,6 +194,18 @@ export class ModuleTracer extends Tracer {
               argumentsList[2],
               "Third argument to define function is present but not a concrete array."
             );
+
+          // Remove if explicitly marked at optimization time
+          let realm = factoryFunction.$Realm;
+          if (realm.removeModuleFactoryFunctions) {
+            let targetFunction = factoryFunction;
+            if (factoryFunction instanceof BoundFunctionValue) targetFunction = factoryFunction.$BoundTargetFunction;
+            invariant(targetFunction instanceof ECMAScriptSourceFunctionValue);
+            let body = ((targetFunction.$ECMAScriptCode: any): FunctionBodyAstNode);
+            let uniqueOrderedTag = body.uniqueOrderedTag;
+            invariant(uniqueOrderedTag !== undefined);
+            realm.moduleFactoryFunctionsToRemove.set(uniqueOrderedTag, "" + moduleId.value);
+          }
         } else
           this.modules.logger.logError(factoryFunction, "First argument to define function is not a function value.");
       } else
@@ -246,6 +261,11 @@ export class Modules {
         }
       }
     }
+
+    let moduleFactoryFunctionsToRemove = this.realm.moduleFactoryFunctionsToRemove;
+    for (let [functionId, moduleIdOfFunction] of this.realm.moduleFactoryFunctionsToRemove) {
+      if (!this.initializedModules.has(moduleIdOfFunction)) moduleFactoryFunctionsToRemove.delete(functionId);
+    }
     this.getStatistics().initializedModules = this.initializedModules.size;
     this.getStatistics().totalModules = this.moduleIds.size;
   }
@@ -436,7 +456,12 @@ export class Modules {
     let count = 0;
     let body = (moduleId: string) => {
       if (this.initializedModules.has(moduleId)) return;
-      let effects = this.tryInitializeModule(moduleId, `Speculative initialization of module ${moduleId}`);
+      let moduleIdNumberIfNumeric = parseInt(moduleId, 10);
+      if (isNaN(moduleIdNumberIfNumeric)) moduleIdNumberIfNumeric = moduleId;
+      let effects = this.tryInitializeModule(
+        moduleIdNumberIfNumeric,
+        `Speculative initialization of module ${moduleId}`
+      );
       if (effects === undefined) return;
       let result = effects.result;
       if (!(result instanceof NormalCompletion)) return; // module might throw

test/serializer/optimizations/require_removefactoryfunctions.js

@@ -0,0 +1,132 @@
+// es6
+// removeModuleFactoryFunctions
+// does not contain:require(0)
+// does not contain:magic-string-1
+// does contain:magic-string-2
+var modules = Object.create(null);
+
+__d = define;
+function require(moduleId) {
+  var moduleIdReallyIsNumber = moduleId;
+  var module = modules[moduleIdReallyIsNumber];
+  return module && module.isInitialized ? module.exports : guardedLoadModule(moduleIdReallyIsNumber, module);
+}
+
+function define(factory, moduleId, dependencyMap) {
+  if (moduleId in modules) {
+    return;
+  }
+  modules[moduleId] = {
+    dependencyMap: dependencyMap,
+    exports: undefined,
+    factory: factory,
+    hasError: false,
+    isInitialized: false,
+  };
+
+  var _verboseName = arguments[3];
+  if (_verboseName) {
+    modules[moduleId].verboseName = _verboseName;
+    global.verboseNamesToModuleIds[_verboseName] = moduleId;
+  }
+}
+
+var inGuard = false;
+function guardedLoadModule(moduleId, module) {
+  if (!inGuard && global.ErrorUtils) {
+    inGuard = true;
+    var returnValue = void 0;
+    try {
+      returnValue = loadModuleImplementation(moduleId, module);
+    } catch (e) {
+      global.ErrorUtils.reportFatalError(e);
+    }
+    inGuard = false;
+    return returnValue;
+  } else {
+    return loadModuleImplementation(moduleId, module);
+  }
+}
+
+function loadModuleImplementation(moduleId, module) {
+  var nativeRequire = global.nativeRequire;
+  if (!module && nativeRequire) {
+    nativeRequire(moduleId);
+    module = modules[moduleId];
+  }
+
+  if (!module) {
+    throw unknownModuleError(moduleId);
+  }
+
+  if (module.hasError) {
+    throw moduleThrewError(moduleId);
+  }
+
+  module.isInitialized = true;
+  var exports = (module.exports = {});
+  var _module = module,
+    factory = _module.factory,
+    dependencyMap = _module.dependencyMap;
+  try {
+    var _moduleObject = { exports: exports };
+
+    factory(global, require, _moduleObject, exports, dependencyMap);
+
+    module.factory = undefined;
+
+    return (module.exports = _moduleObject.exports);
+  } catch (e) {
+    module.hasError = true;
+    module.isInitialized = false;
+    module.exports = undefined;
+    throw e;
+  }
+}
+
+function unknownModuleError(id) {
+  var message = 'Requiring unknown module "' + id + '".';
+  return Error(message);
+}
+
+function moduleThrewError(id) {
+  return Error('Requiring module "' + id + '", which threw an exception.');
+}
+
+// === End require code ===
+
+function defineModules() {
+  define(function(global, require, module, exports) {
+    let z = "magic-string-1";
+    module.exports = { foo: " hello " };
+  }, 0, null);
+
+  define(function(global, require, module, exports) {
+    var x = require(0);
+    var y = require(2);
+    let z = "magic-string-2";
+    module.exports = {
+      bar: " goodbye",
+      foo2: x.foo,
+      baz: y.baz,
+    };
+  }, 1, null);
+
+  define(function(global, require, module, exports) {
+    module.exports = { baz: " foo " };
+  }, 2, null);
+}
+
+defineModules();
+
+var x = require(0);
+
+function f() {
+  return x.foo === " hello " && modules[1].exports === undefined && require(1).bar === " goodbye";
+}
+
+inspect = function() {
+  // the require( 0) should be entirely eliminated from 1's factory function
+  // but the require(2) will remain
+  return f();
+};

test/serializer/optimizations/require_speculatively_specific.js

@@ -1,5 +1,6 @@
 // es6
-// does not contain:1 + 2
+// does not contain:27 + 15
+// does contain:42
 // does contain:3 + 4
 // initialize more modules:0
 
@@ -95,8 +96,7 @@ function moduleThrewError(id) {
 // === End require code ===
 
 define(function(global, require, module, exports) {
-  let useless = 1 + 2;
-  module.exports = { foo: " hello " };
+  module.exports = { foo: 27 + 15 };
 }, 0, null);
 
 define(function(global, r, module, exports) {