perf(common): reduce calls to get_unchecked (#8848)

src/common/benches/bitmap.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
 use itertools::Itertools;
 use risingwave_common::buffer::{Bitmap, BitmapIter};
 
@@ -23,7 +23,27 @@ fn bench_bitmap(c: &mut Criterion) {
     let i = 0x123;
     c.bench_function("zeros", |b| b.iter(|| Bitmap::zeros(CHUNK_SIZE)));
     c.bench_function("ones", |b| b.iter(|| Bitmap::ones(CHUNK_SIZE)));
-    c.bench_function("get", |b| b.iter(|| x.is_set(i)));
+    c.bench_function("get", |b| {
+        b.iter(|| {
+            for _ in 0..1000 {
+                black_box(x.is_set(i));
+            }
+        })
+    });
+    c.bench_function("get_1000", |b| {
+        b.iter(|| {
+            for _ in 0..1000 {
+                black_box(x.is_set(i));
+            }
+        })
+    });
+    c.bench_function("get_1000_000", |b| {
+        b.iter(|| {
+            for _ in 0..1_000_000 {
+                black_box(x.is_set(i));
+            }
+        })
+    });
     c.bench_function("and", |b| b.iter(|| &x & &y));
     c.bench_function("or", |b| b.iter(|| &x | &y));
     c.bench_function("not", |b| b.iter(|| !&x));
@@ -36,17 +56,16 @@ fn bench_bitmap_iter(c: &mut Criterion) {
     fn make_iterators(bitmaps: &[Bitmap]) -> Vec<BitmapIter<'_>> {
         bitmaps.iter().map(|bitmap| bitmap.iter()).collect_vec()
     }
-    fn bench_bitmap_iter(bench_id: &str, bitmap: Bitmap, c: &mut Criterion) {
+    fn bench_bitmap_iter_inner(bench_id: &str, bitmap: Bitmap, c: &mut Criterion) {
         let bitmaps = vec![bitmap; N_CHUNKS];
         let make_iters = || make_iterators(&bitmaps);
-        let mut result = vec![true; CHUNK_SIZE];
         c.bench_function(bench_id, |b| {
             b.iter_batched(
                 make_iters,
                 |iters| {
                     for iter in iters {
-                        for (i, bit_flag) in iter.enumerate() {
-                            result[i] = bit_flag;
+                        for bit_flag in iter {
+                            black_box(bit_flag);
                         }
                     }
                 },
@@ -55,9 +74,9 @@ fn bench_bitmap_iter(c: &mut Criterion) {
         });
     }
     let zeros = Bitmap::zeros(CHUNK_SIZE);
-    bench_bitmap_iter("zeros_iter", zeros, c);
+    bench_bitmap_iter_inner("zeros_iter", zeros, c);
     let ones = Bitmap::ones(CHUNK_SIZE);
-    bench_bitmap_iter("ones_iter", ones, c);
+    bench_bitmap_iter_inner("ones_iter", ones, c);
 }
 
 criterion_group!(benches, bench_bitmap, bench_bitmap_iter);

src/common/src/buffer/bitmap.rs

@@ -173,13 +173,15 @@ impl BitmapBuilder {
 /// An immutable bitmap. Use [`BitmapBuilder`] to build it.
 #[derive(Clone, PartialEq, Eq)]
 pub struct Bitmap {
-    // The useful bits in the bitmap. The total number of bits will usually
-    // be larger than the useful bits due to byte-padding.
+    /// The useful bits in the bitmap. The total number of bits will usually
+    /// be larger than the useful bits due to byte-padding.
     num_bits: usize,
 
-    // The number of high bits in the bitmap.
+    /// The number of high bits in the bitmap.
     count_ones: usize,
 
+    /// Bits are stored in a compact form.
+    /// They are packed into `usize`s.
     bits: Box<[usize]>,
 }
 
@@ -310,6 +312,7 @@ impl Bitmap {
             bits: &self.bits,
             idx: 0,
             num_bits: self.num_bits,
+            current_usize: 0,
         }
     }
 
@@ -556,6 +559,27 @@ pub struct BitmapIter<'a> {
     bits: &'a [usize],
     idx: usize,
     num_bits: usize,
+    current_usize: usize,
+}
+
+impl<'a> BitmapIter<'a> {
+    fn next_always_load_usize(&mut self) -> Option<bool> {
+        if self.idx >= self.num_bits {
+            return None;
+        }
+
+        // Offset of the bit within the usize.
+        let usize_offset = self.idx % BITS;
+
+        // Get the index of usize which the bit is located in
+        let usize_index = self.idx / BITS;
+        self.current_usize = unsafe { *self.bits.get_unchecked(usize_index) };
+
+        let bit_mask = 1 << usize_offset;
+        let bit_flag = self.current_usize & bit_mask != 0;
+        self.idx += 1;
+        Some(bit_flag)
+    }
 }
 
 impl<'a> iter::Iterator for BitmapIter<'a> {
@@ -565,9 +589,21 @@ impl<'a> iter::Iterator for BitmapIter<'a> {
         if self.idx >= self.num_bits {
             return None;
         }
-        let b = unsafe { self.bits.get_unchecked(self.idx / BITS) } & (1 << (self.idx % BITS)) != 0;
+
+        // Offset of the bit within the usize.
+        let usize_offset = self.idx % BITS;
+
+        if usize_offset == 0 {
+            // Get the index of usize which the bit is located in
+            let usize_index = self.idx / BITS;
+            self.current_usize = unsafe { *self.bits.get_unchecked(usize_index) };
+        }
+
+        let bit_mask = 1 << usize_offset;
+
+        let bit_flag = self.current_usize & bit_mask != 0;
         self.idx += 1;
-        Some(b)
+        Some(bit_flag)
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
@@ -577,7 +613,7 @@ impl<'a> iter::Iterator for BitmapIter<'a> {
 
     fn nth(&mut self, n: usize) -> Option<Self::Item> {
         self.idx += n;
-        self.next()
+        self.next_always_load_usize()
     }
 }