From e4103b5886185900b737fe11699c29b3e39f2639 Mon Sep 17 00:00:00 2001
From: Brooks Prumo <brooks@solana.com>
Date: Thu, 16 Sep 2021 16:00:30 -0500
Subject: [PATCH] resurect addref and unref since tests use 'em for now

---
 bucket_map/src/bucket.rs     | 12 ++++++++++++
 bucket_map/src/bucket_map.rs | 14 ++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/bucket_map/src/bucket.rs b/bucket_map/src/bucket.rs
index 4a9933f264..49d5e44ffa 100644
--- a/bucket_map/src/bucket.rs
+++ b/bucket_map/src/bucket.rs
@@ -161,6 +161,18 @@ impl<T: Clone + Copy> Bucket<T> {
         Err(BucketMapError::IndexNoSpace(index.capacity_pow2))
     }
 
+    pub fn addref(&mut self, key: &Pubkey) -> Option<RefCount> {
+        let (elem, _) = self.find_entry_mut(key)?;
+        elem.ref_count += 1;
+        Some(elem.ref_count)
+    }
+
+    pub fn unref(&mut self, key: &Pubkey) -> Option<RefCount> {
+        let (elem, _) = self.find_entry_mut(key)?;
+        elem.ref_count -= 1;
+        Some(elem.ref_count)
+    }
+
     fn create_key(&self, key: &Pubkey, ref_count: u64) -> Result<u64, BucketMapError> {
         Self::bucket_create_key(
             &self.index,
diff --git a/bucket_map/src/bucket_map.rs b/bucket_map/src/bucket_map.rs
index dacd5bf1bb..12a3da78a1 100644
--- a/bucket_map/src/bucket_map.rs
+++ b/bucket_map/src/bucket_map.rs
@@ -196,6 +196,20 @@ impl<T: Clone + Copy + Debug> BucketMap<T> {
             0
         }
     }
+
+    /// Increment the refcount for Pubkey `key`
+    pub fn addref(&self, key: &Pubkey) -> Option<RefCount> {
+        let ix = self.bucket_ix(key);
+        let mut bucket = self.buckets[ix].write().unwrap();
+        bucket.as_mut()?.addref(key)
+    }
+
+    /// Decrement the refcount for Pubkey `key`
+    pub fn unref(&self, key: &Pubkey) -> Option<RefCount> {
+        let ix = self.bucket_ix(key);
+        let mut bucket = self.buckets[ix].write().unwrap();
+        bucket.as_mut()?.unref(key)
+    }
 }
 
 /// Look at the first 8 bytes of the input and reinterpret them as a u64