Reduce polymorphism of DocIdSetIterator#docID(). (#14453)

This introduces a new `AbstractDocIdSetIterator` abstract class that tracks the
current doc IDs and refactors several `DocIdSetIterator` implementations to
extend either this `AbstractDocIdSetIterator` or `FilterDocIdSetIterator`, in
order to reduce polymorphism of call sites of `DocIdSetIterator#docID()`.

Author: jpountz

lucene/core/src/java/org/apache/lucene/codecs/lucene90/IndexedDISI.java

@@ -19,6 +19,7 @@
 import java.io.DataInput;
 import java.io.IOException;
 import org.apache.lucene.index.KnnVectorValues;
+import org.apache.lucene.search.AbstractDocIdSetIterator;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.IndexOutput;
@@ -91,7 +92,7 @@
  *
  * @lucene.internal
  */
-public final class IndexedDISI extends DocIdSetIterator {
+public final class IndexedDISI extends AbstractDocIdSetIterator {
 
   // jump-table time/space trade-offs to consider:
   // The block offsets and the block indexes could be stored in more compressed form with
@@ -422,7 +423,6 @@ public static RandomAccessInput createJumpTable(
   int nextBlockIndex = -1;
   Method method;
 
-  int doc = -1;
   int index = -1;
 
   // SPARSE variables
@@ -474,11 +474,6 @@ public long cost() {
     };
   }
 
-  @Override
-  public int docID() {
-    return doc;
-  }
-
   @Override
   public int advance(int target) throws IOException {
     final int targetBlock = target & 0xFFFF0000;

lucene/core/src/java/org/apache/lucene/search/AbstractDocIdSetIterator.java

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.lucene.search;
+
+/**
+ * Abstract implementation of a {@link DocIdSetIterator} that tracks the current doc ID.
+ * Implementing {@link DocIdSetIterator} by extending either this class or {@link
+ * FilterDocIdSetIterator} is a good idea in order to reduce polymorphism of call sites to {@link
+ * DocIdSetIterator#docID()}.
+ */
+public abstract class AbstractDocIdSetIterator extends DocIdSetIterator {
+
+  /** The current doc ID, initialized at -1. */
+  protected int doc = -1;
+
+  /** Sole constructor, invoked by sub-classes. */
+  protected AbstractDocIdSetIterator() {}
+
+  @Override
+  public final int docID() {
+    return doc;
+  }
+}

lucene/core/src/java/org/apache/lucene/search/BlockMaxConjunctionScorer.java

@@ -97,21 +97,11 @@ public DocIdSetIterator iterator() {
   private DocIdSetIterator approximation() {
     final DocIdSetIterator lead = approximations[0];
 
-    return new DocIdSetIterator() {
+    return new FilterDocIdSetIterator(lead) {
 
       float maxScore;
       int upTo = -1;
 
-      @Override
-      public int docID() {
-        return lead.docID();
-      }
-
-      @Override
-      public long cost() {
-        return lead.cost();
-      }
-
       private void moveToNextBlock(int target) throws IOException {
         if (minScore == 0) {
           upTo = target;

lucene/core/src/java/org/apache/lucene/search/ConjunctionDISI.java

@@ -33,7 +33,7 @@
  *
  * @lucene.internal
  */
-final class ConjunctionDISI extends DocIdSetIterator {
+final class ConjunctionDISI extends FilterDocIdSetIterator {
 
   /**
    * Adds the scorer, possibly splitting up into two phases or collapsing if it is another
@@ -127,6 +127,8 @@ static DocIdSetIterator createConjunction(
     if (iterators.size() == 1) {
       disi = iterators.get(0);
     } else {
+      // Sort the list first to allow the sparser iterator to lead the matching.
+      CollectionUtil.timSort(iterators, (o1, o2) -> Long.compare(o1.cost(), o2.cost()));
       disi = new ConjunctionDISI(iterators);
     }
 
@@ -152,11 +154,9 @@ private static void throwSubIteratorsNotOnSameDocument() {
   final DocIdSetIterator[] others;
 
   private ConjunctionDISI(List<? extends DocIdSetIterator> iterators) {
+    super(iterators.get(0));
     assert iterators.size() >= 2;
 
-    // Sort the array the first time to allow the least frequent DocsEnum to
-    // lead the matching.
-    CollectionUtil.timSort(iterators, (o1, o2) -> Long.compare(o1.cost(), o2.cost()));
     lead1 = iterators.get(0);
     lead2 = iterators.get(1);
     others = iterators.subList(2, iterators.size()).toArray(new DocIdSetIterator[0]);
@@ -205,23 +205,13 @@ assert assertItersOnSameDoc()
     return doNext(lead1.advance(target));
   }
 
-  @Override
-  public int docID() {
-    return lead1.docID();
-  }
-
   @Override
   public int nextDoc() throws IOException {
     assert assertItersOnSameDoc()
         : "Sub-iterators of ConjunctionDISI are not on the same document!";
     return doNext(lead1.nextDoc());
   }
 
-  @Override
-  public long cost() {
-    return lead1.cost(); // overestimate
-  }
-
   // Returns {@code true} if all sub-iterators are on the same doc ID, {@code false} otherwise
   private boolean assertItersOnSameDoc() {
     int curDoc = lead1.docID();
@@ -233,14 +223,15 @@ private boolean assertItersOnSameDoc() {
   }
 
   /** Conjunction between a {@link DocIdSetIterator} and one or more {@link BitSetIterator}s. */
-  private static class BitSetConjunctionDISI extends DocIdSetIterator {
+  private static class BitSetConjunctionDISI extends FilterDocIdSetIterator {
 
     private final DocIdSetIterator lead;
     private final BitSetIterator[] bitSetIterators;
     private final BitSet[] bitSets;
     private final int minLength;
 
     BitSetConjunctionDISI(DocIdSetIterator lead, Collection<BitSetIterator> bitSetIterators) {
+      super(lead);
       this.lead = lead;
       assert bitSetIterators.size() > 0;
 
@@ -257,11 +248,6 @@ private static class BitSetConjunctionDISI extends DocIdSetIterator {
       this.minLength = minLen;
     }
 
-    @Override
-    public int docID() {
-      return lead.docID();
-    }
-
     @Override
     public int nextDoc() throws IOException {
       assert assertItersOnSameDoc()
@@ -297,11 +283,6 @@ private int doNext(int doc) throws IOException {
       }
     }
 
-    @Override
-    public long cost() {
-      return lead.cost();
-    }
-
     // Returns {@code true} if all sub-iterators are on the same doc ID, {@code false} otherwise
     private boolean assertItersOnSameDoc() {
       int curDoc = lead.docID();

lucene/core/src/java/org/apache/lucene/search/DisjunctionDISIApproximation.java

@@ -29,7 +29,7 @@
  *
  * @lucene.internal
  */
-public final class DisjunctionDISIApproximation extends DocIdSetIterator {
+public final class DisjunctionDISIApproximation extends AbstractDocIdSetIterator {
 
   public static DisjunctionDISIApproximation of(
       Collection<? extends DisiWrapper> subIterators, long leadCost) {
@@ -111,11 +111,6 @@ public long cost() {
     return cost;
   }
 
-  @Override
-  public int docID() {
-    return Math.min(minOtherDoc, leadTop.doc);
-  }
-
   @Override
   public int nextDoc() throws IOException {
     if (leadTop.doc < minOtherDoc) {
@@ -124,7 +119,7 @@ public int nextDoc() throws IOException {
         leadTop.doc = leadTop.approximation.nextDoc();
         leadTop = leadIterators.updateTop();
       } while (leadTop.doc == curDoc);
-      return Math.min(leadTop.doc, minOtherDoc);
+      return doc = Math.min(leadTop.doc, minOtherDoc);
     } else {
       return advance(minOtherDoc + 1);
     }
@@ -145,7 +140,7 @@ public int advance(int target) throws IOException {
       minOtherDoc = Math.min(minOtherDoc, w.doc);
     }
 
-    return Math.min(leadTop.doc, minOtherDoc);
+    return doc = Math.min(leadTop.doc, minOtherDoc);
   }
 
   @Override
@@ -162,6 +157,8 @@ public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOExcept
       w.doc = w.approximation.docID();
       minOtherDoc = Math.min(minOtherDoc, w.doc);
     }
+
+    doc = Math.min(leadTop.doc, minOtherDoc);
   }
 
   /** Return the linked list of iterators positioned on the current doc. */

lucene/core/src/java/org/apache/lucene/search/DocIdSetIterator.java

@@ -55,10 +55,9 @@ public static DocIdSetIterator range(int minDoc, int maxDoc) {
     return new RangeDocIdSetIterator(minDoc, maxDoc);
   }
 
-  private static class RangeDocIdSetIterator extends DocIdSetIterator {
+  private static class RangeDocIdSetIterator extends AbstractDocIdSetIterator {
 
     private final int minDoc, maxDoc;
-    private int doc = -1;
 
     RangeDocIdSetIterator(int minDoc, int maxDoc) {
       // advance relies on minDoc <= maxDoc for correctness
@@ -67,11 +66,6 @@ private static class RangeDocIdSetIterator extends DocIdSetIterator {
       this.maxDoc = maxDoc;
     }
 
-    @Override
-    public int docID() {
-      return doc;
-    }
-
     @Override
     public int nextDoc() {
       return advance(doc + 1);

lucene/core/src/java/org/apache/lucene/search/DocValuesRangeIterator.java

@@ -57,16 +57,14 @@ public DocValuesRangeIterator(
     this.innerTwoPhase = twoPhase;
   }
 
-  abstract static class Approximation extends DocIdSetIterator {
+  abstract static class Approximation extends AbstractDocIdSetIterator {
 
     private final DocIdSetIterator innerApproximation;
 
     protected final DocValuesSkipper skipper;
     protected final long lowerValue;
     protected final long upperValue;
 
-    private int doc = -1;
-
     // Track a decision for all doc IDs between the current doc ID and upTo inclusive.
     Match match = Match.MAYBE;
     int upTo = -1;
@@ -82,11 +80,6 @@ abstract static class Approximation extends DocIdSetIterator {
       this.upperValue = upperValue;
     }
 
-    @Override
-    public int docID() {
-      return doc;
-    }
-
     @Override
     public int nextDoc() throws IOException {
       return advance(docID() + 1);

lucene/core/src/java/org/apache/lucene/search/FilteredDocIdSetIterator.java

@@ -22,9 +22,8 @@
  * Abstract decorator class of a DocIdSetIterator implementation that provides on-demand
  * filter/validation mechanism on an underlying DocIdSetIterator.
  */
-public abstract class FilteredDocIdSetIterator extends DocIdSetIterator {
+public abstract class FilteredDocIdSetIterator extends AbstractDocIdSetIterator {
   protected DocIdSetIterator _innerIter;
-  private int doc;
 
   /**
    * Constructor.
@@ -53,11 +52,6 @@ public DocIdSetIterator getDelegate() {
    */
   protected abstract boolean match(int doc) throws IOException;
 
-  @Override
-  public int docID() {
-    return doc;
-  }
-
   @Override
   public int nextDoc() throws IOException {
     while ((doc = _innerIter.nextDoc()) != NO_MORE_DOCS) {

lucene/core/src/java/org/apache/lucene/search/ImpactsDISI.java

@@ -25,9 +25,8 @@
  *
  * @lucene.internal
  */
-public final class ImpactsDISI extends DocIdSetIterator {
+public final class ImpactsDISI extends FilterDocIdSetIterator {
 
-  private final DocIdSetIterator in;
   private final MaxScoreCache maxScoreCache;
   private float minCompetitiveScore = 0;
   private int upTo = DocIdSetIterator.NO_MORE_DOCS;
@@ -40,7 +39,7 @@ public final class ImpactsDISI extends DocIdSetIterator {
    * @param maxScoreCache the cache of maximum scores, typically computed from the same ImpactsEnum
    */
   public ImpactsDISI(DocIdSetIterator in, MaxScoreCache maxScoreCache) {
-    this.in = in;
+    super(in);
     this.maxScoreCache = maxScoreCache;
   }
 
@@ -112,14 +111,4 @@ public int nextDoc() throws IOException {
     }
     return advance(in.docID() + 1);
   }
-
-  @Override
-  public int docID() {
-    return in.docID();
-  }
-
-  @Override
-  public long cost() {
-    return in.cost();
-  }
 }

lucene/core/src/java/org/apache/lucene/search/IndexSortSortedNumericDocValuesRangeQuery.java

@@ -685,28 +685,21 @@ static IteratorAndCount sparseRange(int minDoc, int maxDoc, DocIdSetIterator del
    * A doc ID set iterator that wraps a delegate iterator and only returns doc IDs in the range
    * [firstDocInclusive, lastDoc).
    */
-  private static class BoundedDocIdSetIterator extends DocIdSetIterator {
+  private static class BoundedDocIdSetIterator extends AbstractDocIdSetIterator {
     private final int firstDoc;
     private final int lastDoc;
     private final DocIdSetIterator delegate;
 
-    private int docID = -1;
-
     BoundedDocIdSetIterator(int firstDoc, int lastDoc, DocIdSetIterator delegate) {
       assert delegate != null;
       this.firstDoc = firstDoc;
       this.lastDoc = lastDoc;
       this.delegate = delegate;
     }
 
-    @Override
-    public int docID() {
-      return docID;
-    }
-
     @Override
     public int nextDoc() throws IOException {
-      return advance(docID + 1);
+      return advance(doc + 1);
     }
 
     @Override
@@ -717,11 +710,11 @@ public int advance(int target) throws IOException {
 
       int result = delegate.advance(target);
       if (result < lastDoc) {
-        docID = result;
+        doc = result;
       } else {
-        docID = NO_MORE_DOCS;
+        doc = NO_MORE_DOCS;
       }
-      return docID;
+      return doc;
     }
 
     @Override