src/lib.rs

@@ -14,12 +14,7 @@ pub trait Reclaim: Sized {
 unsafe fn reclaim_shim<T: Reclaim>(erased_ptr: NonNull<Root<Erased>>) {
     let ptr = erased_ptr.cast::<Root<T>>();
 
-    unsafe {
-        debug_assert_eq!((*ptr.as_ptr()).state.get(), STATE_SHARED_INIT);
-        (*ptr.as_ptr()).state.set(STATE_UNCLAIMED);
-    };
-
-    let value = unsafe { (*ptr.as_ptr()).target.assume_init_read().value };
+    let value = unsafe { ptr::addr_of!((*(*ptr.as_ptr()).target.as_ptr()).value).read() };
 
     let unclaimed = Unclaimed { ptr };
     value.reclaim(unclaimed);
@@ -48,7 +43,7 @@ pub struct Root<T> {
 impl<T> Drop for Root<T> {
     fn drop(&mut self) {
         let state = self.state.get();
-        if state <= STATE_SHARED_MAX {
+        if state != STATE_UNPINNED {
             struct Abort;
             impl Drop for Abort {
                 fn drop(&mut self) {
@@ -132,6 +127,15 @@ pub struct Unclaimed<T> {
     ptr: NonNull<Root<T>>,
 }
 
+impl<T> Drop for Unclaimed<T> {
+    fn drop(&mut self) {
+        unsafe {
+            debug_assert_eq!((*self.ptr.as_ptr()).state.get(), STATE_UNCLAIMED);
+            (*self.ptr.as_ptr()).state.set(STATE_UNPINNED);
+        }
+    }
+}
+
 impl<T> Unclaimed<T> {
     pub fn claim(self, value: T) -> Xrc<T> {
         unsafe {
@@ -147,6 +151,7 @@ impl<T> Unclaimed<T> {
         // shared read-only reborrow
         let target = unsafe { (*self.ptr.as_ptr()).target.assume_init_ref() };
         let ptr = NonNull::from(target);
+        mem::forget(self);
         Xrc { ptr }
     }
 }
@@ -169,6 +174,7 @@ impl<T: ?Sized> Drop for Xrc<T> {
             state_ref.set(state - 1);
             return;
         }
+        state_ref.set(STATE_UNCLAIMED);
 
         let reclaim_shim = unsafe { (*root.as_ptr()).reclaim_shim };
         unsafe {
@@ -280,7 +286,16 @@ mod tests {
 
     #[test]
     #[ignore = "will abort"]
-    fn abort() {
+    fn abort_unclaimed() {
+        #[allow(clippy::needless_late_init)]
+        let _unclaimed;
+        let root = pin!(Root::<NoReclaim<()>>::new());
+        _unclaimed = root.unclaimed();
+    }
+
+    #[test]
+    #[ignore = "will abort"]
+    fn abort_xrc() {
         #[allow(clippy::needless_late_init)]
         let _xrc;
         let root = pin!(Root::new());
@@ -288,10 +303,17 @@ mod tests {
     }
 
     #[test]
-    #[should_panic = "unclaimed called multiple times"]
-    fn multiple_unclaimed() {
-        let mut root = pin!(Root::<NoReclaim<i32>>::new());
+    fn multiple_unclaimed_separate() {
+        let mut root = pin!(Root::<NoReclaim<()>>::new());
+        let _ = root.as_mut().unclaimed();
         let _ = root.as_mut().unclaimed();
+    }
+
+    #[test]
+    #[should_panic = "unclaimed called multiple times"]
+    fn multiple_unclaimed_concurrent() {
+        let mut root = pin!(Root::<NoReclaim<()>>::new());
+        let _unclaimed = root.as_mut().unclaimed();
         let _ = root.as_mut().unclaimed();
     }