Fix DISIDocIdStream::count so that it does not try to count beyond max (apache#14522)

Fix DISIDocIdStream::count so that it does not try to count beyond max.

From the perspective of DISIDocIdStream it should not be necessary to check the iterator position before calling count.

Author: ChrisHegarty

lucene/CHANGES.txt

@@ -7,7 +7,9 @@ http://s.apache.org/luceneversions
 
 Bug Fixes
 ---------------------
-(No changes)
+
+* GITHUB#14522: Fix DISIDocIdStream::count so that it does not try to count beyond max.
+  (Chris Hegarty}
 
 ======================= Lucene 10.2.0 =======================
 

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

@@ -51,13 +51,13 @@ public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws I
 
   @Override
   public int count(int upTo) throws IOException {
+    upTo = Math.min(upTo, max);
     if (iterator.docID() >= upTo) {
       return 0;
     }
     // If the collector is just interested in the count, loading in a bit set and counting bits is
     // often faster than incrementing a counter on every call to nextDoc().
     assert spare.scanIsEmpty();
-    upTo = Math.min(upTo, max);
     int offset = iterator.docID();
     iterator.intoBitSet(upTo, spare, offset);
     int count = spare.cardinality(0, upTo - offset);

lucene/core/src/test/org/apache/lucene/search/TestDISIDocIdStream.java

@@ -111,6 +111,24 @@ public void testCountUpTo() throws IOException {
     assertEquals(bitSet.cardinality(60, 80), stream.count(120));
 
     assertFalse(stream.mayHaveRemaining());
+
+    assertEquals(0, stream.count(40)); // we're well past this doc
+    assertEquals(0, stream.count(100)); // we're well past this doc/max
+    assertEquals(0, stream.count()); // we're well past this doc/max
+  }
+
+  // Tests that count when the iterator, but not the external interaction with the stream,
+  // is past max does not fail, returns 0.
+  public void testCountUpTo2() throws IOException {
+    FixedBitSet bitSet = new FixedBitSet(100);
+    bitSet.set(1);
+    bitSet.set(85);
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+    DocIdStream stream = new DISIDocIdStream(iterator, 80, new FixedBitSet(100));
+    iterator.advance(1);
+
+    assertEquals(1, stream.count(79));
+    assertEquals(0, stream.count());
   }
 
   public void testMixForEachCountUpTo() throws IOException {