Port performance optimizations to speed up reading large collections from Android

packages/firestore/src/local/indexeddb_mutation_queue.ts

@@ -25,6 +25,7 @@ import { BATCHID_UNKNOWN, MutationBatch } from '../model/mutation_batch';
 import { ResourcePath } from '../model/path';
 import { assert, fail } from '../util/assert';
 import { primitiveComparator } from '../util/misc';
+import { SortedMap } from '../util/sorted_map';
 import { SortedSet } from '../util/sorted_set';
 
 import * as EncodedResourcePath from './encoded_resource_path';
@@ -46,6 +47,8 @@ import { PersistenceTransaction, ReferenceDelegate } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { SimpleDbStore, SimpleDbTransaction } from './simple_db';
 
+import { AnyJs } from '../../src/util/misc';
+
 /** A mutation queue for a specific user, backed by IndexedDB. */
 export class IndexedDbMutationQueue implements MutationQueue {
   /**
@@ -325,7 +328,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
 
   getAllMutationBatchesAffectingDocumentKeys(
     transaction: PersistenceTransaction,
-    documentKeys: DocumentKeySet
+    documentKeys: SortedMap<DocumentKey, AnyJs>
   ): PersistencePromise<MutationBatch[]> {
     let uniqueBatchIDs = new SortedSet<BatchId>(primitiveComparator);
 

packages/firestore/src/local/indexeddb_remote_document_cache.ts

@@ -20,16 +20,16 @@ import {
   documentKeySet,
   DocumentMap,
   documentMap,
-  DocumentSizeEntry,
   DocumentSizeEntries,
-  nullableMaybeDocumentMap,
-  NullableMaybeDocumentMap,
+  DocumentSizeEntry,
   MaybeDocumentMap,
-  maybeDocumentMap
+  maybeDocumentMap,
+  nullableMaybeDocumentMap,
+  NullableMaybeDocumentMap
 } from '../model/collections';
 import { Document, MaybeDocument, NoDocument } from '../model/document';
-import { SortedMap } from '../util/sorted_map';
 import { DocumentKey } from '../model/document_key';
+import { SortedMap } from '../util/sorted_map';
 
 import { SnapshotVersion } from '../core/snapshot_version';
 import { assert, fail } from '../util/assert';
@@ -187,9 +187,21 @@ export class IndexedDbRemoteDocumentCache implements RemoteDocumentCache {
     transaction: PersistenceTransaction,
     documentKeys: DocumentKeySet
   ): PersistencePromise<NullableMaybeDocumentMap> {
-    return this.getSizedEntries(transaction, documentKeys).next(
-      result => result.maybeDocuments
-    );
+    let results = nullableMaybeDocumentMap();
+    return this.forEachDbEntry(
+      transaction,
+      documentKeys,
+      (key, dbRemoteDoc) => {
+        if (dbRemoteDoc) {
+          results = results.insert(
+            key,
+            this.serializer.fromDbRemoteDocument(dbRemoteDoc)
+          );
+        } else {
+          results = results.insert(key, null);
+        }
+      }
+    ).next(() => results);
   }
 
   /**
@@ -206,46 +218,72 @@ export class IndexedDbRemoteDocumentCache implements RemoteDocumentCache {
   ): PersistencePromise<DocumentSizeEntries> {
     let results = nullableMaybeDocumentMap();
     let sizeMap = new SortedMap<DocumentKey, number>(DocumentKey.comparator);
+    return this.forEachDbEntry(
+      transaction,
+      documentKeys,
+      (key, dbRemoteDoc) => {
+        if (dbRemoteDoc) {
+          results = results.insert(
+            key,
+            this.serializer.fromDbRemoteDocument(dbRemoteDoc)
+          );
+          sizeMap = sizeMap.insert(key, dbDocumentSize(dbRemoteDoc));
+        } else {
+          results = results.insert(key, null);
+          sizeMap = sizeMap.insert(key, 0);
+        }
+      }
+    ).next(() => {
+      return { maybeDocuments: results, sizeMap };
+    });
+  }
+
+  forEachDbEntry(
+    transaction: PersistenceTransaction,
+    documentKeys: DocumentKeySet,
+    callback: (key: DocumentKey, doc: DbRemoteDocument | null) => void
+  ): PersistencePromise<void> {
     if (documentKeys.isEmpty()) {
-      return PersistencePromise.resolve({ maybeDocuments: results, sizeMap });
+      return PersistencePromise.resolve();
     }
 
     const range = IDBKeyRange.bound(
       documentKeys.first()!.path.toArray(),
       documentKeys.last()!.path.toArray()
     );
-    let key = documentKeys.first();
+    const keyIter = documentKeys.getIterator();
+    let nextKey: DocumentKey | null = keyIter.getNext();
 
     return remoteDocumentsStore(transaction)
       .iterate({ range }, (potentialKeyRaw, dbRemoteDoc, control) => {
         const potentialKey = DocumentKey.fromSegments(potentialKeyRaw);
-        while (DocumentKey.comparator(key!, potentialKey) != 1) {
-          if (key!.isEqual(potentialKey)) {
-            results = results.insert(
-              key!,
-              this.serializer.fromDbRemoteDocument(dbRemoteDoc)
-            );
-            sizeMap = sizeMap.insert(key!, dbDocumentSize(dbRemoteDoc));
-          } else {
-            results = results.insert(key!, null);
-            sizeMap = sizeMap.insert(key!, 0);
-          }
 
-          key = documentKeys.firstAfter(key!);
-          if (!key) {
-            control.done();
-            return;
-          }
-          control.skip(key!.path.toArray());
+        // Go through keys not found in cache.
+        while (nextKey && DocumentKey.comparator(nextKey!, potentialKey) < 0) {
+          callback(nextKey!, null);
+          nextKey = keyIter.getNext();
+        }
+
+        if (nextKey && nextKey!.isEqual(potentialKey)) {
+          // Key found in cache.
+          callback(nextKey!, dbRemoteDoc);
+          nextKey = keyIter.hasNext() ? keyIter.getNext() : null;
+        }
+
+        // Skip to the next key (if there is one).
+        if (nextKey) {
+          control.skip(nextKey!.path.toArray());
+        } else {
+          control.done();
         }
       })
       .next(() => {
-        while (key) {
-          results = results.insert(key, null);
-          sizeMap = sizeMap.insert(key, 0);
-          key = documentKeys.firstAfter(key!);
+        // The rest of the keys are not in the cache. One case where `iterate`
+        // above won't go through them is when the cache is empty.
+        while (nextKey) {
+          callback(nextKey!, null);
+          nextKey = keyIter.hasNext() ? keyIter.getNext() : null;
         }
-        return { maybeDocuments: results, sizeMap };
       });
   }
 

packages/firestore/src/local/local_documents_view.ts

@@ -83,17 +83,15 @@ export class LocalDocumentsView {
     transaction: PersistenceTransaction,
     docs: NullableMaybeDocumentMap,
     batches: MutationBatch[]
-  ): PersistencePromise<NullableMaybeDocumentMap> {
+  ): NullableMaybeDocumentMap {
     let results = nullableMaybeDocumentMap();
-    return new PersistencePromise<NullableMaybeDocumentMap>(resolve => {
-      docs.forEach((key, localView) => {
-        for (const batch of batches) {
-          localView = batch.applyToLocalView(key, localView);
-        }
-        results = results.insert(key, localView);
-      });
-      resolve(results);
+    docs.forEach((key, localView) => {
+      for (const batch of batches) {
+        localView = batch.applyToLocalView(key, localView);
+      }
+      results = results.insert(key, localView);
     });
+    return results;
   }
 
   /**
@@ -119,17 +117,14 @@ export class LocalDocumentsView {
     transaction: PersistenceTransaction,
     baseDocs: NullableMaybeDocumentMap
   ): PersistencePromise<MaybeDocumentMap> {
-    let allKeys = documentKeySet();
-    baseDocs.forEach(key => {
-      allKeys = allKeys.add(key);
-    });
-
     return this.mutationQueue
-      .getAllMutationBatchesAffectingDocumentKeys(transaction, allKeys)
-      .next(batches =>
-        this.applyLocalMutationsToDocuments(transaction, baseDocs, batches)
-      )
-      .next(docs => {
+      .getAllMutationBatchesAffectingDocumentKeys(transaction, baseDocs)
+      .next(batches => {
+        const docs = this.applyLocalMutationsToDocuments(
+          transaction,
+          baseDocs,
+          batches
+        );
         let results = maybeDocumentMap();
         docs.forEach((key, maybeDoc) => {
           // TODO(http://b/32275378): Don't conflate missing / deleted.

packages/firestore/src/local/local_serializer.ts

@@ -71,12 +71,9 @@ export class LocalSerializer {
   /** Encodes a document for storage locally. */
   toDbRemoteDocument(maybeDoc: MaybeDocument): DbRemoteDocument {
     if (maybeDoc instanceof Document) {
-      let doc: api.Document;
-      if (maybeDoc.proto) {
-        doc = maybeDoc.proto;
-      } else {
-        doc = this.remoteSerializer.toDocument(maybeDoc);
-      }
+      const doc = maybeDoc.proto
+        ? maybeDoc.proto
+        : this.remoteSerializer.toDocument(maybeDoc);
       const hasCommittedMutations = maybeDoc.hasCommittedMutations;
       return new DbRemoteDocument(
         /* unknownDocument= */ null,

packages/firestore/src/local/memory_mutation_queue.ts

@@ -24,13 +24,16 @@ import { BATCHID_UNKNOWN, MutationBatch } from '../model/mutation_batch';
 import { emptyByteString } from '../platform/platform';
 import { assert } from '../util/assert';
 import { primitiveComparator } from '../util/misc';
+import { SortedMap } from '../util/sorted_map';
 import { SortedSet } from '../util/sorted_set';
 
 import { MutationQueue } from './mutation_queue';
 import { PersistenceTransaction, ReferenceDelegate } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { DocReference } from './reference_set';
 
+import { AnyJs } from '../../src/util/misc';
+
 export class MemoryMutationQueue implements MutationQueue {
   /**
    * The set of all mutations that have been sent but not yet been applied to
@@ -203,7 +206,7 @@ export class MemoryMutationQueue implements MutationQueue {
 
   getAllMutationBatchesAffectingDocumentKeys(
     transaction: PersistenceTransaction,
-    documentKeys: DocumentKeySet
+    documentKeys: SortedMap<DocumentKey, AnyJs>
   ): PersistencePromise<MutationBatch[]> {
     let uniqueBatchIDs = new SortedSet<number>(primitiveComparator);
 

packages/firestore/src/local/mutation_queue.ts

@@ -21,10 +21,13 @@ import { DocumentKeySet } from '../model/collections';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
 import { MutationBatch } from '../model/mutation_batch';
+import { SortedMap } from '../util/sorted_map';
 
 import { PersistenceTransaction } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 
+import { AnyJs } from '../../src/util/misc';
+
 /** A queue of mutations to apply to the remote store. */
 export interface MutationQueue {
   /** Returns true if this queue contains no mutation batches. */
@@ -133,7 +136,7 @@ export interface MutationQueue {
   // TODO(mcg): This should really return an enumerator
   getAllMutationBatchesAffectingDocumentKeys(
     transaction: PersistenceTransaction,
-    documentKeys: DocumentKeySet
+    documentKeys: SortedMap<DocumentKey, AnyJs>
   ): PersistencePromise<MutationBatch[]>;
 
   /**

packages/firestore/src/local/remote_document_change_buffer.ts

@@ -108,8 +108,7 @@ export abstract class RemoteDocumentChangeBuffer {
   }
 
   /**
-   * Looks up several entries in the cache.
-   * checked, and if no buffered change applies, this will forward to
+   * Looks up several entries in the cache, orwarding to
    * `RemoteDocumentCache.getEntry()`.
    *
    * @param transaction The transaction in which to perform any persistence
@@ -127,12 +126,14 @@ export abstract class RemoteDocumentChangeBuffer {
 
     // Record the size of everything we load from the cache so we can compute
     // a delta later.
-    return this.getAllFromCache(transaction, documentKeys).next(getResult => {
-      getResult.sizeMap.forEach((documentKey, size) => {
-        this.documentSizes.set(documentKey, size);
-      });
-      return getResult.maybeDocuments;
-    });
+    return this.getAllFromCache(transaction, documentKeys).next(
+      ({ maybeDocuments, sizeMap }) => {
+        sizeMap.forEach((documentKey, size) => {
+          this.documentSizes.set(documentKey, size);
+        });
+        return maybeDocuments;
+      }
+    );
   }
 
   /**

packages/firestore/src/util/sorted_set.ts

@@ -97,25 +97,20 @@ export class SortedSet<T> {
     }
   }
 
-  /** Finds the least element greater than `elem`. */
-  firstAfter(elem: T): T | null {
-    const iter = this.data.getIteratorFrom(elem);
-    if (!iter.hasNext()) {
-      return null;
-    }
-    if (this.comparator(iter.peek()!.key, elem) == 0) {
-      // Skip
-      iter.getNext();
-    }
-    return iter.hasNext() ? iter.getNext().key : null;
-  }
-
   /** Finds the least element greater than or equal to `elem`. */
   firstAfterOrEqual(elem: T): T | null {
     const iter = this.data.getIteratorFrom(elem);
     return iter.hasNext() ? iter.getNext().key : null;
   }
 
+  getIterator(): SortedSetIterator<T> {
+    return new SortedSetIterator<T>(this.data.getIterator());
+  }
+
+  getIteratorFrom(key: T): SortedSetIterator<T> {
+    return new SortedSetIterator<T>(this.data.getIteratorFrom(key));
+  }
+
   /** Inserts or updates an element */
   add(elem: T): SortedSet<T> {
     return this.copy(this.data.remove(elem).insert(elem, true));
@@ -173,3 +168,15 @@ export class SortedSet<T> {
     return result;
   }
 }
+
+export class SortedSetIterator<T> {
+  constructor(private iter: SortedMapIterator<T, boolean>) {}
+
+  getNext(): T {
+    return this.iter.getNext().key;
+  }
+
+  hasNext(): boolean {
+    return this.iter.hasNext();
+  }
+}