perf(core): optimize memory allocation when reconcilling lists (#52245)

pkozlowski-opensource · alxhub · commit e3a6bf9b6c3b · 2023-12-08T15:00:13.000-08:00
This change replaces the implementation of the multi-map used to store detached views while reconciling lists. The new implementation optimizes memory allocation for such map and avoid arrays allocation when there are no duplicated keys. PR Close #52245
diff --git a/packages/core/src/render3/list_reconciliation.ts b/packages/core/src/render3/list_reconciliation.ts
@@ -7,6 +7,7 @@
  */
 
 import {TrackByFunction} from '../change_detection';
+import {assertNotSame} from '../util/assert';
 
 /**
  * A type representing the live collection to be reconciled with any new (incoming) collection. This
@@ -86,7 +87,7 @@ function valuesMatching<V>(
 export function reconcile<T, V>(
     liveCollection: LiveCollection<T, V>, newCollection: Iterable<V>|undefined|null,
     trackByFn: TrackByFunction<V>): void {
-  let detachedItems: MultiMap<unknown, T>|undefined = undefined;
+  let detachedItems: UniqueValueMultiKeyMap<unknown, T>|undefined = undefined;
   let liveKeysInTheFuture: Set<unknown>|undefined = undefined;
 
   let liveStartIdx = 0;
@@ -148,7 +149,7 @@ export function reconcile<T, V>(
 
       // Fallback to the slow path: we need to learn more about the content of the live and new
       // collections.
-      detachedItems ??= new MultiMap();
+      detachedItems ??= new UniqueValueMultiKeyMap();
       liveKeysInTheFuture ??=
           initLiveItemsInTheFuture(liveCollection, liveStartIdx, liveEndIdx, trackByFn);
 
@@ -197,7 +198,7 @@ export function reconcile<T, V>(
         liveStartIdx++;
         newIterationResult = newCollectionIterator.next();
       } else {
-        detachedItems ??= new MultiMap();
+        detachedItems ??= new UniqueValueMultiKeyMap();
         liveKeysInTheFuture ??=
             initLiveItemsInTheFuture(liveCollection, liveStartIdx, liveEndIdx, trackByFn);
 
@@ -246,8 +247,9 @@ export function reconcile<T, V>(
 }
 
 function attachPreviouslyDetached<T, V>(
-    prevCollection: LiveCollection<T, V>, detachedItems: MultiMap<unknown, T>|undefined,
-    index: number, key: unknown): boolean {
+    prevCollection: LiveCollection<T, V>,
+    detachedItems: UniqueValueMultiKeyMap<unknown, T>|undefined, index: number,
+    key: unknown): boolean {
   if (detachedItems !== undefined && detachedItems.has(key)) {
     prevCollection.attach(index, detachedItems.get(key)!);
     detachedItems.delete(key);
@@ -257,7 +259,8 @@ function attachPreviouslyDetached<T, V>(
 }
 
 function createOrAttach<T, V>(
-    liveCollection: LiveCollection<T, V>, detachedItems: MultiMap<unknown, T>|undefined,
+    liveCollection: LiveCollection<T, V>,
+    detachedItems: UniqueValueMultiKeyMap<unknown, T>|undefined,
     trackByFn: TrackByFunction<unknown>, index: number, value: V) {
   if (!attachPreviouslyDetached(liveCollection, detachedItems, index, trackByFn(index, value))) {
     const newItem = liveCollection.create(index, value);
@@ -277,43 +280,78 @@ function initLiveItemsInTheFuture<T>(
   return keys;
 }
 
-class MultiMap<K, V> {
-  private map = new Map<K, Array<V>>();
+/**
+ * A specific, partial implementation of the Map interface with the following characteristics:
+ * - allows multiple values for a given key;
+ * - maintain FIFO order for multiple values corresponding to a given key;
+ * - assumes that all values are unique.
+ *
+ * The implementation aims at having the minimal overhead for cases where keys are _not_ duplicated
+ * (the most common case in the list reconciliation algorithm). To achieve this, the first value for
+ * a given key is stored in a regular map. Then, when more values are set for a given key, we
+ * maintain a form of linked list in a separate map. To maintain this linked list we assume that all
+ * values (in the entire collection) are unique.
+ */
+export class UniqueValueMultiKeyMap<K, V> {
+  // A map from a key to the first value corresponding to this key.
+  private kvMap = new Map<K, V>();
+  // A map that acts as a linked list of values - each value maps to the next value in this "linked
+  // list" (this only works if values are unique). Allocated lazily to avoid memory consumption when
+  // there are no duplicated values.
+  private _vMap: Map<V, V>|undefined = undefined;
 
   has(key: K): boolean {
-    const listOfKeys = this.map.get(key);
-    return listOfKeys !== undefined && listOfKeys.length > 0;
+    return this.kvMap.has(key);
   }
 
   delete(key: K): boolean {
-    const listOfKeys = this.map.get(key);
-    if (listOfKeys !== undefined) {
-      listOfKeys.shift();
-      return true;
+    if (!this.has(key)) return false;
+
+    const value = this.kvMap.get(key)!;
+    if (this._vMap !== undefined && this._vMap.has(value)) {
+      this.kvMap.set(key, this._vMap.get(value)!);
+      this._vMap.delete(value);
+    } else {
+      this.kvMap.delete(key);
     }
-    return false;
+
+    return true;
   }
 
   get(key: K): V|undefined {
-    const listOfKeys = this.map.get(key);
-    return listOfKeys !== undefined && listOfKeys.length > 0 ? listOfKeys[0] : undefined;
+    return this.kvMap.get(key);
   }
 
   set(key: K, value: V): void {
-    // if value is array, they we always store it as [value].
-    if (!this.map.has(key)) {
-      this.map.set(key, [value]);
-      return;
+    if (this.kvMap.has(key)) {
+      let prevValue = this.kvMap.get(key)!;
+      ngDevMode &&
+          assertNotSame(
+              prevValue, value, `Detected a duplicated value ${value} for the key ${key}`);
+
+      if (this._vMap === undefined) {
+        this._vMap = new Map();
+      }
+
+      const vMap = this._vMap;
+      while (vMap.has(prevValue)) {
+        prevValue = vMap.get(prevValue)!;
+      }
+      vMap.set(prevValue, value);
+    } else {
+      this.kvMap.set(key, value);
     }
-    // THINK: this allows duplicate values, but I guess this is fine?
-    // Is the existing key an array or not?
-    this.map.get(key)?.push(value);
   }
 
   forEach(cb: (v: V, k: K) => void) {
-    for (const [key, values] of this.map) {
-      for (const value of values) {
-        cb(value, key);
+    for (let [key, value] of this.kvMap) {
+      cb(value, key);
+      if (this._vMap !== undefined) {
+        const vMap = this._vMap;
+        while (vMap.has(value)) {
+          value = vMap.get(value)!;
+          cb(value, key);
+        }
       }
     }
   }
diff --git a/packages/core/test/render3/list_reconciliation_spec.ts b/packages/core/test/render3/list_reconciliation_spec.ts
@@ -261,19 +261,21 @@ describe('list reconciliation', () => {
     });
 
     it('should create / reuse duplicated items as needed', () => {
-      const pc = new LoggingLiveCollection([1, 1, 2, 3]);
-      reconcile(pc, [2, 3, 1, 1, 1, 4], trackByIdentity);
+      const trackByKey = (idx: number, item: {k: number}) => item.k;
+      const pc =
+          new LoggingLiveCollection<{k: number}, {k: number}>([{k: 1}, {k: 1}, {k: 2}, {k: 3}]);
+      reconcile(pc, [{k: 2}, {k: 3}, {k: 1}, {k: 1}, {k: 1}, {k: 4}], trackByKey);
 
-      expect(pc.getCollection()).toEqual([2, 3, 1, 1, 1, 4]);
+      expect(pc.getCollection()).toEqual([{k: 2}, {k: 3}, {k: 1}, {k: 1}, {k: 1}, {k: 4}]);
       expect(pc.getLogs()).toEqual([
-        ['detach', 0, 1],
-        ['detach', 0, 1],
-        ['attach', 2, 1],
-        ['attach', 3, 1],
-        ['create', 4, 1],
-        ['attach', 4, 1],
-        ['create', 5, 4],
-        ['attach', 5, 4],
+        ['detach', 0, {k: 1}],
+        ['detach', 0, {k: 1}],
+        ['attach', 2, {k: 1}],
+        ['attach', 3, {k: 1}],
+        ['create', 4, {k: 1}],
+        ['attach', 4, {k: 1}],
+        ['create', 5, {k: 4}],
+        ['attach', 5, {k: 4}],
       ]);
     });
   });
diff --git a/packages/core/test/render3/multi_map_spec.ts b/packages/core/test/render3/multi_map_spec.ts
@@ -0,0 +1,62 @@
+/**
+ * @license
+ * Copyright Google LLC All Rights Reserved.
+ *
+ * Use of this source code is governed by an MIT-style license that can be
+ * found in the LICENSE file at https://angular.io/license
+ */
+
+import {UniqueValueMultiKeyMap} from '@angular/core/src/render3/list_reconciliation';
+
+describe('MultiMap', () => {
+  it('should set, get and remove items with duplicated keys', () => {
+    const map = new UniqueValueMultiKeyMap();
+
+    map.set('k', 'v1');
+    map.set('k', 'v2');
+
+    expect(map.has('k')).toBeTrue();
+    expect(map.get('k')).toBe('v1');
+
+    map.delete('k');
+    expect(map.has('k')).toBeTrue();
+    expect(map.get('k')).toBe('v2');
+
+    map.delete('k');
+    expect(map.has('k')).toBeFalse();
+  });
+
+  it('should set, get and remove items without duplicated keys', () => {
+    const map = new UniqueValueMultiKeyMap();
+
+    map.set('k', 'v1');
+
+    expect(map.has('k')).toBeTrue();
+    expect(map.get('k')).toBe('v1');
+
+    map.delete('k');
+    expect(map.has('k')).toBeFalse();
+  });
+
+  it('should iterate with forEach', () => {
+    const map = new UniqueValueMultiKeyMap<string, string>();
+
+    map.set('km', 'v1');
+    map.set('km', 'v2');
+    map.set('ks', 'v');
+
+    const items: string[][] = [];
+
+    map.forEach((v, k) => items.push([v, k]));
+    expect(items).toEqual([['v1', 'km'], ['v2', 'km'], ['v', 'ks']]);
+  });
+
+  it('should throw upon detecting duplicate values', () => {
+    const map = new UniqueValueMultiKeyMap();
+
+    map.set('k', 'v');
+    expect(() => {
+      map.set('k', 'v');
+    }).toThrowError(/Detected a duplicated value v for the key k/);
+  });
+});
