From 84eaaae062c0e14380c29990cc84c7bd74b21f80 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <wash678@gmail.com>
Date: Tue, 21 Dec 2021 13:50:04 -0600
Subject: [PATCH] disk_buckets: factor out unsafe code (#22028)

---
 bucket_map/src/bucket_storage.rs | 39 +++++++++++++-------------------
 1 file changed, 16 insertions(+), 23 deletions(-)

diff --git a/bucket_map/src/bucket_storage.rs b/bucket_map/src/bucket_storage.rs
index 3a54f111ba..da50bd8d8e 100644
--- a/bucket_map/src/bucket_storage.rs
+++ b/bucket_map/src/bucket_storage.rs
@@ -126,29 +126,29 @@ impl BucketStorage {
         )
     }
 
-    pub fn uid(&self, ix: u64) -> Uid {
-        assert!(ix < self.capacity(), "bad index size");
+    /// return ref to header of item 'ix' in mmapped file
+    fn header_ptr(&self, ix: u64) -> &Header {
         let ix = (ix * self.cell_size) as usize;
         let hdr_slice: &[u8] = &self.mmap[ix..ix + std::mem::size_of::<Header>()];
         unsafe {
             let hdr = hdr_slice.as_ptr() as *const Header;
-            return hdr.as_ref().unwrap().uid();
+            hdr.as_ref().unwrap()
         }
     }
 
+    pub fn uid(&self, ix: u64) -> Uid {
+        assert!(ix < self.capacity(), "bad index size");
+        self.header_ptr(ix).uid()
+    }
+
     /// 'is_resizing' true if caller is resizing the index (so don't increment count)
     /// 'is_resizing' false if caller is adding an item to the index (so increment count)
     pub fn allocate(&self, ix: u64, uid: Uid, is_resizing: bool) -> Result<(), BucketStorageError> {
         assert!(ix < self.capacity(), "allocate: bad index size");
         assert!(UID_UNLOCKED != uid, "allocate: bad uid");
         let mut e = Err(BucketStorageError::AlreadyAllocated);
-        let ix = (ix * self.cell_size) as usize;
         //debug!("ALLOC {} {}", ix, uid);
-        let hdr_slice: &[u8] = &self.mmap[ix..ix + std::mem::size_of::<Header>()];
-        if unsafe {
-            let hdr = hdr_slice.as_ptr() as *const Header;
-            hdr.as_ref().unwrap().try_lock(uid)
-        } {
+        if self.header_ptr(ix).try_lock(uid) {
             e = Ok(());
             if !is_resizing {
                 self.count.fetch_add(1, Ordering::Relaxed);
@@ -160,20 +160,13 @@ impl BucketStorage {
     pub fn free(&self, ix: u64, uid: Uid) {
         assert!(ix < self.capacity(), "bad index size");
         assert!(UID_UNLOCKED != uid, "free: bad uid");
-        let ix = (ix * self.cell_size) as usize;
-        //debug!("FREE {} {}", ix, uid);
-        let hdr_slice: &[u8] = &self.mmap[ix..ix + std::mem::size_of::<Header>()];
-        unsafe {
-            let hdr = hdr_slice.as_ptr() as *const Header;
-            //debug!("FREE uid: {}", hdr.as_ref().unwrap().uid());
-            let previous_uid = hdr.as_ref().unwrap().unlock();
-            assert_eq!(
-                previous_uid, uid,
-                "free: unlocked a header with a differet uid: {}",
-                previous_uid
-            );
-            self.count.fetch_sub(1, Ordering::Relaxed);
-        }
+        let previous_uid = self.header_ptr(ix).unlock();
+        assert_eq!(
+            previous_uid, uid,
+            "free: unlocked a header with a differet uid: {}",
+            previous_uid
+        );
+        self.count.fetch_sub(1, Ordering::Relaxed);
     }
 
     pub fn get<T: Sized>(&self, ix: u64) -> &T {