Re-run limit queries without limit if they need to be re-filled

This changes the query engine to also drop the limit on limit queries that need to be refilled. It also cleans up some of the Target.getLimit() logic that leads to repetitive code.

Author: schmidt-sebastian

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

@@ -14,6 +14,7 @@
 
 package com.google.firebase.firestore;
 
+import static com.google.firebase.firestore.core.Query.LimitType.LIMIT_TO_LAST;
 import static com.google.firebase.firestore.util.Assert.fail;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 import static com.google.firebase.firestore.util.Preconditions.checkNotNull;
@@ -1216,7 +1217,7 @@ private ListenerRegistration addSnapshotListenerInternal(
   }
 
   private void validateHasExplicitOrderByForLimitToLast() {
-    if (query.hasLimitToLast() && query.getExplicitOrderBy().isEmpty()) {
+    if (query.getLimitType().equals(LIMIT_TO_LAST) && query.getExplicitOrderBy().isEmpty()) {
       throw new IllegalStateException(
           "limitToLast() queries require specifying at least one orderBy() clause");
     }

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

@@ -147,34 +147,16 @@ public List<Filter> getFilters() {
     return filters;
   }
 
-  /**
-   * The maximum number of results to return. If there is no limit on the query, then this will
-   * cause an assertion failure.
-   */
-  public long getLimitToFirst() {
-    hardAssert(hasLimitToFirst(), "Called getLimitToFirst when no limit was set");
-    return limit;
-  }
-
-  public boolean hasLimitToFirst() {
-    return limitType == LimitType.LIMIT_TO_FIRST && limit != Target.NO_LIMIT;
-  }
-
-  /**
-   * The maximum number of last-matching results to return. If there is no limit on the query, then
-   * this will cause an assertion failure.
-   */
-  public long getLimitToLast() {
-    hardAssert(hasLimitToLast(), "Called getLimitToLast when no limit was set");
+  /** The maximum number of results to return or {@link Target#NO_LIMIT} if there is no limit. */
+  public long getLimit() {
     return limit;
   }
 
-  public boolean hasLimitToLast() {
-    return limitType == LimitType.LIMIT_TO_LAST && limit != Target.NO_LIMIT;
+  public boolean hasLimit() {
+    return limit != Target.NO_LIMIT;
   }
 
   public LimitType getLimitType() {
-    hardAssert(hasLimitToLast() || hasLimitToFirst(), "Called getLimitType when no limit was set");
     return limitType;
   }
 

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

@@ -14,6 +14,8 @@
 
 package com.google.firebase.firestore.core;
 
+import static com.google.firebase.firestore.core.Query.LimitType.LIMIT_TO_FIRST;
+import static com.google.firebase.firestore.core.Query.LimitType.LIMIT_TO_LAST;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 import static com.google.firebase.firestore.util.Util.compareIntegers;
 
@@ -147,11 +149,11 @@ public DocumentChanges computeDocChanges(
     // Note that this should never get used in a refill (when previousChanges is set), because there
     // will only be adds -- no deletes or updates.
     Document lastDocInLimit =
-        (query.hasLimitToFirst() && oldDocumentSet.size() == query.getLimitToFirst())
+        (query.getLimitType().equals(LIMIT_TO_FIRST) && oldDocumentSet.size() == query.getLimit())
             ? oldDocumentSet.getLastDocument()
             : null;
     Document firstDocInLimit =
-        (query.hasLimitToLast() && oldDocumentSet.size() == query.getLimitToLast())
+        (query.getLimitType().equals(LIMIT_TO_LAST) && oldDocumentSet.size() == query.getLimit())
             ? oldDocumentSet.getFirstDocument()
             : null;
 
@@ -222,11 +224,10 @@ public DocumentChanges computeDocChanges(
     }
 
     // Drop documents out to meet limitToFirst/limitToLast requirement.
-    if (query.hasLimitToFirst() || query.hasLimitToLast()) {
-      long limit = query.hasLimitToFirst() ? query.getLimitToFirst() : query.getLimitToLast();
-      for (long i = newDocumentSet.size() - limit; i > 0; --i) {
+    if (query.hasLimit()) {
+      for (long i = newDocumentSet.size() - query.getLimit(); i > 0; --i) {
         Document oldDoc =
-            query.hasLimitToFirst()
+            query.getLimitType().equals(LIMIT_TO_FIRST)
                 ? newDocumentSet.getLastDocument()
                 : newDocumentSet.getFirstDocument();
         newDocumentSet = newDocumentSet.remove(oldDoc.getKey());

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

@@ -117,15 +117,14 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       return null;
     }
 
-    if (indexType.equals(IndexType.PARTIAL)) {
+    if (query.hasLimit() && indexType.equals(IndexType.PARTIAL)) {
       // We cannot apply a limit for targets that are served using a partial index.
       // If a partial index will be used to serve the target, the query may return a superset of
       // documents that match the target (e.g. if the index doesn't include all the target's
       // filters), or may return the correct set of documents in the wrong order (e.g. if the index
       // doesn't include a segment for one of the orderBys). Therefore a limit should not be applied
       // in such cases.
-      query = query.limitToFirst(Target.NO_LIMIT);
-      target = query.toTarget();
+      return performQueryUsingIndex(query.limitToFirst(Target.NO_LIMIT));
     }
 
     List<DocumentKey> keys = indexManager.getDocumentsMatchingTarget(target);
@@ -136,9 +135,12 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
     IndexOffset offset = indexManager.getMinOffset(target);
 
     ImmutableSortedSet<Document> previousResults = applyQuery(query, indexedDocuments);
-    if ((query.hasLimitToFirst() || query.hasLimitToLast())
-        && needsRefill(query.getLimitType(), keys.size(), previousResults, offset.getReadTime())) {
-      return null;
+    if (needsRefill(query, keys.size(), previousResults, offset.getReadTime())) {
+      // A limit query whose boundaries change due to local edits can be re-run against the cache
+      // by excluding the limit. This ensures that all documents that match the query's filters are
+      // included in the result set. The SDK can then apply the limit once all local edits are
+      // incorporated.
+      return performQueryUsingIndex(query.limitToFirst(Target.NO_LIMIT));
     }
 
     return appendRemainingResults(values(indexedDocuments), query, offset);
@@ -167,12 +169,7 @@ && needsRefill(query.getLimitType(), keys.size(), previousResults, offset.getRea
         localDocumentsView.getDocuments(remoteKeys);
     ImmutableSortedSet<Document> previousResults = applyQuery(query, documents);
 
-    if ((query.hasLimitToFirst() || query.hasLimitToLast())
-        && needsRefill(
-            query.getLimitType(),
-            remoteKeys.size(),
-            previousResults,
-            lastLimboFreeSnapshotVersion)) {
+    if (needsRefill(query, remoteKeys.size(), previousResults, lastLimboFreeSnapshotVersion)) {
       return null;
     }
 
@@ -211,7 +208,7 @@ private ImmutableSortedSet<Document> applyQuery(
    * Determines if a limit query needs to be refilled from cache, making it ineligible for
    * index-free execution.
    *
-   * @param limitType The type of limit query for refill calculation.
+   * @param query The query.
    * @param expectedDocumentCount The number of documents keys that matched the query at the last
    *     snapshot.
    * @param sortedPreviousResults The documents that match the query based on the previous result,
@@ -221,12 +218,17 @@ private ImmutableSortedSet<Document> applyQuery(
    *     synchronized.
    */
   private boolean needsRefill(
-      Query.LimitType limitType,
+      Query query,
       int expectedDocumentCount,
       ImmutableSortedSet<Document> sortedPreviousResults,
       SnapshotVersion limboFreeSnapshotVersion) {
-    // The query needs to be refilled if a previously matching document no longer matches.
+    if (!query.hasLimit()) {
+      // Queries without limits do not need to be refilled.
+      return false;
+    }
+
     if (expectedDocumentCount != sortedPreviousResults.size()) {
+      // The query needs to be refilled if a previously matching document no longer matches.
       return true;
     }
 
@@ -237,7 +239,7 @@ private boolean needsRefill(
     // did not change and documents from cache will continue to be "rejected" by this boundary.
     // Therefore, we can ignore any modifications that don't affect the last document.
     Document documentAtLimitEdge =
-        limitType == Query.LimitType.LIMIT_TO_FIRST
+        query.getLimitType() == Query.LimitType.LIMIT_TO_FIRST
             ? sortedPreviousResults.getMaxEntry()
             : sortedPreviousResults.getMinEntry();
     if (documentAtLimitEdge == null) {

firebase-firestore/src/test/java/com/google/firebase/firestore/bundle/BundleSerializerTest.java

@@ -712,30 +712,26 @@ private void assertDecodesNamedQuery(String json, Query query) throws JSONExcept
             + json
             + ",\n"
             + "    limitType: '"
-            + (query.hasLimitToLast() ? "LAST" : "FIRST")
+            + (query.getLimitType().equals(Query.LimitType.LIMIT_TO_LAST) ? "LAST" : "FIRST")
             + "'\n"
             + "   },\n"
             + " readTime: '2020-01-01T00:00:01.000000001Z'\n"
             + "}";
     NamedQuery actualNamedQuery = serializer.decodeNamedQuery(new JSONObject(queryJson));
 
-    long limit =
-        query.hasLimitToFirst()
-            ? query.getLimitToFirst()
-            : (query.hasLimitToLast() ? query.getLimitToLast() : Target.NO_LIMIT);
     Target target =
         new Target(
             query.getPath(),
             query.getCollectionGroup(),
             query.getFilters(),
             query.getExplicitOrderBy(),
-            limit,
+            query.getLimit(),
             query.getStartAt(),
             query.getEndAt());
     BundledQuery bundledQuery =
         new BundledQuery(
             target,
-            query.hasLimitToLast()
+            query.getLimitType().equals(Query.LimitType.LIMIT_TO_LAST)
                 ? Query.LimitType.LIMIT_TO_LAST
                 : Query.LimitType.LIMIT_TO_FIRST);
     NamedQuery expectedNamedQuery =