Performance: use Map for jest-runtime module registry. (#8232)

* Performance: use Map for jest-runtime mock registry. * Update CHANGELOG.md
jestjs · Mar 29, 2019 · 384a0d9 · 384a0d9
1 parent e08be02
commit 384a0d9
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 46 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -46,6 +46,7 @@
 - `[jest-core]` Improve performance of SearchSource.findMatchingTests by 15% ([#8184](https://github.com/facebook/jest/pull/8184))
 - `[jest-resolve]` Optimize internal cache lookup performance ([#8183](https://github.com/facebook/jest/pull/8183))
 - `[jest-core]` Dramatically improve watch mode performance ([#8201](https://github.com/facebook/jest/pull/8201))
+- `[jest-runtime]` Use `Map` instead of `Object` for module registry ([#8232](https://github.com/facebook/jest/pull/8232))
 
 ## 24.5.0
 

diff --git a/packages/jest-runtime/src/__tests__/runtime_require_module.test.js b/packages/jest-runtime/src/__tests__/runtime_require_module.test.js
@@ -132,7 +132,7 @@ describe('Runtime requireModule', () => {
   });
   it('provides `require.main` to modules', () =>
     createRuntime(__filename).then(runtime => {
-      runtime._moduleRegistry[__filename] = module;
+      runtime._moduleRegistry.set(__filename, module);
       [
         './test_root/modules_with_main/export_main.js',
         './test_root/modules_with_main/re_export_main.js',

diff --git a/packages/jest-runtime/src/index.ts b/packages/jest-runtime/src/index.ts
@@ -47,7 +47,7 @@ type InternalModuleOptions = {
 
 type InitialModule = Partial<Module> &
   Pick<Module, 'children' | 'exports' | 'filename' | 'id' | 'loaded'>;
-type ModuleRegistry = {[key: string]: InitialModule | Module};
+type ModuleRegistry = Map<string, InitialModule | Module>;
 type ResolveOptions = Parameters<typeof require.resolve>[1];
 
 type BooleanObject = {[key: string]: boolean};
@@ -93,8 +93,8 @@ class Runtime {
   private _mockMetaDataCache: {
     [key: string]: MockFunctionMetadata<unknown, Array<unknown>>;
   };
-  private _mockRegistry: {[key: string]: any};
-  private _isolatedMockRegistry: {[key: string]: any} | null;
+  private _mockRegistry: Map<string, any>;
+  private _isolatedMockRegistry: Map<string, any> | null;
   private _moduleMocker: typeof jestMock;
   private _isolatedModuleRegistry: ModuleRegistry | null;
   private _moduleRegistry: ModuleRegistry;
@@ -127,15 +127,15 @@ class Runtime {
     this._currentlyExecutingModulePath = '';
     this._environment = environment;
     this._explicitShouldMock = Object.create(null);
-    this._internalModuleRegistry = Object.create(null);
+    this._internalModuleRegistry = new Map();
     this._isCurrentlyExecutingManualMock = null;
     this._mockFactories = Object.create(null);
-    this._mockRegistry = Object.create(null);
+    this._mockRegistry = new Map();
     // during setup, this cannot be null (and it's fine to explode if it is)
     this._moduleMocker = this._environment.moduleMocker!;
     this._isolatedModuleRegistry = null;
     this._isolatedMockRegistry = null;
-    this._moduleRegistry = Object.create(null);
+    this._moduleRegistry = new Map();
     this._needsCoverageMapped = new Set();
     this._resolver = resolver;
     this._scriptTransformer = new ScriptTransformer(config);
@@ -291,7 +291,7 @@ class Runtime {
       from,
       moduleName,
     );
-    let modulePath;
+    let modulePath: string | undefined;
 
     // Some old tests rely on this mocking behavior. Ideally we'll change this
     // to be more explicit.
@@ -320,7 +320,10 @@ class Runtime {
     let moduleRegistry;
 
     if (!options || !options.isInternalModule) {
-      if (this._moduleRegistry[modulePath] || !this._isolatedModuleRegistry) {
+      if (
+        this._moduleRegistry.get(modulePath) ||
+        !this._isolatedModuleRegistry
+      ) {
         moduleRegistry = this._moduleRegistry;
       } else {
         moduleRegistry = this._isolatedModuleRegistry;
@@ -329,29 +332,33 @@ class Runtime {
       moduleRegistry = this._internalModuleRegistry;
     }
 
-    if (!moduleRegistry[modulePath]) {
-      // We must register the pre-allocated module object first so that any
-      // circular dependencies that may arise while evaluating the module can
-      // be satisfied.
-      const localModule: InitialModule = {
-        children: [],
-        exports: {},
-        filename: modulePath,
-        id: modulePath,
-        loaded: false,
-      };
-      moduleRegistry[modulePath] = localModule;
-
-      this._loadModule(
-        localModule,
-        from,
-        moduleName,
-        modulePath,
-        options,
-        moduleRegistry,
-      );
+    const module = moduleRegistry.get(modulePath);
+    if (module) {
+      return module.exports;
     }
-    return moduleRegistry[modulePath].exports;
+
+    // We must register the pre-allocated module object first so that any
+    // circular dependencies that may arise while evaluating the module can
+    // be satisfied.
+    const localModule: InitialModule = {
+      children: [],
+      exports: {},
+      filename: modulePath,
+      id: modulePath,
+      loaded: false,
+    };
+    moduleRegistry.set(modulePath, localModule);
+
+    this._loadModule(
+      localModule,
+      from,
+      moduleName,
+      modulePath,
+      options,
+      moduleRegistry,
+    );
+
+    return localModule.exports;
   }
 
   requireInternalModule(from: Config.Path, to?: string) {
@@ -369,16 +376,21 @@ class Runtime {
       moduleName,
     );
 
-    if (this._isolatedMockRegistry && this._isolatedMockRegistry[moduleID]) {
-      return this._isolatedMockRegistry[moduleID];
-    } else if (this._mockRegistry[moduleID]) {
-      return this._mockRegistry[moduleID];
+    if (
+      this._isolatedMockRegistry &&
+      this._isolatedMockRegistry.get(moduleID)
+    ) {
+      return this._isolatedMockRegistry.get(moduleID);
+    } else if (this._mockRegistry.get(moduleID)) {
+      return this._mockRegistry.get(moduleID);
     }
 
     const mockRegistry = this._isolatedMockRegistry || this._mockRegistry;
 
     if (moduleID in this._mockFactories) {
-      return (mockRegistry[moduleID] = this._mockFactories[moduleID]());
+      const module = this._mockFactories[moduleID]();
+      mockRegistry.set(moduleID, module);
+      return module;
     }
 
     const manualMockOrStub = this._resolver.getMockModule(from, moduleName);
@@ -435,13 +447,13 @@ class Runtime {
         mockRegistry,
       );
 
-      mockRegistry[moduleID] = localModule.exports;
+      mockRegistry.set(moduleID, localModule.exports);
     } else {
       // Look for a real module to generate an automock from
-      mockRegistry[moduleID] = this._generateMock(from, moduleName);
+      mockRegistry.set(moduleID, this._generateMock(from, moduleName));
     }
 
-    return mockRegistry[moduleID];
+    return mockRegistry.get(moduleID);
   }
 
   private _loadModule(
@@ -495,8 +507,8 @@ class Runtime {
         'isolateModules cannot be nested inside another isolateModules.',
       );
     }
-    this._isolatedModuleRegistry = Object.create(null);
-    this._isolatedMockRegistry = Object.create(null);
+    this._isolatedModuleRegistry = new Map();
+    this._isolatedMockRegistry = new Map();
     fn();
     this._isolatedModuleRegistry = null;
     this._isolatedMockRegistry = null;
@@ -505,8 +517,8 @@ class Runtime {
   resetModules() {
     this._isolatedModuleRegistry = null;
     this._isolatedMockRegistry = null;
-    this._mockRegistry = Object.create(null);
-    this._moduleRegistry = Object.create(null);
+    this._mockRegistry = new Map();
+    this._moduleRegistry = new Map();
 
     if (this._environment) {
       if (this._environment.global) {
@@ -678,7 +690,7 @@ class Runtime {
       enumerable: true,
       get() {
         const key = from || '';
-        return moduleRegistry[key] || null;
+        return moduleRegistry.get(key) || null;
       },
     });
 
@@ -770,8 +782,8 @@ class Runtime {
       // mocked has calls into side-effectful APIs on another module.
       const origMockRegistry = this._mockRegistry;
       const origModuleRegistry = this._moduleRegistry;
-      this._mockRegistry = Object.create(null);
-      this._moduleRegistry = Object.create(null);
+      this._mockRegistry = new Map();
+      this._moduleRegistry = new Map();
 
       const moduleExports = this.requireModule(from, moduleName);