Add support for RAM bundle groups

Summary: Adds support for “RAM bundle groups” (common section for module groups) to the new Buck integration

Reviewed By: jeanlauliac

Differential Revision: D5112065

fbshipit-source-id: 038c06b8f4133c7fcd39aba8bb04a5ef42594f3e

packages/metro-bundler/src/Bundler/Bundle.js

@@ -19,6 +19,7 @@ const crypto = require('crypto');
 const debug = require('debug')('RNP:Bundle');
 const invariant = require('fbjs/lib/invariant');
 
+const {createRamBundleGroups} = require('./util');
 const {fromRawMappings} = require('./source-map');
 const {isMappingsMap} = require('../lib/SourceMap');
 
@@ -185,7 +186,7 @@ class Bundle extends BundleBase {
         lazyModules,
         get groups() {
           if (!groups) {
-            groups = createGroups(ramGroups || [], lazyModules);
+            groups = createRamBundleGroups(ramGroups || [], lazyModules, subtree);
           }
           return groups;
         },
@@ -344,18 +345,18 @@ function partition(array, predicate) {
   return [included, excluded];
 }
 
-function * filter(iterator, predicate) {
-  for (const value of iterator) {
-    if (predicate(value)) {
-      yield value;
-    }
-  }
-}
-
-function * subtree(moduleTransport: ModuleTransport, moduleTransportsByPath, seen = new Set()) {
+function * subtree(
+  moduleTransport: ModuleTransport,
+  moduleTransportsByPath: Map<string, ModuleTransport>,
+  seen = new Set(),
+) {
   seen.add(moduleTransport.id);
-  /* $FlowFixMe: there may not be a `meta` object */
-  for (const [, {path}] of moduleTransport.meta.dependencyPairs || []) {
+  const {meta} = moduleTransport;
+  invariant(
+    meta != null,
+    'Unexpected module transport without meta information: ' + moduleTransport.sourcePath,
+  );
+  for (const [, {path}] of meta.dependencyPairs || []) {
     const dependency = moduleTransportsByPath.get(path);
     if (dependency && !seen.has(dependency.id)) {
       yield dependency.id;
@@ -364,75 +365,6 @@ function * subtree(moduleTransport: ModuleTransport, moduleTransportsByPath, see
   }
 }
 
-class ArrayMap extends Map {
-  get(key) {
-    let array = super.get(key);
-    if (!array) {
-      array = [];
-      this.set(key, array);
-    }
-    return array;
-  }
-}
-
-function createGroups(ramGroups: Array<string>, lazyModules) {
-  // build two maps that allow to lookup module data
-  // by path or (numeric) module id;
-  const byPath = new Map();
-  const byId = new Map();
-  lazyModules.forEach(m => {
-    byPath.set(m.sourcePath, m);
-    byId.set(m.id, m.sourcePath);
-  });
-
-  // build a map of group root IDs to an array of module IDs in the group
-  const result: Map<number, Set<number>> = new Map(
-    ramGroups
-      .map(modulePath => {
-        const root = byPath.get(modulePath);
-        if (!root) {
-          throw Error(`Group root ${modulePath} is not part of the bundle`);
-        }
-        return [
-          root.id,
-          // `subtree` yields the IDs of all transitive dependencies of a module
-          /* $FlowFixMe: assumes the module is always in the Map */
-          new Set(subtree(byPath.get(root.sourcePath), byPath)),
-        ];
-      })
-  );
-
-  if (ramGroups.length > 1) {
-    // build a map of all grouped module IDs to an array of group root IDs
-    const all = new ArrayMap();
-    for (const [parent, children] of result) {
-      for (const module of children) {
-        all.get(module).push(parent);
-      }
-    }
-
-    // find all module IDs that are part of more than one group
-    const doubles = filter(all, ([, parents]) => parents.length > 1);
-    for (const [moduleId, parents] of doubles) {
-      // remove them from their groups
-      /* $FlowFixMe: this assumes the element exists. */
-      parents.forEach(p => result.get(p).delete(moduleId));
-
-      // print a warning for each removed module
-      const parentNames = parents.map(byId.get, byId);
-      const lastName = parentNames.pop();
-      throw new Error(
-        /* $FlowFixMe: this assumes the element exists. */
-        `Module ${byId.get(moduleId)} belongs to groups ${
-          parentNames.join(', ')}, and ${lastName
-          }. Removing it from all groups.`
-      );
-    }
-  }
-
-  return result;
-}
-
 const isRawMappings = Array.isArray;
 
 module.exports = Bundle;

packages/metro-bundler/src/Bundler/__tests__/Bundle-test.js

@@ -379,7 +379,7 @@ describe('Bundle', () => {
       }).toThrow(
         new Error(
           `Module ${fsLocation('invariant')} belongs to groups ${fsLocation('React')}` +
-            `, and ${fsLocation('OtherFramework')}. Removing it from all groups.`,
+            `, and ${fsLocation('OtherFramework')}. Ensure that each module is only part of one group.`,
         ),
       );
     });

packages/metro-bundler/src/Bundler/util.js

@@ -16,6 +16,12 @@ const babelGenerate = require('babel-generator').default;
 const babylon = require('babylon');
 
 import type {AssetDescriptor} from '.';
+import type {ModuleTransportLike} from '../../../local-cli/bundle/types.flow';
+
+type SubTree<T: ModuleTransportLike> = (
+  moduleTransport: T,
+  moduleTransportsByPath: Map<string, T>,
+) => Generator<number, void, void>;
 
 const assetPropertyBlacklist = new Set([
   'files',
@@ -69,7 +75,82 @@ function filterObject(object, blacklist) {
   return copied;
 }
 
+function createRamBundleGroups<T: ModuleTransportLike>(
+  ramGroups: $ReadOnlyArray<string>,
+  groupableModules: $ReadOnlyArray<T>,
+  subtree: SubTree<T>,
+): Map<number, Set<number>> {
+  // build two maps that allow to lookup module data
+  // by path or (numeric) module id;
+  const byPath = new Map();
+  const byId = new Map();
+  groupableModules.forEach(m => {
+    byPath.set(m.sourcePath, m);
+    byId.set(m.id, m.sourcePath);
+  });
+
+  // build a map of group root IDs to an array of module IDs in the group
+  const result: Map<number, Set<number>> = new Map(
+    ramGroups
+      .map(modulePath => {
+        const root = byPath.get(modulePath);
+        if (root == null) {
+          throw Error(`Group root ${modulePath} is not part of the bundle`);
+        }
+        return [
+          root.id,
+          // `subtree` yields the IDs of all transitive dependencies of a module
+          new Set(subtree(root, byPath)),
+        ];
+      })
+  );
+
+  if (ramGroups.length > 1) {
+    // build a map of all grouped module IDs to an array of group root IDs
+    const all = new ArrayMap();
+    for (const [parent, children] of result) {
+      for (const module of children) {
+        all.get(module).push(parent);
+      }
+    }
+
+    // find all module IDs that are part of more than one group
+    const doubles = filter(all, ([, parents]) => parents.length > 1);
+    for (const [moduleId, parents] of doubles) {
+      const parentNames = parents.map(byId.get, byId);
+      const lastName = parentNames.pop();
+      throw new Error(
+        `Module ${byId.get(moduleId) || moduleId} belongs to groups ${
+          parentNames.join(', ')}, and ${String(lastName)
+          }. Ensure that each module is only part of one group.`
+      );
+    }
+  }
+
+  return result;
+}
+
+function * filter(iterator, predicate) {
+  for (const value of iterator) {
+    if (predicate(value)) {
+      yield value;
+    }
+  }
+}
+
+class ArrayMap extends Map {
+  get(key) {
+    let array = super.get(key);
+    if (!array) {
+      array = [];
+      this.set(key, array);
+    }
+    return array;
+  }
+}
+
 module.exports = {
+  createRamBundleGroups,
   generateAssetCodeFileAst,
   generateAssetTransformResult,
   isAssetTypeAnImage,

packages/metro-bundler/src/ModuleGraph/output/__tests__/indexed-ram-bundle-test.js

@@ -26,14 +26,14 @@ let ids, modules, requireCall;
 const idForPath = ({path}) => getId(path);
 beforeAll(() => {
   modules = [
-    makeModule('a', 'script'),
-    makeModule('b'),
-    makeModule('c'),
-    makeModule('d'),
+    makeModule('a', [], 'script'),
+    makeModule('b', ['c']),
+    makeModule('c', ['f']),
+    makeModule('d', ['e']),
     makeModule('e'),
     makeModule('f'),
   ];
-  requireCall = makeModule('r', 'script', 'require(1);');
+  requireCall = makeModule('r', [], 'script', 'require(1);');
 
   ids = new Map(modules.map(({file}, i) => [file.path, i]));
   ({code, map} = createRamBundle());
@@ -79,7 +79,7 @@ it('creates a source map', () => {
   expect(map.x_facebook_offsets).toEqual([1, 2, 3, 4, 5, 6]);
 });
 
-describe('Optimization:', () => {
+describe('Startup section optimization', () => {
   let last, preloaded;
   beforeAll(() => {
     last = modules[modules.length - 1];
@@ -130,12 +130,62 @@ describe('Optimization:', () => {
           return section;
         }
     ));
-
   });
 });
 
-function createRamBundle(preloadedModules = new Set()) {
-  const build = indexedRamBundle.createBuilder(preloadedModules);
+describe('RAM groups / common sections', () => {
+  let groups, groupHeads;
+  beforeAll(() => {
+    groups = [
+      [modules[1], modules[2], modules[5]],
+      [modules[3], modules[4]],
+    ];
+    groupHeads = groups.map(g => g[0]);
+    ({code, map} = createRamBundle(undefined, groupHeads.map(getPath)));
+  });
+
+  it('supports grouping the transitive dependencies of files into common sections', () => {
+    const {codeOffset, table} = parseOffsetTable(code);
+
+    groups.forEach(group => {
+      const [head, ...deps] = group.map(x => idForPath(x.file));
+      const groupEntry = table[head];
+      deps.forEach(id => expect(table[id]).toEqual(groupEntry));
+
+      const [offset, length] = groupEntry;
+      const groupCode = code.slice(codeOffset + offset, codeOffset + offset + length - 1);
+      expect(groupCode.toString())
+        .toEqual(group.map(m => m.file.code).join('\n'));
+    });
+  });
+
+  it('reflects section groups in the source map', () => {
+    expect(map.x_facebook_offsets).toEqual([1, 2, 2, 5, 5, 2]);
+    const maps = map.sections.slice(-2);
+    const toplevelOffsets = [2, 5];
+
+    maps.map((groupMap, i) => [groups[i], groupMap]).forEach(([group, groupMap], i) => {
+      const offsets = group.reduce(moduleLineOffsets, [])[0];
+      expect(groupMap).toEqual({
+        map: {
+          version: 3,
+          sections: group.map((module, j) => ({
+            map: module.file.map,
+            offset: {line: offsets[j], column: 0},
+          })),
+        },
+        offset: {line: toplevelOffsets[i], column: 0},
+      });
+    });
+  });
+
+  function moduleLineOffsets([offsets = [], line = 0], module) {
+    return [[...offsets, line], line + countLines(module)];
+  }
+});
+
+function createRamBundle(preloadedModules = new Set(), ramGroups) {
+  const build = indexedRamBundle.createBuilder(preloadedModules, ramGroups);
   const result = build({
     filename: 'arbitrary/filename.js',
     idForPath,
@@ -149,10 +199,10 @@ function createRamBundle(preloadedModules = new Set()) {
   return {code: result.code, map: result.map};
 }
 
-function makeModule(name, type = 'module', moduleCode = `var ${name};`) {
-  const path = `/${name}.js`;
+function makeModule(name, deps = [], type = 'module', moduleCode = `var ${name};`) {
+  const path = makeModulePath(name);
   return {
-    dependencies: [],
+    dependencies: deps.map(makeDependency),
     file: {
       code: type === 'module' ? makeModuleCode(moduleCode) : moduleCode,
       map: type !== 'module'
@@ -177,6 +227,18 @@ function makeModuleCode(moduleCode) {
   return `__d(() => {${moduleCode}})`;
 }
 
+function makeModulePath(name) {
+  return `/${name}.js`;
+}
+
+function makeDependency(name) {
+  const path = makeModulePath(name);
+  return {
+    id: name,
+    path,
+  };
+}
+
 function getId(path) {
   if (path === requireCall.file.path) {
     return -1;

packages/metro-bundler/src/ModuleGraph/output/indexed-ram-bundle.js

@@ -11,8 +11,10 @@
 'use strict';
 
 const buildSourceMapWithMetaData = require('../../../../local-cli/bundle/output/unbundle/build-unbundle-sourcemap-with-metadata.js');
+const nullthrows = require('fbjs/lib/nullthrows');
 
 const {buildTableAndContents, createModuleGroups} = require('../../../../local-cli/bundle/output/unbundle/as-indexed-file');
+const {createRamBundleGroups} = require('../../Bundler/util');
 const {concat} = require('./util');
 
 import type {FBIndexMap} from '../../lib/SourceMap.js';
@@ -23,12 +25,14 @@ function asIndexedRamBundle({
   idForPath,
   modules,
   preloadedModules,
+  ramGroupHeads,
   requireCalls,
 }) {
   const [startup, deferred] = partition(modules, preloadedModules);
   const startupModules = Array.from(concat(startup, requireCalls));
-  const deferredModules = deferred.map(m => toModuleTransport(m.file, idForPath));
-  const moduleGroups = createModuleGroups(new Map(), deferredModules);
+  const deferredModules = deferred.map(m => toModuleTransport(m, idForPath));
+  const ramGroups = createRamBundleGroups(ramGroupHeads || [], deferredModules, subtree);
+  const moduleGroups = createModuleGroups(ramGroups, deferredModules);
 
   const tableAndContents = buildTableAndContents(
     startupModules.map(getModuleCode).join('\n'),
@@ -41,15 +45,17 @@ function asIndexedRamBundle({
     code: Buffer.concat(tableAndContents),
     map: buildSourceMapWithMetaData({
       fixWrapperOffset: false,
-      startupModules: startupModules.map(m => toModuleTransport(m.file, idForPath)),
       lazyModules: deferredModules,
+      moduleGroups,
+      startupModules: startupModules.map(m => toModuleTransport(m, idForPath)),
     }),
   };
 }
 
-function toModuleTransport(file, idForPath) {
+function toModuleTransport({dependencies, file}, idForPath) {
   return {
     code: file.code,
+    dependencies,
     id: idForPath(file),
     map: file.map,
     name: file.path,
@@ -71,7 +77,26 @@ function partition(modules, preloadedModules) {
   return [startup, deferred];
 }
 
-function createBuilder(preloadedModules: Set<string>): OutputFn<FBIndexMap> {
-  return x => asIndexedRamBundle({preloadedModules, ...x});
+function *subtree(
+  moduleTransport,
+  moduleTransportsByPath,
+  seen = new Set(),
+) {
+  seen.add(moduleTransport.id);
+  for (const {path} of moduleTransport.dependencies) {
+    const dependency = nullthrows(moduleTransportsByPath.get(path));
+    if (!seen.has(dependency.id)) {
+      yield dependency.id;
+      yield *subtree(dependency, moduleTransportsByPath, seen);
+    }
+  }
 }
+
+function createBuilder(
+  preloadedModules: Set<string>,
+  ramGroupHeads: ?$ReadOnlyArray<string>,
+): OutputFn<FBIndexMap> {
+  return x => asIndexedRamBundle({...x, preloadedModules, ramGroupHeads});
+}
+
 exports.createBuilder = createBuilder;