Backport executeQuery() API changes

firebase-firestore/src/main/java/com/google/firebase/firestore/core/FirestoreClient.java

@@ -22,7 +22,6 @@
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
 import com.google.common.base.Function;
-import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.EventListener;
 import com.google.firebase.firestore.FirebaseFirestoreException;
@@ -38,6 +37,7 @@
 import com.google.firebase.firestore.local.MemoryPersistence;
 import com.google.firebase.firestore.local.Persistence;
 import com.google.firebase.firestore.local.QueryEngine;
+import com.google.firebase.firestore.local.QueryResult;
 import com.google.firebase.firestore.local.SQLitePersistence;
 import com.google.firebase.firestore.local.SimpleQueryEngine;
 import com.google.firebase.firestore.model.Document;
@@ -55,7 +55,6 @@
 import com.google.firebase.firestore.util.AsyncQueue;
 import com.google.firebase.firestore.util.Logger;
 import io.grpc.Status;
-import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -201,14 +200,9 @@ public Task<ViewSnapshot> getDocumentsFromLocalCache(Query query) {
     this.verifyNotTerminated();
     return asyncQueue.enqueue(
         () -> {
-          ImmutableSortedMap<DocumentKey, Document> docs = localStore.executeQuery(query);
-
-          View view =
-              new View(
-                  query,
-                  new ImmutableSortedSet<DocumentKey>(
-                      Collections.emptyList(), DocumentKey::compareTo));
-          View.DocumentChanges viewDocChanges = view.computeDocChanges(docs);
+          QueryResult queryResult = localStore.executeQuery(query, /* usePreviousResults= */ true);
+          View view = new View(query, queryResult.getRemoteKeys());
+          View.DocumentChanges viewDocChanges = view.computeDocChanges(queryResult.getDocuments());
           return view.applyChanges(viewDocChanges).getSnapshot();
         });
   }

firebase-firestore/src/main/java/com/google/firebase/firestore/core/SyncEngine.java

@@ -32,8 +32,8 @@
 import com.google.firebase.firestore.local.LocalWriteResult;
 import com.google.firebase.firestore.local.QueryData;
 import com.google.firebase.firestore.local.QueryPurpose;
+import com.google.firebase.firestore.local.QueryResult;
 import com.google.firebase.firestore.local.ReferenceSet;
-import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.MaybeDocument;
 import com.google.firebase.firestore.model.NoDocument;
@@ -191,13 +191,10 @@ public int listen(Query query) {
   private ViewSnapshot initializeViewAndComputeSnapshot(QueryData queryData) {
     Query query = queryData.getQuery();
 
-    ImmutableSortedSet<DocumentKey> remoteKeys =
-        localStore.getRemoteDocumentKeys(queryData.getTargetId());
-    ImmutableSortedMap<DocumentKey, Document> docs =
-        localStore.executeQuery(query, queryData, remoteKeys);
+    QueryResult queryResult = localStore.executeQuery(query, /* usePreviousResults= */ true);
 
-    View view = new View(query, remoteKeys);
-    View.DocumentChanges viewDocChanges = view.computeDocChanges(docs);
+    View view = new View(query, queryResult.getRemoteKeys());
+    View.DocumentChanges viewDocChanges = view.computeDocChanges(queryResult.getDocuments());
     ViewChange viewChange = view.applyChanges(viewDocChanges);
     hardAssert(
         view.getLimboDocuments().size() == 0,
@@ -530,10 +527,9 @@ private void emitNewSnapsAndNotifyLocalStore(
         // The query has a limit and some docs were removed/updated, so we need to re-run the query
         // against the local store to make sure we didn't lose any good docs that had been past the
         // limit.
-        ImmutableSortedMap<DocumentKey, Document> docs =
-            localStore.executeQuery(
-                queryView.getQuery(), /* queryData= */ null, DocumentKey.emptyKeySet());
-        viewDocChanges = view.computeDocChanges(docs, viewDocChanges);
+        QueryResult queryResult =
+            localStore.executeQuery(queryView.getQuery(), /* usePreviousResults= */ false);
+        viewDocChanges = view.computeDocChanges(queryResult.getDocuments(), viewDocChanges);
       }
       TargetChange targetChange =
           remoteEvent == null ? null : remoteEvent.getTargetChanges().get(queryView.getTargetId());

firebase-firestore/src/main/java/com/google/firebase/firestore/local/IndexFreeQueryEngine.java

@@ -16,7 +16,6 @@
 
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
-import androidx.annotation.Nullable;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.core.Query;
@@ -58,7 +57,9 @@ public void setLocalDocumentsView(LocalDocumentsView localDocuments) {
 
   @Override
   public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
-      Query query, @Nullable QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys) {
+      Query query,
+      SnapshotVersion lastLimboFreeSnapshotVersion,
+      ImmutableSortedSet<DocumentKey> remoteKeys) {
     hardAssert(localDocumentsView != null, "setLocalDocumentsView() not called");
 
     // Queries that match all document don't benefit from using IndexFreeQueries. It is more
@@ -69,8 +70,7 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
 
     // Queries that have never seen a snapshot without limbo free documents should also be run as a
     // full collection scan.
-    if (queryData == null
-        || queryData.getLastLimboFreeSnapshotVersion().equals(SnapshotVersion.NONE)) {
+    if (lastLimboFreeSnapshotVersion.equals(SnapshotVersion.NONE)) {
       return executeFullCollectionScan(query);
     }
 
@@ -79,23 +79,22 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
     ImmutableSortedSet<Document> previousResults = applyQuery(query, documents);
 
     if (query.hasLimit()
-        && needsRefill(previousResults, remoteKeys, queryData.getLastLimboFreeSnapshotVersion())) {
+        && needsRefill(previousResults, remoteKeys, lastLimboFreeSnapshotVersion)) {
       return executeFullCollectionScan(query);
     }
 
     if (Logger.isDebugEnabled()) {
       Logger.debug(
           LOG_TAG,
           "Re-using previous result from %s to execute query: %s",
-          queryData.getLastLimboFreeSnapshotVersion().toString(),
+          lastLimboFreeSnapshotVersion.toString(),
           query.toString());
     }
 
     // Retrieve all results for documents that were updated since the last limbo-document free
     // remote snapshot.
     ImmutableSortedMap<DocumentKey, Document> updatedResults =
-        localDocumentsView.getDocumentsMatchingQuery(
-            query, queryData.getLastLimboFreeSnapshotVersion());
+        localDocumentsView.getDocumentsMatchingQuery(query, lastLimboFreeSnapshotVersion);
 
     // We merge `previousResults` into `updateResults`, since `updateResults` is already a
     // ImmutableSortedMap. If a document is contained in both lists, then its contents are the same.

firebase-firestore/src/main/java/com/google/firebase/firestore/local/IndexedQueryEngine.java

@@ -103,7 +103,9 @@ public void setLocalDocumentsView(LocalDocumentsView localDocuments) {
 
   @Override
   public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
-      Query query, @Nullable QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys) {
+      Query query,
+      SnapshotVersion lastLimboFreeSnapshotVersion,
+      ImmutableSortedSet<DocumentKey> remoteKeys) {
     hardAssert(localDocuments != null, "setLocalDocumentsView() not called");
 
     return query.isDocumentQuery()

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalStore.java

@@ -598,25 +598,29 @@ public void releaseQuery(Query query) {
         });
   }
 
-  /** Runs the given query against all the documents in the local store and returns the results. */
-  public ImmutableSortedMap<DocumentKey, Document> executeQuery(Query query) {
+  /**
+   * Runs the specified query against the local store and returns the results, potentially taking
+   * advantage of query data from previous executions (such as the set of remote keys).
+   *
+   * @param usePreviousResults Whether results from previous executions can be used to optimize this
+   *     query execution.
+   */
+  public QueryResult executeQuery(Query query, boolean usePreviousResults) {
     QueryData queryData = getQueryData(query);
+    SnapshotVersion lastLimboFreeSnapshotVersion = SnapshotVersion.NONE;
+    ImmutableSortedSet<DocumentKey> remoteKeys = DocumentKey.emptyKeySet();
+
     if (queryData != null) {
-      ImmutableSortedSet<DocumentKey> remoteKeys =
-          this.queryCache.getMatchingKeysForTargetId(queryData.getTargetId());
-      return executeQuery(query, queryData, remoteKeys);
-    } else {
-      return executeQuery(query, null, DocumentKey.emptyKeySet());
+      lastLimboFreeSnapshotVersion = queryData.getLastLimboFreeSnapshotVersion();
+      remoteKeys = this.queryCache.getMatchingKeysForTargetId(queryData.getTargetId());
     }
-  }
 
-  /**
-   * Runs the given query against the local store and returns the results, potentially taking
-   * advantage of the provided query data and the set of remote document keys.
-   */
-  public ImmutableSortedMap<DocumentKey, Document> executeQuery(
-      Query query, @Nullable QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys) {
-    return queryEngine.getDocumentsMatchingQuery(query, queryData, remoteKeys);
+    ImmutableSortedMap<DocumentKey, Document> documents =
+        queryEngine.getDocumentsMatchingQuery(
+            query,
+            usePreviousResults ? lastLimboFreeSnapshotVersion : SnapshotVersion.NONE,
+            usePreviousResults ? remoteKeys : DocumentKey.emptyKeySet());
+    return new QueryResult(documents, remoteKeys);
   }
 
   /**

firebase-firestore/src/main/java/com/google/firebase/firestore/local/QueryEngine.java

@@ -14,13 +14,13 @@
 
 package com.google.firebase.firestore.local;
 
-import androidx.annotation.Nullable;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.MaybeDocument;
+import com.google.firebase.firestore.model.SnapshotVersion;
 
 /**
  * Represents a query engine capable of performing queries over the local document cache. You must
@@ -33,7 +33,9 @@ public interface QueryEngine {
 
   /** Returns all local documents matching the specified query. */
   ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
-      Query query, @Nullable QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys);
+      Query query,
+      SnapshotVersion lastLimboFreeSnapshotVersion,
+      ImmutableSortedSet<DocumentKey> remoteKeys);
 
   /**
    * Notifies the query engine of a document change in case it would like to update indexes and the

firebase-firestore/src/main/java/com/google/firebase/firestore/local/QueryResult.java

@@ -0,0 +1,41 @@
+// Copyright 2019 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.firestore.local;
+
+import com.google.firebase.database.collection.ImmutableSortedMap;
+import com.google.firebase.database.collection.ImmutableSortedSet;
+import com.google.firebase.firestore.model.Document;
+import com.google.firebase.firestore.model.DocumentKey;
+
+/** The result of executing a query against the local store. */
+public class QueryResult {
+  private final ImmutableSortedMap<DocumentKey, Document> documents;
+  private final ImmutableSortedSet<DocumentKey> remoteKeys;
+
+  public QueryResult(
+      ImmutableSortedMap<DocumentKey, Document> documents,
+      ImmutableSortedSet<DocumentKey> remoteKeys) {
+    this.documents = documents;
+    this.remoteKeys = remoteKeys;
+  }
+
+  public ImmutableSortedMap<DocumentKey, Document> getDocuments() {
+    return documents;
+  }
+
+  public ImmutableSortedSet<DocumentKey> getRemoteKeys() {
+    return remoteKeys;
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SimpleQueryEngine.java

@@ -16,7 +16,6 @@
 
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
-import androidx.annotation.Nullable;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.core.Query;
@@ -40,7 +39,9 @@ public void setLocalDocumentsView(LocalDocumentsView localDocuments) {
 
   @Override
   public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
-      Query query, @Nullable QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys) {
+      Query query,
+      SnapshotVersion lastLimboFreeSnapshotVersion,
+      ImmutableSortedSet<DocumentKey> remoteKeys) {
     hardAssert(localDocumentsView != null, "setLocalDocumentsView() not called");
 
     // TODO: Once LocalDocumentsView provides a getCollectionDocuments() method, we

firebase-firestore/src/test/java/com/google/firebase/firestore/local/CountingQueryEngine.java

@@ -64,8 +64,10 @@ public void setLocalDocumentsView(LocalDocumentsView localDocuments) {
 
   @Override
   public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
-      Query query, @Nullable QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys) {
-    return queryEngine.getDocumentsMatchingQuery(query, queryData, remoteKeys);
+      Query query,
+      SnapshotVersion lastLimboFreeSnapshotVersion,
+      ImmutableSortedSet<DocumentKey> remoteKeys) {
+    return queryEngine.getDocumentsMatchingQuery(query, lastLimboFreeSnapshotVersion, remoteKeys);
   }
 
   @Override