Updating Dense#intoBitSet to properly set upTo if it exceeds bitset size (#14922)

Author: pmpailis

lucene/CHANGES.txt

@@ -234,6 +234,9 @@ Bug Fixes
 
 * GITHUB#14838: Make it possible to extend Patience/Seeded knn queries (Tommaso Teofili)
 
+* GITHUB#14922: Fix for IndexedDISI.Dense#intoBitSetWithinBlock to take into account the provided bitset size
+  and avoid throwing IndexOutOfBoundsException (Panagiotis Bailis)
+
 Build
 ---------------------
 * Upgrade forbiddenapis to version 3.9.  (Uwe Schindler)

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

@@ -27,6 +27,7 @@
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.BitSetIterator;
 import org.apache.lucene.util.FixedBitSet;
+import org.apache.lucene.util.MathUtil;
 import org.apache.lucene.util.RoaringDocIdSet;
 
 /**
@@ -473,7 +474,7 @@ public int advance(int target) throws IOException {
 
   @Override
   public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
-    assert doc >= offset;
+    assert doc >= offset : "offset=" + offset + " doc=" + doc;
     while (doc < upTo && method.intoBitSetWithinBlock(this, upTo, bitSet, offset) == false) {
       readBlockHeader();
       boolean found = method.advanceWithinBlock(this, block);
@@ -719,10 +720,10 @@ boolean intoBitSetWithinBlock(IndexedDISI disi, int upTo, FixedBitSet bitSet, in
         if (disi.bitSet == null) {
           disi.bitSet = new FixedBitSet(BLOCK_SIZE);
         }
-
-        int sourceFrom = disi.doc & 0xFFFF;
-        int sourceTo = Math.min(upTo - disi.block, BLOCK_SIZE);
         int destFrom = disi.doc - offset;
+        int destTo = MathUtil.unsignedMin(upTo, offset + bitSet.length());
+        int sourceFrom = disi.doc & 0xFFFF;
+        int sourceTo = Math.min(destTo - disi.block, BLOCK_SIZE);
 
         long fp = disi.slice.getFilePointer();
         disi.slice.seek(fp - Long.BYTES); // seek back a long to include current word (disi.word).
@@ -731,13 +732,24 @@ boolean intoBitSetWithinBlock(IndexedDISI disi, int upTo, FixedBitSet bitSet, in
         FixedBitSet.orRange(disi.bitSet, sourceFrom, bitSet, destFrom, sourceTo - sourceFrom);
 
         int blockEnd = disi.block | 0xFFFF;
-        if (upTo > blockEnd) {
+        if (destTo > blockEnd) {
           disi.slice.seek(disi.blockEnd);
           disi.index += disi.bitSet.cardinality(sourceFrom, sourceTo);
           return false;
         } else {
           disi.slice.seek(fp);
-          return advanceWithinBlock(disi, upTo);
+          boolean found = advanceWithinBlock(disi, destTo);
+          if (found && disi.doc < upTo) {
+            throw new IllegalStateException(
+                "There are bits set in the source bitset that are not accounted for."
+                    + " doc="
+                    + disi.doc
+                    + " upTo="
+                    + upTo
+                    + " disi.block="
+                    + disi.block);
+          }
+          return found;
         }
       }
 

lucene/core/src/test/org/apache/lucene/codecs/lucene90/TestIndexedDISI.java

@@ -388,6 +388,36 @@ public void testDenseMultiBlock() throws IOException {
     }
   }
 
+  public void testDenseBitSizeLessThanBlockSize() throws IOException {
+    final byte denseRankPower = (byte) (random().nextInt(7) + 7);
+    try (Directory dir = newDirectory()) {
+      // initialize a maxDoc that is less than IndexedDISI.BLOCK_SIZE
+      int maxDoc = random().nextInt(4096 * 2, 65536);
+      FixedBitSet set = new FixedBitSet(maxDoc);
+      for (int i = 0; i < maxDoc; i += 2) { // Set every other to ensure dense
+        set.set(i);
+      }
+      int jumpTableEntryCount; // this should always be 0 given that maxDoc < BLOCK_SIZE
+      long length;
+      try (IndexOutput out = dir.createOutput("foo", IOContext.DEFAULT)) {
+        jumpTableEntryCount =
+            IndexedDISI.writeBitSet(
+                new BitSetIterator(set, set.cardinality()), out, denseRankPower);
+        length = out.getFilePointer();
+        assertTrue(
+            "jumpTableEntryCount should be 0 for dense bitsets with size < BLOCK_SIZE",
+            0 == jumpTableEntryCount);
+      }
+      try (IndexInput in = dir.openInput("foo", IOContext.DEFAULT)) {
+        IndexedDISI disi =
+            new IndexedDISI(in, 0L, length, jumpTableEntryCount, denseRankPower, set.cardinality());
+        FixedBitSet disiSet = new FixedBitSet(maxDoc);
+        // This would throw IOOB if bitset size is not handled correctly as per #14882
+        disiSet.or(disi);
+      }
+    }
+  }
+
   public void testIllegalDenseRankPower() throws IOException {
 
     // Legal values