Simplify BucketedSort/Teach BitArray a useful trick (backport #53199)

server/src/main/java/org/elasticsearch/common/util/BitArray.java

@@ -37,16 +37,32 @@ public BitArray(int initialSize, BigArrays bigArrays) {
         this.bits = bigArrays.newLongArray(initialSize, true);
     }
 
+    /**
+     * Set the {@code index}th bit.
+     */
     public void set(int index) {
         fill(index, true);
     }
 
+    /**
+     * Clear the {@code index}th bit.
+     */
     public void clear(int index) {
         fill(index, false);
     }
 
+    /**
+     * Is the {@code index}th bit set?
+     */
     public boolean get(int index) {
         int wordNum = index >> 6;
+        if (wordNum >= bits.size()) {
+            /*
+             * If the word is bigger than the array then it could *never* have
+             * been set.
+             */
+            return false;
+        }
         long bitmask = 1L << index;
         return (bits.get(wordNum) & bitmask) != 0;
     }

server/src/main/java/org/elasticsearch/search/sort/BucketedSort.java

@@ -134,13 +134,6 @@ public Loader loader(LeafReaderContext ctx) throws IOException {
      * it is still gathering.
      */
     private final BitArray heapMode;
-    /**
-     * The highest bucket ordinal that has been converted into a heap. This is
-     * required because calling {@link BitArray#get(int)} on an index higher
-     * than the highest one that was {@link BitArray#set(int) set} could throw
-     * and {@link ArrayIndexOutOfBoundsException}. So we check this first.
-     */
-    private long maxHeapBucket = 0;
 
     protected BucketedSort(BigArrays bigArrays, SortOrder order, DocValueFormat format, int bucketSize, ExtraData extra) {
         this.bigArrays = bigArrays;
@@ -216,7 +209,7 @@ public final List<SortValue> getValues(long bucket) {
      * Is this bucket a min heap {@code true} or in gathering mode {@code false}? 
      */
     private boolean inHeapMode(long bucket) {
-        return bucket <= maxHeapBucket && heapMode.get((int) bucket);
+        return heapMode.get((int) bucket);
     }
 
     /**
@@ -430,7 +423,6 @@ public final void collect(int doc, long bucket) throws IOException {
                     throw new UnsupportedOperationException("Bucketed sort doesn't support more than [" + Integer.MAX_VALUE + "] buckets");
                     // BitArray needs int keys and this'd take a ton of memory to use that many buckets. So we just don't.
                 }
-                maxHeapBucket = Math.max(bucket, maxHeapBucket);
                 heapMode.set((int) bucket);
                 heapify(rootIndex);
             } else {

server/src/test/java/org/elasticsearch/common/util/BitArrayTests.java

@@ -50,4 +50,17 @@ public void testRandom() {
             }
         }
     }
+
+    public void testTooBigIsNotSet() {
+        try (BitArray bitArray = new BitArray(1, BigArrays.NON_RECYCLING_INSTANCE)) {
+            for (int i = 0; i < 1000; i++) {
+                /*
+                 * The first few times this is called we check within the
+                 * array. But we quickly go beyond it and those all return
+                 * false as well.
+                 */
+                assertFalse(bitArray.get(i));
+            }
+        }
+    }
 }