Update IndexFreeQueryEngine with changes from Web

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

@@ -28,27 +28,22 @@
 import java.util.Collections;
 import java.util.Map;
 
+// TOOD(b/140938512): Drop SimpleQueryEngine and rename IndexFreeQueryEngine.
+
 /**
  * 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.
+ * IndexFreeQueryEngine optimizes query execution by only reading the documents that 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:
+ * <p>There are some cases where Index-Free queries are not guaranteed to produce the same results
+ * as full collection scans. In these cases, the IndexFreeQueryEngine falls back to 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.
+ *   <il>Limit queries where a document that matched the query previously no longer matches the
+ *   query. <il> Limit queries where a document edit may cause the document to sort below another
+ *   document that is in the local cache. <il>Queries that have never been CURRENT or free of Limbo
+ *   documents.
  * </ol>
  */
 public class IndexFreeQueryEngine implements QueryEngine {
@@ -79,7 +74,9 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       return executeFullCollectionScan(query);
     }
 
-    ImmutableSortedSet<Document> previousResults = getSortedPreviousResults(query, remoteKeys);
+    ImmutableSortedMap<DocumentKey, MaybeDocument> documents =
+        localDocumentsView.getDocuments(remoteKeys);
+    ImmutableSortedSet<Document> previousResults = applyQuery(query, documents);
 
     if (query.hasLimit()
         && needsRefill(previousResults, remoteKeys, queryData.getLastLimboFreeSnapshotVersion())) {
@@ -99,35 +96,31 @@ && needsRefill(previousResults, remoteKeys, queryData.getLastLimboFreeSnapshotVe
     ImmutableSortedMap<DocumentKey, Document> updatedResults =
         localDocumentsView.getDocumentsMatchingQuery(
             query, queryData.getLastLimboFreeSnapshotVersion());
+
+    // 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.
     for (Document result : previousResults) {
       updatedResults = updatedResults.insert(result.getKey(), result);
     }
 
     return updatedResults;
   }
 
-  /**
-   * Returns the documents for the specified remote keys if they still match the query, sorted by
-   * the query's comparator.
-   */
-  private ImmutableSortedSet<Document> getSortedPreviousResults(
-      Query query, ImmutableSortedSet<DocumentKey> remoteKeys) {
-    // Fetch the documents that matched the query at the last snapshot.
-    ImmutableSortedMap<DocumentKey, MaybeDocument> previousResults =
-        localDocumentsView.getDocuments(remoteKeys);
-
+  /** Applies the query filter and sorting to the provided documents. */
+  private ImmutableSortedSet<Document> applyQuery(
+      Query query, ImmutableSortedMap<DocumentKey, MaybeDocument> documents) {
     // Sort the documents and re-apply the query filter since previously matching documents do not
     // necessarily still match the query.
-    ImmutableSortedSet<Document> results =
+    ImmutableSortedSet<Document> queryResults =
         new ImmutableSortedSet<>(Collections.emptyList(), query.comparator());
-    for (Map.Entry<DocumentKey, MaybeDocument> entry : previousResults) {
+    for (Map.Entry<DocumentKey, MaybeDocument> entry : documents) {
       MaybeDocument maybeDoc = entry.getValue();
       if (maybeDoc instanceof Document && query.matches((Document) maybeDoc)) {
         Document doc = (Document) maybeDoc;
-        results = results.insert(doc);
+        queryResults = queryResults.insert(doc);
       }
     }
-    return results;
+    return queryResults;
   }
 
   /**
@@ -149,11 +142,6 @@ private boolean needsRefill(
       return true;
     }
 
-    // We don't need to find a better match from cache if no documents matched the query.
-    if (sortedPreviousResults.isEmpty()) {
-      return false;
-    }
-
     // Limit queries are not eligible for index-free query execution if there is a potential that an
     // older document from cache now sorts before a document that was previously part of the limit.
     // This, however, can only happen if the last document of the limit sorts lower than it did when
@@ -162,6 +150,10 @@ private boolean needsRefill(
     // continue to be "rejected" by this boundary. Therefore, we can ignore any modifications that
     // don't affect the last document.
     Document lastDocumentInLimit = sortedPreviousResults.getMaxEntry();
+    if (lastDocumentInLimit == null) {
+      // We don't need to refill the query if there were already no documents.
+      return false;
+    }
     return lastDocumentInLimit.hasPendingWrites()
         || lastDocumentInLimit.getVersion().compareTo(limboFreeSnapshotVersion) > 0;
   }