LUCENE-10289: Change DocIdSetBuilder#grow() from taking an int to a long (#520)

lucene/CHANGES.txt

@@ -42,6 +42,8 @@ API Changes
 * LUCENE-10244: MultiCollector::getCollectors is now public, allowing users to access the wrapped
   collectors. (Andriy Redko)
   
+* LUCENE-10289: DocIdSetBuilder#grow() takes now a long instead of an int. (Ignacio Vera)  
+
 New Features
 ---------------------
 

lucene/core/src/java/org/apache/lucene/util/DocIdSetBuilder.java

@@ -166,9 +166,9 @@ public final class DocIdSetBuilder {
       bitSet.or(iter);
       return;
     }
-    int cost = (int) Math.min(Integer.MAX_VALUE, iter.cost());
+    long cost = iter.cost();
     BulkAdder adder = grow(cost);
-    for (int i = 0; i < cost; ++i) {
+    for (long i = 0; i < cost; ++i) {
       int doc = iter.nextDoc();
       if (doc == DocIdSetIterator.NO_MORE_DOCS) {
         return;
@@ -184,20 +184,35 @@ public final class DocIdSetBuilder {
    * Reserve space and return a {@link BulkAdder} object that can be used to add up to {@code
    * numDocs} documents.
    */
-  public BulkAdder grow(int numDocs) {
+  public BulkAdder grow(long numDocs) {
     if (bitSet == null) {
-      if ((long) totalAllocated + numDocs <= threshold) {
-        ensureBufferCapacity(numDocs);
+      if ((long) totalAllocated + checkTotalAllocatedOverflow(numDocs) <= threshold) {
+        // For extra safety we use toIntExact
+        ensureBufferCapacity(Math.toIntExact(numDocs));
       } else {
         upgradeToBitSet();
-        counter += numDocs;
+        counter += checkCounterOverflow(numDocs);
       }
     } else {
-      counter += numDocs;
+      counter += checkCounterOverflow(numDocs);
     }
     return adder;
   }
 
+  private long checkTotalAllocatedOverflow(long numDocs) {
+    if ((long) totalAllocated + numDocs < totalAllocated) {
+      throw new ArithmeticException("long overflow");
+    }
+    return numDocs;
+  }
+
+  private long checkCounterOverflow(long numDocs) {
+    if (counter + numDocs < counter) {
+      throw new ArithmeticException("long overflow");
+    }
+    return numDocs;
+  }
+
   private void ensureBufferCapacity(int numDocs) {
     if (buffers.isEmpty()) {
       addBuffer(additionalCapacity(numDocs));

lucene/core/src/test/org/apache/lucene/util/TestDocIdSetBuilder.java

@@ -128,9 +128,9 @@ public class TestDocIdSetBuilder extends LuceneTestCase {
       for (j = 0; j < array.length; ) {
         final int l = TestUtil.nextInt(random(), 1, array.length - j);
         DocIdSetBuilder.BulkAdder adder = null;
-        for (int k = 0, budget = 0; k < l; ++k) {
+        for (long k = 0, budget = 0; k < l; ++k) {
           if (budget == 0 || rarely()) {
-            budget = TestUtil.nextInt(random(), 1, l - k + 5);
+            budget = TestUtil.nextLong(random(), 1, l - k + 5);
             adder = builder.grow(budget);
           }
           adder.add(array[j++]);
@@ -241,6 +241,38 @@ public class TestDocIdSetBuilder extends LuceneTestCase {
     assertTrue(builder.multivalued);
   }
 
+  @Nightly
+  public void testLotsOfDocs() throws IOException {
+    final int docCount = 1;
+    final long numDocs = (long) Integer.MAX_VALUE + 1;
+    PointValues values = new DummyPointValues(docCount, numDocs);
+    DocIdSetBuilder builder = new DocIdSetBuilder(100, values, "foo");
+    DocIdSetBuilder.BulkAdder adder = builder.grow(numDocs);
+    for (long i = 0; i < numDocs; ++i) {
+      adder.add(0);
+    }
+    DocIdSet result = builder.build();
+    assertTrue(result instanceof BitDocIdSet);
+    assertEquals(1, result.iterator().cost());
+  }
+
+  public void testLongOverflow() throws IOException {
+    {
+      DocIdSetBuilder builder = new DocIdSetBuilder(100);
+      builder.grow(1L);
+      Exception ex = expectThrows(ArithmeticException.class, () -> builder.grow(Long.MAX_VALUE));
+      assertEquals("long overflow", ex.getMessage());
+    }
+    {
+      DocIdSetBuilder builder = new DocIdSetBuilder(100);
+      builder.grow((long) Integer.MAX_VALUE + 1);
+      Exception ex =
+          expectThrows(
+              ArithmeticException.class, () -> builder.grow(Long.MAX_VALUE - Integer.MAX_VALUE));
+      assertEquals("long overflow", ex.getMessage());
+    }
+  }
+
   private static class DummyTerms extends Terms {
 
     private final int docCount;