ref_count: add weak count

Needed for weakref support.

Author: quark-zju

src/ref_count.rs

@@ -33,14 +33,19 @@ pub trait RefCount: 'static {
     ) -> Option<parking_lot::lock_api::RwLockReadGuard<'_, parking_lot::RawRwLock, ()>> {
         None
     }
+
+    // Weakref support.
+    fn inc_weak(&self) -> usize;
+    fn dec_weak(&self) -> usize;
+    fn weak_count(&self) -> usize;
 }
 
-pub struct SingleThreadRefCount(Cell<usize>);
+pub struct SingleThreadRefCount(Cell<usize>, Cell<usize>);
 
 impl SingleThreadRefCount {
     pub fn new(tracked: bool) -> Self {
         let value = (1 << REF_COUNT_SHIFT) | if tracked { REF_COUNT_MASK_TRACKED } else { 0 };
-        Self(Cell::new(value))
+        Self(Cell::new(value), Cell::new(0))
     }
 }
 
@@ -80,4 +85,23 @@ impl RefCount for SingleThreadRefCount {
         self.0.set(value - (1 << REF_COUNT_SHIFT));
         value >> REF_COUNT_SHIFT
     }
+
+    #[inline]
+    fn inc_weak(&self) -> usize {
+        let value = Cell::get(&self.1);
+        self.1.set(value + 1);
+        value
+    }
+
+    #[inline]
+    fn dec_weak(&self) -> usize {
+        let value = Cell::get(&self.1);
+        self.1.set(value - 1);
+        value
+    }
+
+    #[inline]
+    fn weak_count(&self) -> usize {
+        Cell::get(&self.1)
+    }
 }

src/sync/ref_count.rs

@@ -8,6 +8,7 @@ use std::sync::Arc;
 
 pub struct ThreadedRefCount {
     ref_count: AtomicUsize,
+    weak_count: AtomicUsize,
     pub(crate) collector_lock: Arc<RwLock<()>>,
 }
 
@@ -19,6 +20,7 @@ impl ThreadedRefCount {
             ref_count: AtomicUsize::new(
                 (1 << REF_COUNT_SHIFT) | if tracked { REF_COUNT_MASK_TRACKED } else { 0 },
             ),
+            weak_count: AtomicUsize::new(0),
         }
     }
 }
@@ -59,4 +61,19 @@ impl RefCount for ThreadedRefCount {
     fn locked(&self) -> Option<RwLockReadGuard<'_, RawRwLock, ()>> {
         Some(self.collector_lock.read_recursive())
     }
+
+    #[inline]
+    fn inc_weak(&self) -> usize {
+        self.weak_count.fetch_add(1, AcqRel)
+    }
+
+    #[inline]
+    fn dec_weak(&self) -> usize {
+        self.weak_count.fetch_sub(1, AcqRel)
+    }
+
+    #[inline]
+    fn weak_count(&self) -> usize {
+        self.weak_count.load(Acquire)
+    }
 }