feat: skip alocation for arc_t inputs (#77)

Author: filipdulic

src/flavors/arc_swap.rs

@@ -8,8 +8,8 @@ pub struct Slot<T> {
 }
 
 impl<T> SwapSlot<T> for Slot<T> {
-    fn store(&self, item: T) {
-        self.shared.store(Some(Arc::new(item)))
+    fn store(&self, item: impl Into<Arc<T>>) {
+        self.shared.store(Some(item.into()));
     }
 
     fn load(&self) -> Option<Arc<T>> {
@@ -70,3 +70,46 @@ mod test {
         assert_eq!(Arc::strong_count(&arc.unwrap()), 2)
     }
 }
+
+#[cfg(test)]
+mod allocation_tests {
+    use crate::flavors::allocation_tests::{allocs_current_thread, reset_allocs_current_thread};
+
+    use super::*;
+
+    #[test]
+    fn store_with_arc_does_not_allocate_new_arc() {
+        let slot = Slot::<u32>::none();
+        let arc = Arc::new(123u32);
+
+        // Ignore allocations from constructing `slot` and `arc`.
+        reset_allocs_current_thread();
+
+        // This should only move / clone the Arc; no new heap allocation for T.
+        slot.store(arc.clone());
+
+        // Might still be 0 or some tiny number depending on RwLock internals,
+        // but definitely shouldn't be "one Arc allocation vs another" difference.
+        let after = allocs_current_thread();
+        assert_eq!(
+            after, 0,
+            "expected no additional allocations when storing an Arc"
+        );
+    }
+
+    #[test]
+    fn store_with_value_allocates_arc() {
+        let slot = Slot::<u32>::none();
+
+        reset_allocs_current_thread();
+
+        // This goes through `impl From<T> for Arc<T>` and must allocate.
+        slot.store(5u32);
+
+        let after = allocs_current_thread();
+        assert!(
+            after == 1,
+            "expected at least one allocation when storing a bare value T"
+        );
+    }
+}

src/flavors/mod.rs

@@ -3,3 +3,42 @@ pub mod arc_swap;
 
 #[cfg(feature = "rwlock")]
 pub mod rw_lock;
+
+#[cfg(test)]
+mod allocation_tests {
+    use std::alloc::{GlobalAlloc, Layout, System};
+    use std::cell::Cell;
+
+    struct CountingAlloc;
+
+    #[global_allocator]
+    static GLOBAL: CountingAlloc = CountingAlloc;
+
+    thread_local! {
+        // Each OS thread gets its own counter.
+        static ALLOCS_THIS_THREAD: Cell<usize> = Cell::new(0);
+    }
+
+    unsafe impl GlobalAlloc for CountingAlloc {
+        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+            let ptr = System.alloc(layout);
+            if !ptr.is_null() {
+                // Only bump the counter for the *current* thread.
+                ALLOCS_THIS_THREAD.with(|c| c.set(c.get() + 1));
+            }
+            ptr
+        }
+
+        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+            System.dealloc(ptr, layout);
+        }
+    }
+
+    pub(crate) fn reset_allocs_current_thread() {
+        ALLOCS_THIS_THREAD.with(|c| c.set(0));
+    }
+
+    pub(crate) fn allocs_current_thread() -> usize {
+        ALLOCS_THIS_THREAD.with(|c| c.get())
+    }
+}

src/flavors/rw_lock.rs

@@ -7,8 +7,8 @@ pub struct Slot<T> {
 }
 
 impl<T> SwapSlot<T> for Slot<T> {
-    fn store(&self, item: T) {
-        *self.lock.write().unwrap() = Some(Arc::new(item));
+    fn store(&self, item: impl Into<Arc<T>>) {
+        *self.lock.write().unwrap() = Some(item.into());
     }
 
     fn load(&self) -> Option<Arc<T>> {
@@ -69,3 +69,46 @@ mod test {
         assert_eq!(Arc::strong_count(&arc.unwrap()), 2)
     }
 }
+
+#[cfg(test)]
+mod allocation_tests {
+    use crate::flavors::allocation_tests::{allocs_current_thread, reset_allocs_current_thread};
+
+    use super::*;
+
+    #[test]
+    fn store_with_arc_does_not_allocate_new_arc() {
+        let slot = Slot::<u32>::none();
+        let arc = Arc::new(123u32);
+
+        // Ignore allocations from constructing `slot` and `arc`.
+        reset_allocs_current_thread();
+
+        // This should only move / clone the Arc; no new heap allocation for T.
+        slot.store(arc.clone());
+
+        // Might still be 0 or some tiny number depending on RwLock internals,
+        // but definitely shouldn't be "one Arc allocation vs another" difference.
+        let after = allocs_current_thread();
+        assert_eq!(
+            after, 0,
+            "expected no additional allocations when storing an Arc"
+        );
+    }
+
+    #[test]
+    fn store_with_value_allocates_arc() {
+        let slot = Slot::<u32>::none();
+
+        reset_allocs_current_thread();
+
+        // This goes through `impl From<T> for Arc<T>` and must allocate.
+        slot.store(5u32);
+
+        let after = allocs_current_thread();
+        assert!(
+            after == 1,
+            "expected at least one allocation when storing a bare value T"
+        );
+    }
+}

src/swap_slot.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 pub trait SwapSlot<T> {
     /// Creates a new Arc around item and stores it,
     /// dropping the previously held item's Arc.
-    fn store(&self, item: T);
+    fn store(&self, item: impl Into<Arc<T>>);
 
     /// Returns a clone of the held Arc,
     /// incrementing the ref count atomically