From fc2bcdffe239e88f1f0272a9818661de11fa8e15 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Fri, 31 Mar 2023 14:37:21 -0500
Subject: [PATCH] fix up refcounts on disk bucket test (#31001)

---
 bucket_map/src/bucket_map.rs  | 14 +++++++++++++-
 bucket_map/src/index_entry.rs |  3 +++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/bucket_map/src/bucket_map.rs b/bucket_map/src/bucket_map.rs
index 309eabf37..478f5500a 100644
--- a/bucket_map/src/bucket_map.rs
+++ b/bucket_map/src/bucket_map.rs
@@ -170,6 +170,7 @@ fn read_be_u64(input: &[u8]) -> u64 {
 mod tests {
     use {
         super::*,
+        crate::index_entry::MAX_LEGAL_REFCOUNT,
         rand::{thread_rng, Rng},
         std::{collections::HashMap, sync::RwLock},
     };
@@ -365,7 +366,18 @@ mod tests {
             let v = (0..count)
                 .map(|x| (x as usize, x as usize /*thread_rng().gen::<usize>()*/))
                 .collect::<Vec<_>>();
-            let rc = thread_rng().gen_range(0, RefCount::MAX >> 2);
+            let range = thread_rng().gen_range(0, 100);
+            // pick ref counts that are useful and common
+            let rc = if range < 50 {
+                1
+            } else if range < 60 {
+                0
+            } else if range < 70 {
+                2
+            } else {
+                thread_rng().gen_range(0, MAX_LEGAL_REFCOUNT)
+            };
+
             (v, rc)
         };
 
diff --git a/bucket_map/src/index_entry.rs b/bucket_map/src/index_entry.rs
index a784bd2fa..1a68bb822 100644
--- a/bucket_map/src/index_entry.rs
+++ b/bucket_map/src/index_entry.rs
@@ -93,6 +93,9 @@ pub struct IndexEntry<T: 'static> {
     _phantom: PhantomData<&'static T>,
 }
 
+/// 62 bits available for ref count
+pub(crate) const MAX_LEGAL_REFCOUNT: RefCount = RefCount::MAX >> 2;
+
 /// hold a big `RefCount` while leaving room for extra bits to be used for things like 'Occupied'
 #[bitfield(bits = 64)]
 #[repr(C)]