Simplify segment count calculation

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

@@ -23,7 +23,6 @@
 import static java.lang.Math.max;
 
 import android.text.TextUtils;
-import android.util.Pair;
 import androidx.annotation.Nullable;
 import com.google.firebase.Timestamp;
 import com.google.firebase.database.collection.ImmutableSortedMap;
@@ -315,13 +314,12 @@ public IndexOffset getMinOffset(String collectionGroup) {
   @Override
   public IndexType getIndexType(Target target) {
     for (Target subTarget : getSubTargets(target)) {
-      Pair<FieldIndex, Integer> indexAndSegmentCount = getFieldIndexAndSegmentCount(subTarget);
-      if (indexAndSegmentCount == null) {
+      FieldIndex index = getFieldIndex(subTarget);
+      if (index == null) {
         return IndexType.NONE;
       }
 
-      Integer indexSegmentCount = indexAndSegmentCount.second;
-      if (indexSegmentCount < subTarget.getSegmentCount()) {
+      if (index.getSegments().size() < subTarget.getSegmentCount()) {
         return IndexType.PARTIAL;
       }
     }
@@ -332,9 +330,9 @@ public IndexType getIndexType(Target target) {
   public IndexOffset getMinOffset(Target target) {
     List<FieldIndex> fieldIndexes = new ArrayList<>();
     for (Target subTarget : getSubTargets(target)) {
-      Pair<FieldIndex, Integer> indexAndSegmentCount = getFieldIndexAndSegmentCount(subTarget);
-      if (indexAndSegmentCount != null) {
-        fieldIndexes.add(indexAndSegmentCount.first);
+      FieldIndex index = getFieldIndex(subTarget);
+      if (index != null) {
+        fieldIndexes.add(index);
       }
     }
     return getMinOffset(fieldIndexes);
@@ -468,12 +466,11 @@ public List<DocumentKey> getDocumentsMatchingTarget(Target target) {
     List<Object> bindings = new ArrayList<>();
 
     for (Target subTarget : getSubTargets(target)) {
-      Pair<FieldIndex, Integer> indexAndSegmentCount = getFieldIndexAndSegmentCount(subTarget);
-      if (indexAndSegmentCount == null) {
+      FieldIndex fieldIndex = getFieldIndex(subTarget);
+      if (fieldIndex == null) {
         return null;
       }
 
-      FieldIndex fieldIndex = indexAndSegmentCount.first;
       @Nullable List<Value> arrayValues = subTarget.getArrayValues(fieldIndex);
       @Nullable Collection<Value> notInValues = subTarget.getNotInValues(fieldIndex);
       @Nullable Bound lowerBound = subTarget.getLowerBound(fieldIndex);
@@ -627,12 +624,11 @@ private Object[] fillBounds(
   }
 
   /**
-   * Returns an index that can be used to serve the provided target, as well as the number of index
-   * segments that matched the target filters and orderBys. Returns {@code null} if no index is
-   * configured.
+   * Returns an index that can be used to serve the provided target. Returns {@code null} if no
+   * index is configured.
    */
   @Nullable
-  public Pair<FieldIndex, Integer> getFieldIndexAndSegmentCount(Target target) {
+  private FieldIndex getFieldIndex(Target target) {
     hardAssert(started, "IndexManager not started");
 
     TargetIndexMatcher targetIndexMatcher = new TargetIndexMatcher(target);
@@ -647,19 +643,16 @@ public Pair<FieldIndex, Integer> getFieldIndexAndSegmentCount(Target target) {
     }
 
     // Return the index with the most number of segments.
-    int maxNumMatchingSegments = -1;
     FieldIndex matchingIndex = null;
     for (FieldIndex fieldIndex : collectionIndexes) {
-      int numMatchingSegments = targetIndexMatcher.servedByIndex(fieldIndex);
-      if (numMatchingSegments > maxNumMatchingSegments) {
-        maxNumMatchingSegments = numMatchingSegments;
+      boolean matches = targetIndexMatcher.servedByIndex(fieldIndex);
+      if (matches
+          && (matchingIndex == null
+              || fieldIndex.getSegments().size() > matchingIndex.getSegments().size())) {
         matchingIndex = fieldIndex;
       }
     }
-    if (maxNumMatchingSegments == -1) {
-      return null;
-    }
-    return new Pair<>(matchingIndex, maxNumMatchingSegments);
+    return matchingIndex;
   }
 
   /**

firebase-firestore/src/main/java/com/google/firebase/firestore/model/TargetIndexMatcher.java

@@ -111,8 +111,7 @@ public TargetIndexMatcher(Target target) {
   }
 
   /**
-   * Returns the number of index segments that match the TargetIndexMatcher's target. Returns -1 if
-   * the index cannot be used to serve the target.
+   * Returns whether the index can be used to serve the TargetIndexMatcher's target.
    *
    * <p>An index is considered capable of serving the target when:
    *
@@ -132,26 +131,21 @@ public TargetIndexMatcher(Target target) {
    *
    * @throws AssertionError if the index is for a different collection
    */
-  public int servedByIndex(FieldIndex index) {
+  public boolean servedByIndex(FieldIndex index) {
     hardAssert(index.getCollectionGroup().equals(collectionId), "Collection IDs do not match");
-    int numMatchingSegments = 0;
 
     // If there is an array element, find a matching filter.
     FieldIndex.Segment arraySegment = index.getArraySegment();
-    if (arraySegment != null) {
-      if (hasMatchingEqualityFilter(arraySegment)) {
-        ++numMatchingSegments;
-      } else {
-        return -1;
-      }
+    if (arraySegment != null && !hasMatchingEqualityFilter(arraySegment)) {
+      return false;
     }
 
     Iterator<OrderBy> orderBys = this.orderBys.iterator();
     List<FieldIndex.Segment> segments = index.getDirectionalSegments();
     int segmentIndex = 0;
 
     // Process all equalities first. Equalities can appear out of order.
-    for (; segmentIndex < segments.size(); ++segmentIndex, ++numMatchingSegments) {
+    for (; segmentIndex < segments.size(); ++segmentIndex) {
       // We attempt to greedily match all segments to equality filters. If a filter matches an
       // index segment, we can mark the segment as used. Since it is not possible to use the same
       // field path in both an equality and inequality/oderBy clause, we do not have to consider the
@@ -168,29 +162,28 @@ public int servedByIndex(FieldIndex index) {
     // filters and we do not need to map any segments to the target's inequality and orderBy
     // clauses.
     if (segmentIndex == segments.size()) {
-      return numMatchingSegments;
+      return true;
     }
 
     // If there is an inequality filter, the next segment must match both the filter and the first
     // orderBy clause.
     if (inequalityFilter != null) {
       FieldIndex.Segment segment = segments.get(segmentIndex);
       if (!matchesFilter(inequalityFilter, segment) || !matchesOrderBy(orderBys.next(), segment)) {
-        return -1;
+        return false;
       }
       ++segmentIndex;
-      ++numMatchingSegments;
     }
 
     // All remaining segments need to represent the prefix of the target's orderBy.
-    for (; segmentIndex < segments.size(); ++segmentIndex, ++numMatchingSegments) {
+    for (; segmentIndex < segments.size(); ++segmentIndex) {
       FieldIndex.Segment segment = segments.get(segmentIndex);
       if (!orderBys.hasNext() || !matchesOrderBy(orderBys.next(), segment)) {
-        return -1;
+        return false;
       }
     }
 
-    return numMatchingSegments;
+    return true;
   }
 
   private boolean hasMatchingEqualityFilter(FieldIndex.Segment segment) {

firebase-firestore/src/test/java/com/google/firebase/firestore/model/TargetIndexMatcherTest.java

@@ -21,7 +21,7 @@
 import static com.google.firebase.firestore.testutil.TestUtil.orderBy;
 import static com.google.firebase.firestore.testutil.TestUtil.path;
 import static com.google.firebase.firestore.testutil.TestUtil.query;
-import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 import com.google.firebase.firestore.core.Query;
@@ -569,13 +569,13 @@ private void validateServesTarget(
       Query query, String field, FieldIndex.Segment.Kind kind, Object... fieldsAndKind) {
     FieldIndex expectedIndex = fieldIndex("collId", field, kind, fieldsAndKind);
     TargetIndexMatcher targetIndexMatcher = new TargetIndexMatcher(query.toTarget());
-    assertTrue(targetIndexMatcher.servedByIndex(expectedIndex) > 0);
+    assertTrue(targetIndexMatcher.servedByIndex(expectedIndex));
   }
 
   private void validateDoesNotServeTarget(
       Query query, String field, FieldIndex.Segment.Kind kind, Object... fieldsAndKind) {
     FieldIndex expectedIndex = fieldIndex("collId", field, kind, fieldsAndKind);
     TargetIndexMatcher targetIndexMatcher = new TargetIndexMatcher(query.toTarget());
-    assertEquals(targetIndexMatcher.servedByIndex(expectedIndex), -1);
+    assertFalse(targetIndexMatcher.servedByIndex(expectedIndex));
   }
 }