metro: fix collect-dependencies for the async/BundleSegments case

Reviewed By: cpojer

Differential Revision: D6462925

fbshipit-source-id: ca83176a7c7f47346daf16703cefb86172dc35b5

packages/metro/src/ModuleGraph/worker/__tests__/collect-dependencies-test.js

@@ -137,14 +137,17 @@ describe('Dependency collection from optimized ASTs', () => {
     ast = astFromCode(`
       const a = require(${dependencyMapName}[0], 'b/lib/a');
       exports.do = () => require(${dependencyMapName}[1], "do");
+      require(${dependencyMapName}[2], "BundleSegments").loadForModule(${dependencyMapName}[3]).then(function () { return require(${dependencyMapName}[3], "some/async/module"); }).then(foo => {});
       if (!something) {
-        require(${dependencyMapName}[2], "setup/something");
+        require(${dependencyMapName}[4], "setup/something");
       }
     `);
     deps = [
       {name: 'b/lib/a', isAsync: false},
       {name: 'do', isAsync: false},
-      {name: 'setup/something', isAsync: true},
+      {name: 'BundleSegments', isAsync: false},
+      {name: 'some/async/module', isAsync: true},
+      {name: 'setup/something', isAsync: false},
     ];
   });
 
@@ -164,12 +167,20 @@ describe('Dependency collection from optimized ASTs', () => {
     expect(result.dependencies).toEqual([{name: 'do', isAsync: false}]);
   });
 
-  it('only returns async dependencies that are in the code', () => {
+  it('only returns dependencies that are in the code, and properly translate async dependencies', () => {
-    ast = astFromCode(`require(${dependencyMapName}[2], "setup/something")`);
+    ast = astFromCode(`
+      require(${dependencyMapName}[2], "BundleSegments").loadForModule(${dependencyMapName}[3]).then(function () { return require(${dependencyMapName}[3], "some/async/module"); }).then(foo => {});
+    `);
     const result = forOptimization(ast, deps, dependencyMapName);
     expect(result.dependencies).toEqual([
-      {name: 'setup/something', isAsync: true},
+      {name: 'BundleSegments', isAsync: false},
+      {name: 'some/async/module', isAsync: true},
     ]);
+    expect(codeFromAst(ast)).toEqual(
+      comparableCode(`
+        require(${dependencyMapName}[0]).loadForModule(${dependencyMapName}[1]).then(function () { return require(${dependencyMapName}[1]); }).then(foo => {});
+      `),
+    );
   });
 
   it('replaces all call signatures inserted by a prior call to `collectDependencies`', () => {
@@ -178,8 +189,9 @@ describe('Dependency collection from optimized ASTs', () => {
       comparableCode(`
       const a = require(${dependencyMapName}[0]);
       exports.do = () => require(${dependencyMapName}[1]);
+      require(${dependencyMapName}[2]).loadForModule(${dependencyMapName}[3]).then(function () { return require(${dependencyMapName}[3]); }).then(foo => {});
       if (!something) {
-        require(${dependencyMapName}[2]);
+        require(${dependencyMapName}[4]);
       }
     `),
     );

packages/metro/src/ModuleGraph/worker/collect-dependencies.js

@@ -62,6 +62,8 @@ class Replacement {
     const mapLookup = createMapLookup(dependencyMapIdentifier, newId);
     return [mapLookup, oldId];
   }
+
+  processMemberExp(path, state, depMapIdent) {}
 }
 
 function getInvalidProdRequireMessage(node) {
@@ -101,6 +103,34 @@ class ProdReplacement {
   }
 
   getIndex(memberExpression, _: boolean) {
+    return memberExpression.property.value;
+  }
+
+  getDependencies(): $ReadOnlyArray<TransformResultDependency> {
+    return this.replacement.getDependencies();
+  }
+
+  makeArgs(newId, _, dependencyMapIdentifier) {
+    const mapLookup = createMapLookup(dependencyMapIdentifier, newId);
+    return [mapLookup];
+  }
+
+  processMemberExp(path, state, depMapIdent) {
+    const node = path.node;
+    if (
+      node.computed &&
+      node.object.type === 'Identifier' &&
+      node.object.name === depMapIdent &&
+      node.property.type === 'NumericLiteral'
+    ) {
+      const newIx = this._getTranslatedIndex(node);
+      if (newIx !== node.property.value) {
+        node.property.value = newIx;
+      }
+    }
+  }
+
+  _getTranslatedIndex(memberExpression) {
     const id = memberExpression.property.value;
     if (id in this.dependencies) {
       const dependency = this.dependencies[id];
@@ -114,15 +144,6 @@ class ProdReplacement {
         .join(', ')}`,
     );
   }
-
-  getDependencies(): $ReadOnlyArray<TransformResultDependency> {
-    return this.replacement.getDependencies();
-  }
-
-  makeArgs(newId, _, dependencyMapIdentifier) {
-    const mapLookup = createMapLookup(dependencyMapIdentifier, newId);
-    return [mapLookup];
-  }
 }
 
 function createMapLookup(dependencyMapIdentifier, propertyIdentifier) {
@@ -168,6 +189,12 @@ function collectDependencies(
           return;
         }
         const arg = replacement.getRequireCallArg(node);
+        // FIXME: in prod mode, this is a no-op, the actual translation of the
+        // index if done by the MemberExpression instead, so that it also
+        // translates calls to ex. BundleSegments.loadForModule() (but it
+        // could be anything that refers to that particular ID.)
+        // We should just get rid of the generic logic and have a specific
+        // prod/optimize babel visitor instead.
         const index = replacement.getIndex(arg, false);
         node.arguments = replacement.makeArgs(
           types.numericLiteral(index),
@@ -176,6 +203,17 @@ function collectDependencies(
         );
         visited.add(node);
       },
+
+      // FIXME: we should remove ID-translation completely from CallExpression
+      // handler for ProdReplacement, and just do everything here once.
+      // Refactor necessary.
+      MemberExpression(path, state) {
+        replacement.processMemberExp(
+          path,
+          state,
+          state.dependencyMapIdentifier.name,
+        );
+      },
     },
     null,
     traversalState,