add occupy_if_matches to bucket map (#33318)

bucket_map/src/bucket.rs

@@ -756,7 +756,7 @@ impl<'b, T: Clone + Copy + PartialEq + std::fmt::Debug + 'static> Bucket<T> {
 
 #[cfg(test)]
 mod tests {
-    use {super::*, tempfile::tempdir};
+    use {super::*, crate::index_entry::OccupyIfMatches, tempfile::tempdir};
 
     #[test]
     fn test_index_entries() {
@@ -950,6 +950,116 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_occupy_if_matches() {
+        let random = 1;
+        let k = Pubkey::from([1u8; 32]);
+        let k2 = Pubkey::from([2u8; 32]);
+        let v = 12u64;
+        let v2 = 13u64;
+        let raw = vec![(k, v)];
+        let hashed = Bucket::index_entries(&raw, random);
+        let hashed_raw = hashed.clone();
+
+        let mut index = create_test_index(None);
+
+        let single_hashed_raw_inserted = hashed_raw.last().unwrap();
+        let elem = IndexEntryPlaceInBucket::new(single_hashed_raw_inserted.0 % index.capacity());
+
+        assert_eq!(elem.get_slot_count_enum(&index), OccupiedEnum::Free);
+        elem.init(&mut index, &k);
+        elem.set_slot_count_enum_value(&mut index, OccupiedEnum::OneSlotInIndex(&v));
+        assert_eq!(
+            elem.get_slot_count_enum(&index),
+            OccupiedEnum::OneSlotInIndex(&v)
+        );
+        // clear it
+        elem.set_slot_count_enum_value(&mut index, OccupiedEnum::Free);
+        assert_eq!(elem.get_slot_count_enum(&index), OccupiedEnum::Free);
+
+        assert_eq!(
+            elem.occupy_if_matches(&mut index, &v, &k),
+            OccupyIfMatches::SuccessfulInit
+        );
+        assert_eq!(
+            elem.get_slot_count_enum(&index),
+            OccupiedEnum::OneSlotInIndex(&v)
+        );
+        // clear it
+        elem.set_slot_count_enum_value(&mut index, OccupiedEnum::Free);
+        assert_eq!(elem.get_slot_count_enum(&index), OccupiedEnum::Free);
+
+        // v2 but will still write it
+        assert_eq!(
+            elem.occupy_if_matches(&mut index, &v2, &k),
+            OccupyIfMatches::SuccessfulInit
+        );
+        assert_eq!(
+            elem.get_slot_count_enum(&index),
+            OccupiedEnum::OneSlotInIndex(&v2)
+        );
+
+        // already a different occupied value for this pubkey in the index, so found duplicate
+        assert_eq!(
+            elem.occupy_if_matches(&mut index, &v, &k),
+            OccupyIfMatches::FoundDuplicate
+        );
+        assert_eq!(
+            elem.get_slot_count_enum(&index),
+            OccupiedEnum::OneSlotInIndex(&v2)
+        );
+
+        // k2 is pubkey mismatch
+        assert_eq!(
+            elem.occupy_if_matches(&mut index, &v, &k2),
+            OccupyIfMatches::PubkeyMismatch
+        );
+
+        // clear it
+        elem.set_slot_count_enum_value(&mut index, OccupiedEnum::Free);
+        assert_eq!(elem.get_slot_count_enum(&index), OccupiedEnum::Free);
+
+        // k2 is pubkey mismatch
+        assert_eq!(
+            elem.occupy_if_matches(&mut index, &v, &k2),
+            OccupyIfMatches::PubkeyMismatch
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "index asked to insert the same data twice")]
+    fn test_occupy_if_matches_panic() {
+        solana_logger::setup();
+        let random = 1;
+        let k = Pubkey::from([1u8; 32]);
+        let v = 12u64;
+        let raw = vec![(k, v)];
+        let hashed = Bucket::index_entries(&raw, random);
+        let hashed_raw = hashed.clone();
+
+        let mut index = create_test_index(None);
+
+        let single_hashed_raw_inserted = hashed_raw.last().unwrap();
+        let elem = IndexEntryPlaceInBucket::new(single_hashed_raw_inserted.0 % index.capacity());
+
+        assert_eq!(elem.get_slot_count_enum(&index), OccupiedEnum::Free);
+        elem.init(&mut index, &k);
+        elem.set_slot_count_enum_value(&mut index, OccupiedEnum::OneSlotInIndex(&v));
+        assert_eq!(
+            elem.get_slot_count_enum(&index),
+            OccupiedEnum::OneSlotInIndex(&v)
+        );
+
+        assert_eq!(
+            elem.occupy_if_matches(&mut index, &v, &k),
+            OccupyIfMatches::SuccessfulInit
+        );
+        assert_eq!(
+            elem.get_slot_count_enum(&index),
+            OccupiedEnum::OneSlotInIndex(&v)
+        );
+    }
+
     #[should_panic(expected = "batch insertion can only occur prior to any deletes")]
     #[test]
     fn batch_insert_after_delete() {

bucket_map/src/index_entry.rs

@@ -36,6 +36,16 @@ struct DataBucketRefCountOccupiedHeader {
     packed_ref_count: PackedRefCount,
 }
 
+#[derive(Debug, PartialEq)]
+pub enum OccupyIfMatches {
+    /// this entry is occupied and contains a pubkey with a different value, so this entry could not be updated
+    FoundDuplicate,
+    /// this entry was free and contains this pubkey and either value matched or the value was written to match
+    SuccessfulInit,
+    /// this entry had a different pubkey
+    PubkeyMismatch,
+}
+
 /// allocated in `contents` in a BucketStorage
 #[derive(Copy, Clone)]
 #[repr(C)]
@@ -279,7 +289,7 @@ pub(crate) union SingleElementOrMultipleSlots<T: Clone + Copy> {
 
 /// just the values for `OccupiedEnum`
 /// This excludes the contents of any enum value.
-#[derive(PartialEq, FromPrimitive)]
+#[derive(PartialEq, FromPrimitive, Debug)]
 #[repr(u8)]
 enum OccupiedEnumTag {
     #[default]
@@ -377,6 +387,39 @@ impl<T: Copy + PartialEq + 'static> IndexEntryPlaceInBucket<T> {
         index_entry.key = *pubkey;
     }
 
+    /// If the entry matches the pubkey and is unoccupied, then store `data` here and occupy the entry.
+    pub(crate) fn occupy_if_matches(
+        &self,
+        index_bucket: &mut BucketStorage<IndexBucket<T>>,
+        data: &T,
+        k: &Pubkey,
+    ) -> OccupyIfMatches {
+        let index_entry = index_bucket.get::<IndexEntry<T>>(self.ix);
+        if &index_entry.key == k {
+            let enum_tag = index_bucket.contents.get_enum_tag(self.ix);
+            if unsafe { &index_entry.contents.single_element } == data {
+                assert_eq!(
+                    enum_tag,
+                    OccupiedEnumTag::Free,
+                    "index asked to insert the same data twice"
+                );
+                index_bucket
+                    .contents
+                    .set_enum_tag(self.ix, OccupiedEnumTag::OneSlotInIndex);
+                OccupyIfMatches::SuccessfulInit
+            } else if enum_tag == OccupiedEnumTag::Free {
+                // pubkey is same, but value is different, so update value
+                self.set_slot_count_enum_value(index_bucket, OccupiedEnum::OneSlotInIndex(data));
+                OccupyIfMatches::SuccessfulInit
+            } else {
+                // found occupied duplicate of this pubkey
+                OccupyIfMatches::FoundDuplicate
+            }
+        } else {
+            OccupyIfMatches::PubkeyMismatch
+        }
+    }
+
     pub(crate) fn read_value<'a>(
         &self,
         index_bucket: &'a BucketStorage<IndexBucket<T>>,