Use ObjectMap instead of Map with proper comparator. (#6018)

* Use ObjectMap instead of Map with proper comparator. * use documentKeySet. * Address comments. * prettier. * Update names.
firebase · Feb 28, 2022 · 09e6243 · 09e6243
1 parent 27ed845
commit 09e6243
Show file tree

Hide file tree

Showing 8 changed files with 110 additions and 49 deletions.
diff --git a/packages/firestore/src/local/document_overlay_cache.ts b/packages/firestore/src/local/document_overlay_cache.ts
@@ -15,9 +15,8 @@
  * limitations under the License.
  */
 
-import { DocumentKeySet } from '../model/collections';
+import { DocumentKeySet, MutationMap, OverlayMap } from '../model/collections';
 import { DocumentKey } from '../model/document_key';
-import { Mutation } from '../model/mutation';
 import { Overlay } from '../model/overlay';
 import { ResourcePath } from '../model/path';
 
@@ -51,7 +50,7 @@ export interface DocumentOverlayCache {
   saveOverlays(
     transaction: PersistenceTransaction,
     largestBatchId: number,
-    overlays: Map<DocumentKey, Mutation>
+    overlays: MutationMap
   ): PersistencePromise<void>;
 
   /** Removes overlays for the given document keys and batch ID. */
@@ -74,7 +73,7 @@ export interface DocumentOverlayCache {
     transaction: PersistenceTransaction,
     collection: ResourcePath,
     sinceBatchId: number
-  ): PersistencePromise<Map<DocumentKey, Overlay>>;
+  ): PersistencePromise<OverlayMap>;
 
   /**
    * Returns `count` overlays with a batch ID higher than `sinceBatchId` for the
@@ -95,5 +94,5 @@ export interface DocumentOverlayCache {
     collectionGroup: string,
     sinceBatchId: number,
     count: number
-  ): PersistencePromise<Map<DocumentKey, Overlay>>;
+  ): PersistencePromise<OverlayMap>;
 }
diff --git a/packages/firestore/src/local/indexeddb_document_overlay_cache.ts b/packages/firestore/src/local/indexeddb_document_overlay_cache.ts
@@ -16,9 +16,13 @@
  */
 
 import { User } from '../auth/user';
-import { DocumentKeySet } from '../model/collections';
+import {
+  DocumentKeySet,
+  MutationMap,
+  OverlayMap,
+  newOverlayMap
+} from '../model/collections';
 import { DocumentKey } from '../model/document_key';
-import { Mutation } from '../model/mutation';
 import { Overlay } from '../model/overlay';
 import { ResourcePath } from '../model/path';
 
@@ -80,10 +84,10 @@ export class IndexedDbDocumentOverlayCache implements DocumentOverlayCache {
   saveOverlays(
     transaction: PersistenceTransaction,
     largestBatchId: number,
-    overlays: Map<DocumentKey, Mutation>
+    overlays: MutationMap
   ): PersistencePromise<void> {
     const promises: Array<PersistencePromise<void>> = [];
-    overlays.forEach(mutation => {
+    overlays.forEach((_, mutation) => {
       const overlay = new Overlay(largestBatchId, mutation);
       promises.push(this.saveOverlay(transaction, overlay));
     });
@@ -124,8 +128,8 @@ export class IndexedDbDocumentOverlayCache implements DocumentOverlayCache {
     transaction: PersistenceTransaction,
     collection: ResourcePath,
     sinceBatchId: number
-  ): PersistencePromise<Map<DocumentKey, Overlay>> {
-    const result = new Map<DocumentKey, Overlay>();
+  ): PersistencePromise<OverlayMap> {
+    const result = newOverlayMap();
     const collectionPath = encodeResourcePath(collection);
     // We want batch IDs larger than `sinceBatchId`, and so the lower bound
     // is not inclusive.
@@ -150,8 +154,8 @@ export class IndexedDbDocumentOverlayCache implements DocumentOverlayCache {
     collectionGroup: string,
     sinceBatchId: number,
     count: number
-  ): PersistencePromise<Map<DocumentKey, Overlay>> {
-    const result = new Map<DocumentKey, Overlay>();
+  ): PersistencePromise<OverlayMap> {
+    const result = newOverlayMap();
     let currentBatchId: number | undefined = undefined;
     // We want batch IDs larger than `sinceBatchId`, and so the lower bound
     // is not inclusive.
@@ -173,7 +177,7 @@ export class IndexedDbDocumentOverlayCache implements DocumentOverlayCache {
           // `count` if there are more overlays from the `currentBatchId`.
           const overlay = fromDbDocumentOverlay(this.serializer, dbOverlay);
           if (
-            result.size < count ||
+            result.size() < count ||
             overlay.largestBatchId === currentBatchId
           ) {
             result.set(overlay.getKey(), overlay);

diff --git a/packages/firestore/src/local/memory_document_overlay_cache.ts b/packages/firestore/src/local/memory_document_overlay_cache.ts
@@ -15,7 +15,13 @@
  * limitations under the License.
  */
 
-import { DocumentKeySet } from '../model/collections';
+import {
+  documentKeySet,
+  DocumentKeySet,
+  MutationMap,
+  OverlayMap,
+  newOverlayMap
+} from '../model/collections';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
 import { Overlay } from '../model/overlay';
@@ -35,7 +41,7 @@ export class MemoryDocumentOverlayCache implements DocumentOverlayCache {
   private overlays = new SortedMap<DocumentKey, Overlay>(
     DocumentKey.comparator
   );
-  private overlayByBatchId = new Map<number, Set<DocumentKey>>();
+  private overlayByBatchId = new Map<number, DocumentKeySet>();
 
   getOverlay(
     transaction: PersistenceTransaction,
@@ -47,9 +53,9 @@ export class MemoryDocumentOverlayCache implements DocumentOverlayCache {
   saveOverlays(
     transaction: PersistenceTransaction,
     largestBatchId: number,
-    overlays: Map<DocumentKey, Mutation>
+    overlays: MutationMap
   ): PersistencePromise<void> {
-    overlays.forEach(mutation => {
+    overlays.forEach((_, mutation) => {
       this.saveOverlay(transaction, largestBatchId, mutation);
     });
     return PersistencePromise.resolve();
@@ -72,8 +78,8 @@ export class MemoryDocumentOverlayCache implements DocumentOverlayCache {
     transaction: PersistenceTransaction,
     collection: ResourcePath,
     sinceBatchId: number
-  ): PersistencePromise<Map<DocumentKey, Overlay>> {
-    const result = new Map<DocumentKey, Overlay>();
+  ): PersistencePromise<OverlayMap> {
+    const result = newOverlayMap();
 
     const immediateChildrenPathLength = collection.length + 1;
     const prefix = new DocumentKey(collection.child(''));
@@ -102,8 +108,8 @@ export class MemoryDocumentOverlayCache implements DocumentOverlayCache {
     collectionGroup: string,
     sinceBatchId: number,
     count: number
-  ): PersistencePromise<Map<DocumentKey, Overlay>> {
-    let batchIdToOverlays = new SortedMap<number, Map<DocumentKey, Overlay>>(
+  ): PersistencePromise<OverlayMap> {
+    let batchIdToOverlays = new SortedMap<number, OverlayMap>(
       (key1: number, key2: number) => key1 - key2
     );
 
@@ -118,7 +124,7 @@ export class MemoryDocumentOverlayCache implements DocumentOverlayCache {
       if (overlay.largestBatchId > sinceBatchId) {
         let overlaysForBatchId = batchIdToOverlays.get(overlay.largestBatchId);
         if (overlaysForBatchId === null) {
-          overlaysForBatchId = new Map<DocumentKey, Overlay>();
+          overlaysForBatchId = newOverlayMap();
           batchIdToOverlays = batchIdToOverlays.insert(
             overlay.largestBatchId,
             overlaysForBatchId
@@ -128,13 +134,13 @@ export class MemoryDocumentOverlayCache implements DocumentOverlayCache {
       }
     }
 
-    const result = new Map<DocumentKey, Overlay>();
+    const result = newOverlayMap();
     const batchIter = batchIdToOverlays.getIterator();
     while (batchIter.hasNext()) {
       const entry = batchIter.getNext();
       const overlays = entry.value;
-      overlays.forEach((overlay, key) => result.set(key, overlay));
-      if (result.size >= count) {
+      overlays.forEach((key, overlay) => result.set(key, overlay));
+      if (result.size() >= count) {
         break;
       }
     }
@@ -153,7 +159,10 @@ export class MemoryDocumentOverlayCache implements DocumentOverlayCache {
     // Remove the association of the overlay to its batch id.
     const existing = this.overlays.get(mutation.key);
     if (existing !== null) {
-      this.overlayByBatchId.get(existing.largestBatchId)!.delete(mutation.key);
+      const newSet = this.overlayByBatchId
+        .get(existing.largestBatchId)!
+        .delete(mutation.key);
+      this.overlayByBatchId.set(existing.largestBatchId, newSet);
     }
 
     this.overlays = this.overlays.insert(
@@ -164,9 +173,9 @@ export class MemoryDocumentOverlayCache implements DocumentOverlayCache {
     // Create the association of this overlay to the given largestBatchId.
     let batch = this.overlayByBatchId.get(largestBatchId);
     if (batch === undefined) {
-      batch = new Set<DocumentKey>();
+      batch = documentKeySet();
       this.overlayByBatchId.set(largestBatchId, batch);
     }
-    batch.add(mutation.key);
+    this.overlayByBatchId.set(largestBatchId, batch.add(mutation.key));
   }
 }
diff --git a/packages/firestore/src/model/collections.ts b/packages/firestore/src/model/collections.ts
@@ -18,11 +18,14 @@
 import { SnapshotVersion } from '../core/snapshot_version';
 import { TargetId } from '../core/types';
 import { primitiveComparator } from '../util/misc';
+import { ObjectMap } from '../util/obj_map';
 import { SortedMap } from '../util/sorted_map';
 import { SortedSet } from '../util/sorted_set';
 
 import { Document, MutableDocument } from './document';
 import { DocumentKey } from './document_key';
+import { Mutation } from './mutation';
+import { Overlay } from './overlay';
 
 /** Miscellaneous collection types / constants. */
 
@@ -47,6 +50,22 @@ export function documentMap(): DocumentMap {
   return EMPTY_DOCUMENT_MAP;
 }
 
+export type OverlayMap = ObjectMap<DocumentKey, Overlay>;
+export function newOverlayMap(): OverlayMap {
+  return new ObjectMap<DocumentKey, Overlay>(
+    key => key.toString(),
+    (l, r) => l.isEqual(r)
+  );
+}
+
+export type MutationMap = ObjectMap<DocumentKey, Mutation>;
+export function newMutationMap(): MutationMap {
+  return new ObjectMap<DocumentKey, Mutation>(
+    key => key.toString(),
+    (l, r) => l.isEqual(r)
+  );
+}
+
 export type DocumentVersionMap = SortedMap<DocumentKey, SnapshotVersion>;
 const EMPTY_DOCUMENT_VERSION_MAP = new SortedMap<DocumentKey, SnapshotVersion>(
   DocumentKey.comparator

diff --git a/packages/firestore/src/util/obj_map.ts b/packages/firestore/src/util/obj_map.ts
@@ -36,6 +36,9 @@ export class ObjectMap<KeyType, ValueType> {
     [canonicalId: string]: Array<Entry<KeyType, ValueType>>;
   } = {};
 
+  /** The number of entries stored in the map */
+  private innerSize = 0;
+
   constructor(
     private mapKeyFn: (key: KeyType) => string,
     private equalsFn: (l: KeyType, r: KeyType) => boolean
@@ -66,15 +69,18 @@ export class ObjectMap<KeyType, ValueType> {
     const matches = this.inner[id];
     if (matches === undefined) {
       this.inner[id] = [[key, value]];
+      this.innerSize++;
       return;
     }
     for (let i = 0; i < matches.length; i++) {
       if (this.equalsFn(matches[i][0], key)) {
+        // This is updating an existing entry and does not increase `innerSize`.
         matches[i] = [key, value];
         return;
       }
     }
     matches.push([key, value]);
+    this.innerSize++;
   }
 
   /**
@@ -93,6 +99,7 @@ export class ObjectMap<KeyType, ValueType> {
         } else {
           matches.splice(i, 1);
         }
+        this.innerSize--;
         return true;
       }
     }
@@ -110,4 +117,8 @@ export class ObjectMap<KeyType, ValueType> {
   isEmpty(): boolean {
     return isEmpty(this.inner);
   }
+
+  size(): number {
+    return this.innerSize;
+  }
 }
diff --git a/packages/firestore/test/unit/local/document_overlay_cache.test.ts b/packages/firestore/test/unit/local/document_overlay_cache.test.ts
@@ -19,10 +19,13 @@ import { expect } from 'chai';
 import { User } from '../../../src/auth/user';
 import { IndexedDbPersistence } from '../../../src/local/indexeddb_persistence';
 import { Persistence } from '../../../src/local/persistence';
-import { documentKeySet, DocumentKeySet } from '../../../src/model/collections';
-import { DocumentKey } from '../../../src/model/document_key';
+import {
+  documentKeySet,
+  MutationMap,
+  OverlayMap,
+  newMutationMap
+} from '../../../src/model/collections';
 import { Mutation, mutationEquals } from '../../../src/model/mutation';
-import { Overlay } from '../../../src/model/overlay';
 import { addEqualityMatcher } from '../../util/equality_matcher';
 import {
   deleteMutation,
@@ -86,7 +89,7 @@ function genericDocumentOverlayCacheTests(): void {
     largestBatch: number,
     ...mutations: Mutation[]
   ): Promise<void> {
-    const data: Map<DocumentKey, Mutation> = new Map<DocumentKey, Mutation>();
+    const data: MutationMap = newMutationMap();
     for (const mutation of mutations) {
       data.set(mutation.key, mutation);
     }
@@ -97,7 +100,7 @@ function genericDocumentOverlayCacheTests(): void {
     largestBatch: number,
     ...overlayKeys: string[]
   ): Promise<void> {
-    const data: Map<DocumentKey, Mutation> = new Map<DocumentKey, Mutation>();
+    const data: MutationMap = newMutationMap();
     for (const overlayKey of overlayKeys) {
       data.set(key(overlayKey), setMutation(overlayKey, {}));
     }
@@ -115,16 +118,14 @@ function genericDocumentOverlayCacheTests(): void {
   }
 
   function verifyOverlayContains(
-    overlays: Map<DocumentKey, Overlay>,
+    overlays: OverlayMap,
     ...keys: string[]
   ): void {
-    const overlayKeys: DocumentKeySet = documentKeySet(
-      ...Array.from(overlays.keys())
-    );
-    const expectedKeys: DocumentKeySet = documentKeySet(
-      ...Array.from(keys.map(value => key(value)))
-    );
-    expect(overlayKeys.isEqual(expectedKeys)).to.equal(true);
+    let overlayKeys = documentKeySet();
+    overlays.forEach(overlayKey => (overlayKeys = overlayKeys.add(overlayKey)));
+    let expectedKeys = documentKeySet();
+    keys.forEach(value => (expectedKeys = expectedKeys.add(key(value))));
+    expect(overlayKeys.isEqual(expectedKeys)).to.deep.equal(true);
   }
 
   it('returns null when overlay is not found', async () => {

diff --git a/packages/firestore/test/unit/local/test_document_overlay_cache.ts b/packages/firestore/test/unit/local/test_document_overlay_cache.ts
@@ -17,7 +17,11 @@
 
 import { DocumentOverlayCache } from '../../../src/local/document_overlay_cache';
 import { Persistence } from '../../../src/local/persistence';
-import { DocumentKeySet } from '../../../src/model/collections';
+import {
+  DocumentKeySet,
+  MutationMap,
+  OverlayMap
+} from '../../../src/model/collections';
 import { DocumentKey } from '../../../src/model/document_key';
 import { Mutation } from '../../../src/model/mutation';
 import { Overlay } from '../../../src/model/overlay';
@@ -34,10 +38,7 @@ export class TestDocumentOverlayCache {
     private cache: DocumentOverlayCache
   ) {}
 
-  saveOverlays(
-    largestBatch: number,
-    data: Map<DocumentKey, Mutation>
-  ): Promise<void> {
+  saveOverlays(largestBatch: number, data: MutationMap): Promise<void> {
     return this.persistence.runTransaction('saveOverlays', 'readwrite', txn => {
       return this.cache.saveOverlays(txn, largestBatch, data);
     });
@@ -61,7 +62,7 @@ export class TestDocumentOverlayCache {
   getOverlaysForCollection(
     path: ResourcePath,
     sinceBatchId: number
-  ): Promise<Map<DocumentKey, Overlay>> {
+  ): Promise<OverlayMap> {
     return this.persistence.runTransaction(
       'getOverlaysForCollection',
       'readonly',
@@ -75,7 +76,7 @@ export class TestDocumentOverlayCache {
     collectionGroup: string,
     sinceBatchId: number,
     count: number
-  ): Promise<Map<DocumentKey, Overlay>> {
+  ): Promise<OverlayMap> {
     return this.persistence.runTransaction(
       'getOverlaysForCollectionGroup',
       'readonly',