Merge branch 'main' into kdy1/dev-3

crates/hstr/Cargo.toml

@@ -5,7 +5,7 @@ edition     = "2021"
 license     = "Apache-2.0"
 name        = "hstr"
 repository  = "https://github.com/dudykr/ddbase.git"
-version     = "0.2.9"
+version     = "0.2.10"
 
 [lib]
 bench = false

crates/hstr/src/dynamic.rs

@@ -4,7 +4,10 @@ use std::{
     hash::{BuildHasherDefault, Hash, Hasher},
     num::NonZeroU32,
     ptr::NonNull,
-    sync::atomic::{AtomicU32, Ordering::SeqCst},
+    sync::atomic::{
+        AtomicU32,
+        Ordering::{self, SeqCst},
+    },
 };
 
 use rustc_hash::FxHasher;
@@ -23,6 +26,16 @@ pub(crate) struct Entry {
     pub alias: AtomicTaggedValue,
 }
 
+impl Drop for Entry {
+    fn drop(&mut self) {
+        if let Some(old) = self.alias.load(Ordering::SeqCst) {
+            unsafe {
+                Self::restore_arc(old);
+            }
+        }
+    }
+}
+
 impl Entry {
     pub unsafe fn cast(ptr: TaggedValue) -> *const Entry {
         ptr.get_ptr().cast()
@@ -32,6 +45,14 @@ impl Entry {
         &*Self::cast(ptr)
     }
 
+    #[cfg(test)]
+    pub unsafe fn ref_count(ptr: TaggedValue) -> usize {
+        let arc = Arc::from_raw(ptr.get_ptr() as *const Entry);
+        let count = Arc::count(&arc);
+        std::mem::forget(arc);
+        count
+    }
+
     pub unsafe fn restore_arc(v: TaggedValue) -> Arc<Entry> {
         let ptr = v.get_ptr() as *const Entry;
         Arc::from_raw(ptr)
@@ -83,9 +104,13 @@ impl AtomStore {
         for entry in other.data.keys() {
             let cur_entry = self.insert_entry(Cow::Borrowed(&entry.string), entry.hash);
 
-            let ptr = unsafe { NonNull::new_unchecked(Arc::as_ptr(&cur_entry) as *mut Entry) };
+            let ptr =
+                unsafe { NonNull::new_unchecked(Arc::into_raw(cur_entry.clone()) as *mut Entry) };
 
-            entry.alias.store(TaggedValue::new_ptr(ptr), SeqCst);
+            let old = entry.alias.swap(TaggedValue::new_ptr(ptr), SeqCst);
+            if let Some(old) = old {
+                unsafe { Entry::restore_arc(old) };
+            }
         }
     }
 

crates/hstr/src/lib.rs

@@ -261,6 +261,15 @@ impl Atom {
         }
     }
 
+    #[cfg(test)]
+    fn ref_count(&self) -> usize {
+        if self.tag() == DYNAMIC_TAG {
+            unsafe { Entry::ref_count(self.unsafe_data) }
+        } else {
+            0
+        }
+    }
+
     #[inline(always)]
     fn simple_eq(&self, other: &Self) -> Option<bool> {
         if self.unsafe_data == other.unsafe_data {

crates/hstr/src/tagged_value.rs

@@ -172,10 +172,15 @@ impl AtomicTaggedValue {
         unsafe { std::mem::transmute(self.value.load(ordering)) }
     }
 
-    pub fn store(&self, value: TaggedValue, ordering: std::sync::atomic::Ordering) {
+    #[must_use]
+    pub fn swap(
+        &self,
+        value: TaggedValue,
+        ordering: std::sync::atomic::Ordering,
+    ) -> Option<TaggedValue> {
         let value = Some(value);
         // The niche guarantees that Option<TaggedValue> has the same layout as the
         // atomic
-        unsafe { self.value.store(std::mem::transmute(value), ordering) }
+        unsafe { std::mem::transmute(self.value.swap(std::mem::transmute(value), ordering)) }
     }
 }

crates/hstr/src/tests.rs

@@ -21,6 +21,35 @@ fn simple_usage() {
     assert_eq!(a1.unsafe_data, a2.unsafe_data);
 }
 
+#[test]
+fn store_merge_and_drop() {
+    let (mut s1, atoms1) = store_with_atoms(vec!["Hello, world!!!!"]);
+    let a1 = atoms1.into_iter().next().unwrap();
+    assert_eq!(2, a1.ref_count());
+    let (mut s2, atoms2) = store_with_atoms(vec!["Hello, world!!!!"]);
+    let a2 = atoms2.into_iter().next().unwrap();
+    let (s3, atoms3) = store_with_atoms(vec!["Hello, world!!!!"]);
+    let a3 = atoms3.into_iter().next().unwrap();
+    s2.merge(s3);
+    assert_eq!(2, a1.ref_count());
+    s1.merge(s2);
+    assert_eq!(3, a1.ref_count());
+
+    assert_eq!("Hello, world!!!!", a1.as_str());
+    assert_eq!("Hello, world!!!!", a2.as_str());
+    assert_eq!("Hello, world!!!!", a3.as_str());
+
+    drop(s1);
+    assert_eq!(2, a1.ref_count());
+
+    assert_eq!("Hello, world!!!!", a1.as_str());
+    assert_eq!(a1, a2);
+    drop(a1);
+    assert_eq!(a2, a3);
+    assert_eq!("Hello, world!!!!", a2.as_str());
+    assert_eq!("Hello, world!!!!", a3.as_str());
+}
+
 #[test]
 fn eager_drop() {
     let (_, atoms1) = store_with_atoms(vec!["Hello, world!!!!"]);