Add Index-Free query engine

firebase-database-collection/src/main/java/com/google/firebase/database/collection/ImmutableSortedMap.java

@@ -56,6 +56,14 @@ public abstract class ImmutableSortedMap<K, V> implements Iterable<Map.Entry<K,
 
   public abstract Comparator<K> getComparator();
 
+  public ImmutableSortedMap<K, V> insertAll(ImmutableSortedMap<K, V> source) {
+    ImmutableSortedMap<K, V> result = this;
+    for (Map.Entry<K, V> entry : source) {
+      result = result.insert(entry.getKey(), entry.getValue());
+    }
+    return result;
+  }
+
   @Override
   @SuppressWarnings("unchecked")
   public boolean equals(Object o) {

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/remote/RemoteStoreTest.java

@@ -24,6 +24,7 @@
 import com.google.firebase.firestore.local.LocalStore;
 import com.google.firebase.firestore.local.MemoryPersistence;
 import com.google.firebase.firestore.local.Persistence;
+import com.google.firebase.firestore.local.SimpleQueryEngine;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.mutation.MutationBatchResult;
 import com.google.firebase.firestore.testutil.IntegrationTestUtil;
@@ -72,9 +73,10 @@ public ImmutableSortedSet<DocumentKey> getRemoteKeysForTarget(int targetId) {
         };
 
     FakeConnectivityMonitor connectivityMonitor = new FakeConnectivityMonitor();
+    SimpleQueryEngine queryEngine = new SimpleQueryEngine();
     Persistence persistence = MemoryPersistence.createEagerGcMemoryPersistence();
     persistence.start();
-    LocalStore localStore = new LocalStore(persistence, User.UNAUTHENTICATED);
+    LocalStore localStore = new LocalStore(persistence, queryEngine, User.UNAUTHENTICATED);
     RemoteStore remoteStore =
         new RemoteStore(callback, localStore, datastore, testQueue, connectivityMonitor);
 

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

@@ -37,7 +37,9 @@
 import com.google.firebase.firestore.local.LruGarbageCollector;
 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.SQLitePersistence;
+import com.google.firebase.firestore.local.SimpleQueryEngine;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.MaybeDocument;
@@ -263,7 +265,9 @@ private void initialize(Context context, User user, boolean usePersistence, long
     }
 
     persistence.start();
-    localStore = new LocalStore(persistence, user);
+    // TODO(index-free): Use IndexFreeQueryEngine/IndexedQueryEngine as appropriate.
+    QueryEngine queryEngine = new SimpleQueryEngine();
+    localStore = new LocalStore(persistence, queryEngine, user);
     if (gc != null) {
       lruScheduler = gc.newScheduler(asyncQueue, localStore);
       lruScheduler.start();

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

@@ -117,6 +117,13 @@ public boolean isCollectionGroupQuery() {
     return collectionGroup != null;
   }
 
+  /**
+   * Returns true if this query does not specify any query constraints that could remove results.
+   */
+  public boolean matchesAllDocuments() {
+    return filters.isEmpty() && limit == NO_LIMIT && startAt == null && endAt == null;
+  }
+
   /** The filters on the documents returned by the query. */
   public List<Filter> getFilters() {
     return filters;

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

@@ -193,9 +193,10 @@ public int listen(Query query) {
   private ViewSnapshot initializeViewAndComputeSnapshot(QueryData queryData) {
     Query query = queryData.getQuery();
 
-    ImmutableSortedMap<DocumentKey, Document> docs = localStore.executeQuery(query);
     ImmutableSortedSet<DocumentKey> remoteKeys =
         localStore.getRemoteDocumentKeys(queryData.getTargetId());
+    ImmutableSortedMap<DocumentKey, Document> docs =
+        localStore.executeQuery(query, queryData, remoteKeys);
 
     View view = new View(query, remoteKeys);
     View.DocumentChanges viewDocChanges = view.computeDocChanges(docs);
@@ -557,7 +558,8 @@ private void emitNewSnapsAndNotifyLocalStore(
         // 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());
+            localStore.executeQuery(
+                queryView.getQuery(), /* queryData= */ null, DocumentKey.emptyKeySet());
         viewDocChanges = view.computeDocChanges(docs, viewDocChanges);
       }
       TargetChange targetChange =

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

@@ -0,0 +1,150 @@
+// 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 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;
+import com.google.firebase.firestore.model.Document;
+import com.google.firebase.firestore.model.DocumentCollections;
+import com.google.firebase.firestore.model.DocumentKey;
+import com.google.firebase.firestore.model.MaybeDocument;
+import com.google.firebase.firestore.model.SnapshotVersion;
+import java.util.Map;
+
+/**
+ * A query engine that takes advantage of the target document mapping in the QueryCache. The
+ * IndexFreeQueryEngine optimizes query execution by only reading the documents previously matched a
+ * query plus any documents that were edited after the query was last listened to.
+ *
+ * <p>There are some cases where Index-Free queries are not guaranteed to produce to the same
+ * results as full collection scans. In these case, the IndexFreeQueryEngine falls back to a full
+ * query processing. These cases are:
+ *
+ * <ol>
+ *   <li>Limit queries where a document that matched the query previously no longer matches the
+ *       query. In this case, we have to scan all local documents since a document that was sent to
+ *       us as part of a different query result may now fall into the limit.
+ *   <li>Limit queries that include edits that occurred after the last remote snapshot (both
+ *       latency-compensated or committed). Even if an edited document continues to match the query,
+ *       an edit may cause a document to sort below another document that is in the local cache.
+ *   <li>Queries where the last snapshot contained Limbo documents. Even though a Limbo document is
+ *       not part of the backend result set, we need to include Limbo documents in local views to
+ *       ensure consistency between different Query views. If there exists a previous query snapshot
+ *       that contained no limbo documents, we can instead use the older snapshot version for
+ *       Index-Free processing.
+ * </ol>
+ */
+public class IndexFreeQueryEngine implements QueryEngine {
+  private LocalDocumentsView localDocumentsView;
+
+  @Override
+  public void setLocalDocumentsView(LocalDocumentsView localDocuments) {
+    this.localDocumentsView = localDocuments;
+  }
+
+  @Override
+  public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
+      Query query, @Nullable QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys) {
+    hardAssert(localDocumentsView != null, "setLocalDocumentsView() not called");
+
+    // Queries that match all document don't benefit from using IndexFreeQueries. It is more
+    // efficient to scan all documents in a collection, rather than to perform individual lookups.
+    if (query.matchesAllDocuments()) {
+      return executeFullCollectionScan(query);
+    }
+
+    // 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)) {
+      return executeFullCollectionScan(query);
+    }
+
+    ImmutableSortedMap<DocumentKey, Document> result =
+        executeIndexFreeQuery(query, queryData, remoteKeys);
+
+    return result != null ? result : executeFullCollectionScan(query);
+  }
+
+  /**
+   * Attempts index-free query execution. Returns the set of query results on success, otherwise
+   * returns null.
+   */
+  private @Nullable ImmutableSortedMap<DocumentKey, Document> executeIndexFreeQuery(
+      Query query, QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys) {
+    // Fetch the documents that matched the query at the last snapshot.
+    ImmutableSortedMap<DocumentKey, MaybeDocument> previousResults =
+        localDocumentsView.getDocuments(remoteKeys);
+
+    // Limit queries are not eligible for index-free query execution if any part of the result was
+    // modified after we received the last query snapshot. This makes sure that we re-populate the
+    // view with older documents that may sort before the modified document.
+    if (query.hasLimit()
+        && containsUpdatesSinceSnapshotVersion(previousResults, queryData.getSnapshotVersion())) {
+      return null;
+    }
+
+    ImmutableSortedMap<DocumentKey, Document> results = DocumentCollections.emptyDocumentMap();
+
+    // Re-apply the query filter since previously matching documents do not necessarily still
+    // match the query.
+    for (Map.Entry<DocumentKey, MaybeDocument> entry : previousResults) {
+      MaybeDocument maybeDoc = entry.getValue();
+      if (maybeDoc instanceof Document && query.matches((Document) maybeDoc)) {
+        Document doc = (Document) maybeDoc;
+        results = results.insert(entry.getKey(), doc);
+      } else if (query.hasLimit()) {
+        // Limit queries with documents that no longer match need to be re-filled from cache.
+        return null;
+      }
+    }
+
+    // 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());
+
+    results = results.insertAll(updatedResults);
+
+    return results;
+  }
+
+  @Override
+  public void handleDocumentChange(MaybeDocument oldDocument, MaybeDocument newDocument) {
+    // No indexes to update.
+  }
+
+  private boolean containsUpdatesSinceSnapshotVersion(
+      ImmutableSortedMap<DocumentKey, MaybeDocument> previousResults,
+      SnapshotVersion sinceSnapshotVersion) {
+    for (Map.Entry<DocumentKey, MaybeDocument> doc : previousResults) {
+      if (doc.getValue().hasPendingWrites()
+          || doc.getValue().getVersion().compareTo(sinceSnapshotVersion) > 0) {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
+  private ImmutableSortedMap<DocumentKey, Document> executeFullCollectionScan(Query query) {
+    return localDocumentsView.getDocumentsMatchingQuery(query, SnapshotVersion.NONE);
+  }
+}

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

@@ -19,6 +19,7 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.firebase.database.collection.ImmutableSortedMap;
+import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.core.FieldFilter;
 import com.google.firebase.firestore.core.Filter;
 import com.google.firebase.firestore.core.Filter.Operator;
@@ -88,17 +89,23 @@ public class IndexedQueryEngine implements QueryEngine {
   private static final List<Class> lowCardinalityTypes =
       Arrays.asList(BooleanValue.class, ArrayValue.class, ObjectValue.class);
 
-  private final LocalDocumentsView localDocuments;
   private final SQLiteCollectionIndex collectionIndex;
+  private LocalDocumentsView localDocuments;
 
-  public IndexedQueryEngine(
-      LocalDocumentsView localDocuments, SQLiteCollectionIndex collectionIndex) {
-    this.localDocuments = localDocuments;
+  public IndexedQueryEngine(SQLiteCollectionIndex collectionIndex) {
     this.collectionIndex = collectionIndex;
   }
 
   @Override
-  public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(Query query) {
+  public void setLocalDocumentsView(LocalDocumentsView localDocuments) {
+    this.localDocuments = localDocuments;
+  }
+
+  @Override
+  public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
+      Query query, @Nullable QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys) {
+    hardAssert(localDocuments != null, "setLocalDocumentsView() not called");
+
     return query.isDocumentQuery()
         ? localDocuments.getDocumentsMatchingQuery(query, SnapshotVersion.NONE)
         : performCollectionQuery(query);

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

@@ -40,7 +40,7 @@
  * mutations in the MutationQueue to the RemoteDocumentCache.
  */
 // TODO: Turn this into the UnifiedDocumentCache / whatever.
-final class LocalDocumentsView {
+class LocalDocumentsView {
 
   private final RemoteDocumentCache remoteDocumentCache;
   private final MutationQueue mutationQueue;

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

@@ -19,6 +19,7 @@
 
 import android.util.SparseArray;
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 import com.google.firebase.Timestamp;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.collection.ImmutableSortedSet;
@@ -125,7 +126,7 @@ public final class LocalStore {
   /** Used to generate targetIds for queries tracked locally. */
   private final TargetIdGenerator targetIdGenerator;
 
-  public LocalStore(Persistence persistence, User initialUser) {
+  public LocalStore(Persistence persistence, QueryEngine queryEngine, User initialUser) {
     hardAssert(
         persistence.isStarted(), "LocalStore was passed an unstarted persistence implementation");
     this.persistence = persistence;
@@ -135,8 +136,9 @@ public LocalStore(Persistence persistence, User initialUser) {
     remoteDocuments = persistence.getRemoteDocumentCache();
     localDocuments =
         new LocalDocumentsView(remoteDocuments, mutationQueue, persistence.getIndexManager());
-    // TODO: Use IndexedQueryEngine as appropriate.
-    queryEngine = new SimpleQueryEngine(localDocuments);
+
+    this.queryEngine = queryEngine;
+    queryEngine.setLocalDocumentsView(localDocuments);
 
     localViewReferences = new ReferenceSet();
     persistence.getReferenceDelegate().setInMemoryPins(localViewReferences);
@@ -170,8 +172,7 @@ public ImmutableSortedMap<DocumentKey, MaybeDocument> handleUserChange(User user
     // Recreate our LocalDocumentsView using the new MutationQueue.
     localDocuments =
         new LocalDocumentsView(remoteDocuments, mutationQueue, persistence.getIndexManager());
-    // TODO: Use IndexedQueryEngine as appropriate.
-    queryEngine = new SimpleQueryEngine(localDocuments);
+    queryEngine.setLocalDocumentsView(localDocuments);
 
     // Union the old/new changed keys.
     ImmutableSortedSet<DocumentKey> changedKeys = DocumentKey.emptyKeySet();
@@ -561,6 +562,21 @@ public QueryData allocateQuery(Query query) {
     return cached;
   }
 
+  /**
+   * Returns the QueryData as seen by the LocalStore, including updates that may have not yet been
+   * persisted to the QueryCache.
+   */
+  @VisibleForTesting
+  @Nullable
+  QueryData getQueryData(Query query) {
+    QueryData queryData = queryCache.getQueryData(query);
+    if (queryData == null) {
+      return null;
+    }
+    QueryData updatedQueryData = targetIds.get(queryData.getTargetId());
+    return updatedQueryData != null ? updatedQueryData : queryData;
+  }
+
   /** Mutable state for the transaction in allocateQuery. */
   private static class AllocateQueryHolder {
     QueryData cached;
@@ -611,7 +627,23 @@ 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) {
-    return queryEngine.getDocumentsMatchingQuery(query);
+    QueryData queryData = getQueryData(query);
+    if (queryData != null) {
+      ImmutableSortedSet<DocumentKey> remoteKeys =
+          this.queryCache.getMatchingKeysForTargetId(queryData.getTargetId());
+      return executeQuery(query, queryData, remoteKeys);
+    } else {
+      return executeQuery(query, null, DocumentKey.emptyKeySet());
+    }
+  }
+
+  /**
+   * 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);
   }
 
   /**