From f289a0726be610bde1e2ef8c7f144999faeee6fc Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Mon, 3 Apr 2023 13:40:49 -0500
Subject: [PATCH] disk index: refactoring get/get_mut fns (#31025)

---
 bucket_map/src/bucket.rs         | 11 +++++++---
 bucket_map/src/bucket_storage.rs | 37 +++++++++++++++++++++++---------
 bucket_map/src/index_entry.rs    |  8 +++++--
 3 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/bucket_map/src/bucket.rs b/bucket_map/src/bucket.rs
index dcad861aa..bcca7b605 100644
--- a/bucket_map/src/bucket.rs
+++ b/bucket_map/src/bucket.rs
@@ -3,7 +3,7 @@ use {
         bucket_item::BucketItem,
         bucket_map::BucketMapError,
         bucket_stats::BucketMapStats,
-        bucket_storage::{BucketOccupied, BucketStorage, DEFAULT_CAPACITY_POW2},
+        bucket_storage::{BucketOccupied, BucketStorage, IncludeHeader, DEFAULT_CAPACITY_POW2},
         index_entry::{
             DataBucket, IndexBucket, IndexEntry, IndexEntryPlaceInBucket, MultipleSlots,
             OccupiedEnum,
@@ -328,7 +328,11 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
             if best_fit_bucket == bucket_ix as u64 {
                 // in place update in same data file
                 assert!(!current_bucket.is_free(elem_loc));
-                let slice: &mut [T] = current_bucket.get_mut_cell_slice(elem_loc, data_len as u64);
+                let slice: &mut [T] = current_bucket.get_mut_cell_slice(
+                    elem_loc,
+                    data_len as u64,
+                    IncludeHeader::NoHeader,
+                );
                 multiple_slots.set_num_slots(num_slots);
 
                 slice.iter_mut().zip(data).for_each(|(dest, src)| {
@@ -375,7 +379,8 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
                     // copy slotlist into the data bucket
                     let best_bucket = &mut self.data[best_fit_bucket as usize];
                     best_bucket.occupy(ix, false).unwrap();
-                    let slice = best_bucket.get_mut_cell_slice(ix, num_slots);
+                    let slice =
+                        best_bucket.get_mut_cell_slice(ix, num_slots, IncludeHeader::NoHeader);
                     slice.iter_mut().zip(data).for_each(|(dest, src)| {
                         *dest = *src;
                     });
diff --git a/bucket_map/src/bucket_storage.rs b/bucket_map/src/bucket_storage.rs
index 3bded3578..c9820ae43 100644
--- a/bucket_map/src/bucket_storage.rs
+++ b/bucket_map/src/bucket_storage.rs
@@ -86,6 +86,14 @@ impl<O: BucketOccupied> Drop for BucketStorage<O> {
     }
 }
 
+#[allow(dead_code)]
+pub(crate) enum IncludeHeader {
+    /// caller wants header
+    Header,
+    /// caller wants header skipped
+    NoHeader,
+}
+
 impl<O: BucketOccupied> BucketStorage<O> {
     pub fn new_with_capacity(
         drives: Arc<Vec<PathBuf>>,
@@ -198,18 +206,22 @@ impl<O: BucketOccupied> BucketStorage<O> {
         (self.cell_size * ix) as usize
     }
 
-    fn get_start_offset_no_header(&self, ix: u64) -> usize {
-        self.get_start_offset_with_header(ix) + O::offset_to_first_data()
+    fn get_start_offset(&self, ix: u64, header: IncludeHeader) -> usize {
+        self.get_start_offset_with_header(ix)
+            + match header {
+                IncludeHeader::Header => 0,
+                IncludeHeader::NoHeader => O::offset_to_first_data(),
+            }
     }
 
-    pub fn get<T>(&self, ix: u64) -> &T {
-        let slice = self.get_cell_slice::<T>(ix, 1);
+    pub(crate) fn get<T>(&self, ix: u64) -> &T {
+        let slice = self.get_cell_slice::<T>(ix, 1, IncludeHeader::NoHeader);
         // SAFETY: `get_cell_slice` ensures there's at least one element in the slice
         unsafe { slice.get_unchecked(0) }
     }
 
-    pub fn get_mut<T>(&mut self, ix: u64) -> &mut T {
-        let slice = self.get_mut_cell_slice::<T>(ix, 1);
+    pub(crate) fn get_mut<T>(&mut self, ix: u64) -> &mut T {
+        let slice = self.get_mut_cell_slice::<T>(ix, 1, IncludeHeader::NoHeader);
         // SAFETY: `get_mut_cell_slice` ensures there's at least one element in the slice
         unsafe { slice.get_unchecked_mut(0) }
     }
@@ -226,8 +238,8 @@ impl<O: BucketOccupied> BucketStorage<O> {
         unsafe { &*item }
     }
 
-    pub fn get_cell_slice<T>(&self, ix: u64, len: u64) -> &[T] {
-        let start = self.get_start_offset_no_header(ix);
+    pub(crate) fn get_cell_slice<T>(&self, ix: u64, len: u64, header: IncludeHeader) -> &[T] {
+        let start = self.get_start_offset(ix, header);
         let slice = {
             let size = std::mem::size_of::<T>() * len as usize;
             let slice = &self.mmap[start..];
@@ -242,8 +254,13 @@ impl<O: BucketOccupied> BucketStorage<O> {
         unsafe { std::slice::from_raw_parts(ptr, len as usize) }
     }
 
-    pub fn get_mut_cell_slice<T>(&mut self, ix: u64, len: u64) -> &mut [T] {
-        let start = self.get_start_offset_no_header(ix);
+    pub(crate) fn get_mut_cell_slice<T>(
+        &mut self,
+        ix: u64,
+        len: u64,
+        header: IncludeHeader,
+    ) -> &mut [T] {
+        let start = self.get_start_offset(ix, header);
         let slice = {
             let size = std::mem::size_of::<T>() * len as usize;
             let slice = &mut self.mmap[start..];
diff --git a/bucket_map/src/index_entry.rs b/bucket_map/src/index_entry.rs
index db22e2765..81dbd4c70 100644
--- a/bucket_map/src/index_entry.rs
+++ b/bucket_map/src/index_entry.rs
@@ -2,7 +2,7 @@
 
 use {
     crate::{
-        bucket_storage::{BucketOccupied, BucketStorage},
+        bucket_storage::{BucketOccupied, BucketStorage, IncludeHeader},
         RefCount,
     },
     bv::BitVec,
@@ -308,7 +308,11 @@ impl<T: Copy> IndexEntryPlaceInBucket<T> {
                     let data_bucket = &data_buckets[data_bucket_ix as usize];
                     let loc = multiple_slots.data_loc(data_bucket);
                     assert!(!data_bucket.is_free(loc));
-                    data_bucket.get_cell_slice::<T>(loc, multiple_slots.num_slots)
+                    data_bucket.get_cell_slice::<T>(
+                        loc,
+                        multiple_slots.num_slots,
+                        IncludeHeader::NoHeader,
+                    )
                 }
                 _ => {
                     panic!("trying to read data from a free entry");