Fix usages of Map/Set (switch to ObjectMap/SortedSet). Don't use documentMap, mutableDocumentMap, and documentKeySet.

ehsannas · ehsannas · commit c3fbc23d495e · 2022-02-18T11:59:03.000-06:00
diff --git a/packages/firestore/src/local/local_documents_view.ts b/packages/firestore/src/local/local_documents_view.ts
@@ -24,27 +24,28 @@ import {
 } from '../core/query';
 import { Timestamp } from '../lite-api/timestamp';
 import {
-  documentKeySet,
   DocumentKeySet,
+  DocumentKeyToOverlayMap,
   DocumentMap,
-  documentMap,
-  mutableDocumentMap,
-  MutableDocumentMap
+  MutableDocumentMap,
+  newDocumentKeyMap,
+  newDocumentKeyToMutationMap,
+  newDocumentKeyToOverlayMap
 } from '../model/collections';
 import { Document, MutableDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
 import { IndexOffset } from '../model/field_index';
 import { FieldMask } from '../model/field_mask';
 import {
   calculateOverlayMutation,
-  Mutation,
   mutationApplyToLocalView,
   PatchMutation
 } from '../model/mutation';
 import { Overlay } from '../model/overlay';
 import { ResourcePath } from '../model/path';
 import { debugAssert } from '../util/assert';
 import { SortedMap } from '../util/sorted_map';
+import { SortedSet } from '../util/sorted_set';
 
 import { DocumentOverlayCache } from './document_overlay_cache';
 import { IndexManager } from './index_manager';
@@ -110,9 +111,11 @@ export class LocalDocumentsView {
     return this.remoteDocumentCache
       .getEntries(transaction, keys)
       .next(docs =>
-        this.getLocalViewOfDocuments(transaction, docs, documentKeySet()).next(
-          () => docs as DocumentMap
-        )
+        this.getLocalViewOfDocuments(
+          transaction,
+          docs,
+          new SortedSet(DocumentKey.comparator)
+        ).next(() => docs as DocumentMap)
       );
   }
 
@@ -134,7 +137,7 @@ export class LocalDocumentsView {
     return this.computeViews(
       transaction,
       docs,
-      new Map<DocumentKey, Overlay>(),
+      newDocumentKeyToOverlayMap(),
       existenceStateChanged
     );
   }
@@ -146,11 +149,13 @@ export class LocalDocumentsView {
   computeViews(
     transaction: PersistenceTransaction,
     docs: MutableDocumentMap,
-    memoizedOverlays: Map<DocumentKey, Overlay>,
+    memoizedOverlays: DocumentKeyToOverlayMap,
     existenceStateChanged: DocumentKeySet
   ): PersistencePromise<DocumentMap> {
-    let results = documentMap();
-    let recalculateDocuments = mutableDocumentMap();
+    let results = new SortedMap<DocumentKey, Document>(DocumentKey.comparator);
+    let recalculateDocuments = new SortedMap<DocumentKey, MutableDocument>(
+      DocumentKey.comparator
+    );
     const promises: Array<PersistencePromise<void>> = [];
     docs.forEach((_, doc) => {
       const overlayPromise = memoizedOverlays.has(doc.key)
@@ -199,12 +204,12 @@ export class LocalDocumentsView {
     transaction: PersistenceTransaction,
     docs: MutableDocumentMap
   ): PersistencePromise<void> {
-    const masks = new Map<DocumentKey, FieldMask | null>();
+    const masks = newDocumentKeyMap<FieldMask | null>();
     // A reverse lookup map from batch id to the documents within that batch.
-    let documentsByBatchId = new SortedMap<number, Set<DocumentKey>>(
+    let documentsByBatchId = new SortedMap<number, DocumentKeySet>(
       (key1: number, key2: number) => key1 - key2
     );
-    const processed = new Set<DocumentKey>();
+    let processed = new SortedSet(DocumentKey.comparator);
     return this.mutationQueue
       .getAllMutationBatchesAffectingDocumentKeys(transaction, docs)
       .next(batches => {
@@ -218,10 +223,11 @@ export class LocalDocumentsView {
             if (documentsByBatchId.get(batch.batchId) === null) {
               documentsByBatchId = documentsByBatchId.insert(
                 batch.batchId,
-                new Set<DocumentKey>()
+                new SortedSet(DocumentKey.comparator)
               );
             }
-            documentsByBatchId.get(batch.batchId)!.add(key);
+            const newSet = documentsByBatchId.get(batch.batchId)!.add(key);
+            documentsByBatchId.insert(batch.batchId, newSet);
           });
         });
       })
@@ -234,7 +240,7 @@ export class LocalDocumentsView {
           const entry = iter.getNext();
           const batchId = entry.key;
           const keys = entry.value;
-          const overlays = new Map<DocumentKey, Mutation>();
+          const overlays = newDocumentKeyToMutationMap();
           keys.forEach(key => {
             if (!processed.has(key)) {
               // TODO: Should we change `overlays` type to Map<DK, Mutation|null>
@@ -246,7 +252,7 @@ export class LocalDocumentsView {
               if (overlayMutation !== null) {
                 overlays.set(key, overlayMutation);
               }
-              processed.add(key);
+              processed = processed.add(key);
             }
           });
           promises.push(
@@ -313,7 +319,9 @@ export class LocalDocumentsView {
     // Just do a simple document lookup.
     return this.getDocument(transaction, new DocumentKey(docPath)).next(
       document => {
-        let result = documentMap();
+        let result = new SortedMap<DocumentKey, Document>(
+          DocumentKey.comparator
+        );
         if (document.isFoundDocument()) {
           result = result.insert(document.key, document);
         }
@@ -332,7 +340,7 @@ export class LocalDocumentsView {
       'Currently we only support collection group queries at the root.'
     );
     const collectionId = query.collectionGroup!;
-    let results = documentMap();
+    let results = new SortedMap<DocumentKey, Document>(DocumentKey.comparator);
     return this.indexManager
       .getCollectionParents(transaction, collectionId)
       .next(parents => {
@@ -376,7 +384,7 @@ export class LocalDocumentsView {
       .next(overlays => {
         // As documents might match the query because of their overlay we need to
         // include documents for all overlays in the initial document set.
-        overlays.forEach(overlay => {
+        overlays.forEach((_, overlay) => {
           const key = overlay.getKey();
           if (remoteDocuments.get(key) === null) {
             remoteDocuments = remoteDocuments.insert(
@@ -387,7 +395,9 @@ export class LocalDocumentsView {
         });
 
         // Apply the overlays and match against the query.
-        let results = documentMap();
+        let results = new SortedMap<DocumentKey, Document>(
+          DocumentKey.comparator
+        );
         remoteDocuments.forEach((key, document) => {
           const overlay = overlays.get(key);
           if (overlay !== undefined) {
diff --git a/packages/firestore/src/model/collections.ts b/packages/firestore/src/model/collections.ts
@@ -66,6 +66,14 @@ export function newDocumentKeyToMutationMap(): DocumentKeyToMutationMap {
   );
 }
 
+export type DocumentKeyMap<T> = ObjectMap<DocumentKey, T>;
+export function newDocumentKeyMap<T>(): DocumentKeyMap<T> {
+  return new ObjectMap<DocumentKey, T>(
+    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/model/mutation.ts b/packages/firestore/src/model/mutation.ts
@@ -20,6 +20,7 @@ import { Timestamp } from '../lite-api/timestamp';
 import { Value as ProtoValue } from '../protos/firestore_proto_api';
 import { debugAssert, hardAssert } from '../util/assert';
 import { arrayEquals } from '../util/misc';
+import { SortedSet } from '../util/sorted_set';
 
 import { Document, MutableDocument } from './document';
 import { DocumentKey } from './document_key';
@@ -238,7 +239,7 @@ export function calculateOverlayMutation(
   } else {
     const docValue = doc.data;
     const patchValue = ObjectValue.empty();
-    const maskSet = new Set<FieldPath>();
+    let maskSet = new SortedSet<FieldPath>(FieldPath.comparator);
     mask.fields.forEach(path => {
       if (!maskSet.has(path)) {
         const value = docValue.field(path);
@@ -255,13 +256,13 @@ export function calculateOverlayMutation(
           path = path.popLast();
         }
         patchValue.set(path, docValue.field(path)!);
-        maskSet.add(path);
+        maskSet = maskSet.add(path);
       }
     });
     return new PatchMutation(
       doc.key,
       patchValue,
-      new FieldMask(Array.from(maskSet)),
+      new FieldMask(maskSet.toArray()),
       Precondition.none()
     );
   }
@@ -567,12 +568,17 @@ function patchMutationApplyToLocalView(
     return null;
   }
 
-  const mergedMaskSet = new Set<FieldPath>(previousMask.fields);
-  mutation.fieldMask.fields.forEach(fieldPath => mergedMaskSet.add(fieldPath));
+  let mergedMaskSet = new SortedSet<FieldPath>(FieldPath.comparator);
+  previousMask.fields.forEach(
+    fieldPath => (mergedMaskSet = mergedMaskSet.add(fieldPath))
+  );
+  mutation.fieldMask.fields.forEach(
+    fieldPath => (mergedMaskSet = mergedMaskSet.add(fieldPath))
+  );
   mutation.fieldTransforms
     .map(transform => transform.field)
-    .forEach(fieldPath => mergedMaskSet.add(fieldPath));
-  return new FieldMask(Array.from(mergedMaskSet));
+    .forEach(fieldPath => (mergedMaskSet = mergedMaskSet.add(fieldPath)));
+  return new FieldMask(mergedMaskSet.toArray());
 }
 
 /**
diff --git a/packages/firestore/src/model/mutation_batch.ts b/packages/firestore/src/model/mutation_batch.ts
@@ -24,12 +24,13 @@ import { arrayEquals } from '../util/misc';
 import {
   documentKeySet,
   DocumentKeySet,
+  DocumentKeyToMutationMap,
   DocumentMap,
   DocumentVersionMap,
-  documentVersionMap
+  documentVersionMap,
+  newDocumentKeyToMutationMap
 } from './collections';
 import { MutableDocument } from './document';
-import { DocumentKey } from './document_key';
 import { FieldMask } from './field_mask';
 import {
   calculateOverlayMutation,
@@ -148,13 +149,11 @@ export class MutationBatch {
    * this batch. Returns a `DocumentKey` to `Mutation` map which can be used to
    * replace all the mutation applications.
    */
-  applyToLocalDocumentSet(
-    documentMap: DocumentMap
-  ): Map<DocumentKey, Mutation> {
+  applyToLocalDocumentSet(documentMap: DocumentMap): DocumentKeyToMutationMap {
     // TODO(mrschmidt): This implementation is O(n^2). If we apply the mutations
     // directly (as done in `applyToLocalView()`), we can reduce the complexity
     // to O(n).
-    const overlays = new Map<DocumentKey, Mutation>();
+    const overlays = newDocumentKeyToMutationMap();
     this.mutations.forEach(m => {
       const document = documentMap.get(m.key)!;
       // TODO(mutabledocuments): This method should take a MutableDocumentMap
diff --git a/packages/firestore/test/unit/local/query_engine.test.ts b/packages/firestore/test/unit/local/query_engine.test.ts
@@ -32,7 +32,11 @@ import { PersistenceTransaction } from '../../../src/local/persistence_transacti
 import { QueryEngine } from '../../../src/local/query_engine';
 import { RemoteDocumentCache } from '../../../src/local/remote_document_cache';
 import { TargetCache } from '../../../src/local/target_cache';
-import { documentKeySet, DocumentMap } from '../../../src/model/collections';
+import {
+  documentKeySet,
+  DocumentMap,
+  newDocumentKeyToMutationMap
+} from '../../../src/model/collections';
 import { Document, MutableDocument } from '../../../src/model/document';
 import { DocumentKey } from '../../../src/model/document_key';
 import { DocumentSet } from '../../../src/model/document_set';
@@ -133,7 +137,7 @@ describe('QueryEngine', () => {
       return mutationQueue
         .addMutationBatch(txn, Timestamp.now(), [], [mutation])
         .next(batch => {
-          const overlayMap = new Map<DocumentKey, Mutation>();
+          const overlayMap = newDocumentKeyToMutationMap();
           overlayMap.set(mutation.key, mutation);
           return documentOverlayCache.saveOverlays(
             txn,
