From b0585f11de1aca1a5d730408793b268daf2f9a90 Mon Sep 17 00:00:00 2001
From: Adrien Grand <jpountz@gmail.com>
Date: Mon, 23 Dec 2024 09:39:59 +0100
Subject: [PATCH] Make inlining of FixedBitSet#get more predictable when
 checking live docs. (#14077)

This helps make calls sites of `Bits#get` bimorphic at most when checking live
docs. This helps because calls to `FixedBitSet#get` can then be inlined when
live docs are stored in a `FixedBitSet`. Another reason why this is important
is because these calls could be subject to auto-vectorizing when applied to an
array of doc IDs, which cannot apply if the `FixedBitSet#get` call is not
inlined.

While live docs are stored in a `FixedBitSet` by default, some features like
document-level security sometimes create wrappers.
---
 .../apache/lucene/search/IndexSearcher.java   |  4 +-
 .../org/apache/lucene/search/ScorerUtil.java  | 38 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java b/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
index 6e8bbf81966..e3074a96d88 100644
--- a/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
+++ b/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
@@ -845,7 +845,9 @@ public class IndexSearcher {
         scorer = new TimeLimitingBulkScorer(scorer, queryTimeout);
       }
       try {
-        scorer.score(leafCollector, ctx.reader().getLiveDocs(), minDocId, maxDocId);
+        // Optimize for the case when live docs are stored in a FixedBitSet.
+        Bits acceptDocs = ScorerUtil.likelyFixedBitSet(ctx.reader().getLiveDocs());
+        scorer.score(leafCollector, acceptDocs, minDocId, maxDocId);
       } catch (
           @SuppressWarnings("unused")
           CollectionTerminatedException e) {
diff --git a/lucene/core/src/java/org/apache/lucene/search/ScorerUtil.java b/lucene/core/src/java/org/apache/lucene/search/ScorerUtil.java
index 1154dc1b154..2081f9ac717 100644
--- a/lucene/core/src/java/org/apache/lucene/search/ScorerUtil.java
+++ b/lucene/core/src/java/org/apache/lucene/search/ScorerUtil.java
@@ -30,7 +30,9 @@ import org.apache.lucene.index.PostingsEnum;
 import org.apache.lucene.index.TermsEnum;
 import org.apache.lucene.store.ByteBuffersDirectory;
 import org.apache.lucene.store.Directory;
+import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.FixedBitSet;
 import org.apache.lucene.util.PriorityQueue;
 
 /** Util class for Scorer related methods */
@@ -108,4 +110,40 @@ class ScorerUtil {
     }
     return scorable;
   }
+
+  /**
+   * Optimize {@link Bits} representing the set of accepted documents for the case when it is likely
+   * implemented via a {@link FixedBitSet}. This helps make calls to {@link Bits#get(int)}
+   * inlinable, which in-turn helps speed up query evaluation. This is especially helpful as
+   * inlining will sometimes enable auto-vectorizing shifts and masks that are done in {@link
+   * FixedBitSet#get(int)}.
+   */
+  static Bits likelyFixedBitSet(Bits acceptDocs) {
+    if (acceptDocs instanceof FixedBitSet) {
+      return acceptDocs;
+    } else if (acceptDocs != null) {
+      return new FilterBits(acceptDocs);
+    } else {
+      return null;
+    }
+  }
+
+  private static class FilterBits implements Bits {
+
+    private final Bits in;
+
+    FilterBits(Bits in) {
+      this.in = in;
+    }
+
+    @Override
+    public boolean get(int index) {
+      return in.get(index);
+    }
+
+    @Override
+    public int length() {
+      return in.length();
+    }
+  }
 }