Avoid Map lookup serializing WeakMap/WeakSet [perf]

lib/init/index.js

@@ -50,13 +50,13 @@ module.exports = (filename, module, require, nextBlockId, prefixNum) => {
 // Imports
 // These imports are after export to avoid circular requires in Jest tests
 const captureFunctions = require('./functions.js'),
-	{getWeakSets, getWeakMaps} = require('./weak.js'),
+	{shimWeakSet, shimWeakMap} = require('./weak.js'),
 	{populateGlobals} = require('./globals.js'),
 	patchModule = require('./module.js');
 
 // Init internal vars
 captureFunctions(specialFunctions);
 patchModule();
-internal.weakSets = getWeakSets();
-internal.weakMaps = getWeakMaps();
+internal.getWeakSetEntries = shimWeakSet();
+internal.getWeakMapEntries = shimWeakMap();
 populateGlobals(globals);

lib/init/weak.js

@@ -9,22 +9,20 @@
 
 // Exports
 
-module.exports = {getWeakSets, getWeakMaps};
+module.exports = {shimWeakSet, shimWeakMap};
 
 const WeakMapOriginal = WeakMap;
 
 /**
- * Shim `WeakSet` to record all `WeakSet`s created in program.
- * The shimmed `WeakSet` class captures entries, so they can be serialized.
- * Implementation uses `WeakRef` to avoid holding strong references to entries,
+ * Shim `WeakSet` with implementation which allow getting entries for serializing.
+ * Uses `WeakRef` to avoid holding strong references to entries,
  * and allows them to be garbage collected.
  * Adapted from https://github.com/Lunuy/iterable-weak/blob/master/src/IterableWeakSet.ts
  *
- * @returns {WeakMap} - `WeakMap` of all `WeakSet`s.
- *   `WeakMap` returned maps `WeakSet`s to a `Set` of `WeakRef`s to the `WeakSet`'s entries.
+ * @returns {Function} - Function which returns entries of a `WeakSet`
  */
-function getWeakSets() {
-	const weakSets = new WeakMapOriginal();
+function shimWeakSet() {
+	let getWeakSetEntries;
 
 	WeakSet = class WeakSet { // eslint-disable-line no-global-assign
 		// Set of `WeakRefs` to entries
@@ -35,8 +33,6 @@ function getWeakSets() {
 		#finalizationRegistry = new FinalizationRegistry(ref => this.#refs.delete(ref));
 
 		constructor(iterable = undefined) { // `= undefined` so `WeakSet.length === 0`
-			weakSets.set(this, this.#refs);
-
 			if (iterable != null) {
 				for (const element of iterable) {
 					// NB: Original calls `WeakSet.prototype.add` even if it's been overwritten by user
@@ -68,30 +64,37 @@ function getWeakSets() {
 			}
 			return this;
 		}
+
+		static {
+			getWeakSetEntries = (weakSet) => {
+				const entries = [];
+				for (const ref of weakSet.#refs) {
+					const value = ref.deref();
+					if (value) entries.push(value);
+				}
+				return entries;
+			};
+		}
 	};
 
 	// eslint-disable-next-line no-extend-native
 	Object.defineProperty(WeakSet.prototype, Symbol.toStringTag, {value: 'WeakSet', configurable: true});
 
-	return weakSets;
+	return getWeakSetEntries;
 }
 
 /**
- * Shim `WeakMap` to record all `WeakMap`s created in program.
- * The shimmed `WeakMap` class captures `WeakMap` keys, so they can be serialized.
- * Implementation uses `WeakRef` to avoid holding strong references to keys,
+ * Shim `WeakMap` with implementation which allow getting entries for serializing.
+ * Uses `WeakRef` to avoid holding strong references to keys,
  * and allows them to be garbage collected.
  *
  * Adapted from https://github.com/tc39/proposal-weakrefs#iterable-weakmaps
  * and https://github.com/Lunuy/iterable-weak/blob/master/src/IterableWeakMap.ts
  *
- * @returns {WeakMap} - `WeakMap` of all `WeakMap`s.
- *   `WeakMap` returned maps `WeakMap`s to objects with properties `.refs` and `.mappings`.
- *   `.refs` is a `Set` of `WeakRef`s to the `WeakMap`'s keys.
- *   `.mappings` is a `WeakMap` of keys to values.
+ * @returns {Function} - Function which returns entries of a `WeakMap`
  */
-function getWeakMaps() {
-	const weakMaps = new WeakMapOriginal();
+function shimWeakMap() {
+	let getWeakMapEntries;
 
 	WeakMap = class WeakMap { // eslint-disable-line no-global-assign
 		// Set of `WeakRefs` to keys
@@ -102,8 +105,6 @@ function getWeakMaps() {
 		#finalizationRegistry = new FinalizationRegistry(ref => this.#refs.delete(ref));
 
 		constructor(iterable = undefined) { // `= undefined` so `WeakMap.length === 0`
-			weakMaps.set(this, {refs: this.#refs, mapping: this.#mapping});
-
 			if (iterable != null) {
 				for (const [key, value] of iterable) {
 					// NB: Original calls `WeakMap.prototype.set` even if it's been overwritten by user
@@ -145,10 +146,22 @@ function getWeakMaps() {
 		has(key) {
 			return this.#mapping.has(key);
 		}
+
+		static {
+			getWeakMapEntries = (weakMap) => {
+				const mapping = weakMap.#mapping,
+					entries = [];
+				for (const ref of weakMap.#refs) {
+					const key = ref.deref();
+					if (key) entries.push([key, mapping.get(key).value]);
+				}
+				return entries;
+			};
+		}
 	};
 
 	// eslint-disable-next-line no-extend-native
 	Object.defineProperty(WeakMap.prototype, Symbol.toStringTag, {value: 'WeakMap', configurable: true});
 
-	return weakMaps;
+	return getWeakMapEntries;
 }

lib/serialize/setsMaps.js

@@ -10,7 +10,7 @@ const t = require('@babel/types');
 
 // Imports
 const {createDependency, createAssignment} = require('./records.js'),
-	{weakSets, weakMaps} = require('../shared/internal.js'),
+	{getWeakSetEntries, getWeakMapEntries} = require('../shared/internal.js'),
 	{recordIsCircular} = require('./utils.js');
 
 // Exports
@@ -24,14 +24,7 @@ module.exports = {
 	},
 
 	serializeWeakSet(set, record) {
-		const refs = weakSets.get(set);
-		const entries = [];
-		for (const ref of refs) {
-			const value = ref.deref();
-			if (value) entries.push(value);
-		}
-
-		return this.serializeSetLike(set, entries, WeakSet, true, record);
+		return this.serializeSetLike(set, getWeakSetEntries(set), WeakSet, true, record);
 	},
 
 	serializeSetLike(set, entries, ctor, isUnordered, record) {
@@ -69,14 +62,7 @@ module.exports = {
 	},
 
 	serializeWeakMap(map, record) {
-		const {refs, mapping} = weakMaps.get(map);
-		const entries = [];
-		for (const ref of refs) {
-			const key = ref.deref();
-			if (key) entries.push([key, mapping.get(key).value]);
-		}
-
-		return this.serializeMapLike(map, entries, WeakMap, true, record);
+		return this.serializeMapLike(map, getWeakMapEntries(map), WeakMap, true, record);
 	},
 
 	serializeMapLike(map, entries, ctor, isUnordered, record) {