1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
//! Macros, attributes, and traits for invasively reference-counted structs in Rust.
//!
//! This crate is centered around manipulating invasively reference counted
//! structs. These structs are declared using the `#[refcounted]` attribute,
//! constructed with the [`make_refptr`] macro, and have their lifetimes managed
//! using the [`RefPtr`] and [`WeakPtr`] smart pointer types.
//!
//! # Declaring a refcounted struct
//!
//! The `#[refcounted]` attribute can be applied to a `struct` declaration to
//! mark it as refcounted. Refcounted structs are always allocated on the heap,
//! and are constructed using the `make_refptr` helper macro.
//!
//! ## Example
//!
//! ```
//! # use refptr::*;
//! # use std::cell::Cell;
//! #[refcounted(local)]
//! struct HeapInteger {
//!     value: Cell<i32>,
//! }
//!
//! let orig = make_refptr!(HeapInteger { value: Cell::new(10) });
//! let copy = orig.clone();
//! orig.value.set(20);
//! assert_eq!(copy.value.get(), 20);
//! ```
//!
//! # Allocating
//!
//! Structs declared with `#[refcounted]` are constructed on the heap using the
//! [`make_refptr!`] macro. This macro accepts struct literal syntax, but
//! constructs the value onto the heap.
//!
//! This is required in order to ensure that the type always lives on the heap
//! for invasive reference counting.
//!
//! ## Example
//!
//! ```
//! # use refptr::*;
//! # #[refcounted(local)] struct HeapPair<T, U> { t: T, u: U }
//! let ptr = make_refptr!(HeapPair { t: 10, u: 20 });
//! assert_eq!(ptr.t, 10);
//! assert_eq!(ptr.u, 20);
//! ```
//!
//! # Finalization and `Drop`
//!
//! Types annotated with `#[refcounted]` cannot manually implement `Drop`, as it
//! would allow recovering a `RefPtr<Self>` while the object is being dropped,
//! leading to a use-after-free.
//!
//! If a finalization method is needed, the `#[refcounted(finalize)]` attribute
//! provides support for custom finalization. If finalization is enabled, a `fn
//! finalize(&self)` method is called before dropping any fields.
//!
//! It is possible for code to acquire a new strong reference during the
//! `finalize` method, which may cause the struct to not be dropped after it
//! returns. Because of this, `finalize` may be called on the same struct
//! multiple times over it's lifetime.
//!
//! # Configuration
//!
//! ## `#[refcounted(atomic)]` and `#[refcounted(local)]`
//!
//! Select between atomic reference counting, like [`Arc`], or thread local
//! reference counting, like [`Rc`]. Atomically refcounted types may be shared
//! between threads, so long as all fields are also sharable.
//!
//! The atomicity of the refcount must be specified.
//!
//! ### Example
//!
//! ```
//! # use refptr::*;
//! # use std::thread;
//! #[refcounted(atomic)]
//! struct HeapInt { i: i32 }
//!
//! let here = make_refptr!(HeapInt { i: 10 });
//! let thread = thread::spawn(move || here.i);
//! assert_eq!(thread.join().unwrap(), 10);
//! ```
//!
//! [`Arc`]: alloc::sync::Arc
//! [`Rc`]: alloc::rc::Rc
//!
//! ## `#[refcounted(weak)]`
//!
//! Adds support for weak reference counts and the [`WeakPtr`] smart pointer
//! type. This annotation may be combined with other annotations.
//!
//! ### Example
//!
//! ```
//! # use refptr::*;
//! # use std::thread;
//! #[refcounted(atomic, weak)]
//! struct HeapInt { i: i32 }
//!
//! let here = make_refptr!(HeapInt { i: 10 });
//! let weak = WeakPtr::new(&*here);
//! assert_eq!(weak.upgrade().unwrap().i, 10);
//! drop(here);
//! assert!(weak.upgrade().is_none());
//! ```
//!
//! ## `#[refcounted(finalize)]`
//!
//! Calls a `fn finalize(&self)` method on the struct before attempting to
//! destroy it. See the "Finalization" section for more details. This annotation
//! may be combined with other annotations.
//!
//! Structs which support being referenced using [`RefPtr`] are annotated with the
//! `#[refcounted(...)]` attribute. This attribute generates the necessary unsafe
//! code, extra members, and trait implementations required.
//!
//! ```
//! # use refptr::*;
//! # use std::sync::atomic::{AtomicBool, Ordering::SeqCst};
//! #[refcounted(atomic, finalize)]
//! struct FinalizeExample {}
//!
//! static FINALIZED: AtomicBool = AtomicBool::new(false);
//! impl FinalizeExample {
//!     fn finalize(&self) {
//!         FINALIZED.store(true, SeqCst);
//!     }
//! }
//!
//! let orig = make_refptr!(FinalizeExample {});
//! assert_eq!(FINALIZED.load(SeqCst), false);
//! let copy = orig.clone();
//! assert_eq!(FINALIZED.load(SeqCst), false);
//! drop(orig);
//! assert_eq!(FINALIZED.load(SeqCst), false);
//! drop(copy);
//! assert_eq!(FINALIZED.load(SeqCst), true);
//! ```
//!
//! # Trait Objects
//!
//! `#[refcounted]` can also be used for managing the lifecycles of trait
//! objects, by including the `Refcounted` trait in your trait object's
//! hierarchy. The `Rc` associated type will need to be specified in order to
//! maintain object safety.
//!
//! The [`refcnt`] module contains the specific reference count types used by
//! this crate.
//!
//! ## Example
//!
//! ```
//! # use refptr::*;
//! trait MyTrait : Refcounted<Rc = refcnt::AtomicWeak> {
//!     fn my_trait_method(&self) -> i32;
//! }
//!
//! #[refcounted(atomic, weak)]
//! struct MyStruct { i: i32 }
//!
//! impl MyTrait for MyStruct {
//!     fn my_trait_method(&self) -> i32 { self.i }
//! }
//!
//! fn takes_trait_object(obj: &dyn MyTrait) -> i32 {
//!     let strong_ref: RefPtr<dyn MyTrait> = RefPtr::new(obj);
//!     strong_ref.my_trait_method()
//! }
//!
//! let concrete = make_refptr!(MyStruct { i: 10 });
//! let i = takes_trait_object(&*concrete);
//! assert_eq!(i, 10);
//! ```

#![no_std]

extern crate alloc;

use core::cmp::Ordering;
use core::fmt;
use core::hash::{Hash, Hasher};
use core::marker::PhantomData;
use core::mem;
use core::ops::Deref;
use core::ptr::NonNull;

pub mod refcnt;
use refcnt::{Inner, Refcount, WeakRefcount};

// Not public API.
#[doc(hidden)]
#[path = "runtime.rs"]
pub mod __rt;

/// Attribute for declaring [`Refcounted`] structs.
///
/// See the [module level documentation](self) for usage.
pub use refptr_macros::refcounted;

/// An invasively reference counted type.
///
/// Objects implementing this trait are always allocated on the heap, and have
/// their lifecycle managed using the [`RefPtr`] smart pointer.
///
/// ## Safety
///
///  * `Refcounted` objects are always heap-allocated
///  * Only shared references may exist to `Refcounted` objects
pub unsafe trait Refcounted {
    /// Reference count used by this type.
    type Rc: Refcount;

    /// Metadata used internally by `Refcount` implementations.
    ///
    /// This metadata can be used to implement dynamic extensions to the
    /// refcount type, such as `finalize` support or RTTI.
    unsafe fn refcount_metadata(&self) -> <Self::Rc as Refcount>::Metadata;
}

/// Strong reference to a [`Refcounted`] object.
pub struct RefPtr<T: ?Sized + Refcounted> {
    ptr: NonNull<Inner<T>>,
    _marker: PhantomData<T>,
}

impl<T: ?Sized + Refcounted> RefPtr<T> {
    /// Obtain a strong reference to a `Refcounted` object.
    pub fn new(val: &T) -> RefPtr<T> {
        unsafe {
            let ptr = Inner::cast(val as *const T as *mut T);
            T::Rc::inc_strong(ptr);
            RefPtr::from_inner(ptr)
        }
    }

    /// Recover a `RefPtr` from a raw pointer which was previously returned from
    /// `into_raw`. This does not increment the reference count.
    pub unsafe fn from_raw(val: *const T) -> RefPtr<T> {
        RefPtr::from_inner(Inner::cast(val as *mut T))
    }

    /// Acquire a raw pointer to the allocation, consuming the `RefPtr`.
    pub fn into_raw(this: Self) -> *const T {
        let ptr = this.deref() as *const T;
        mem::forget(this);
        ptr
    }

    unsafe fn from_inner(ptr: *mut Inner<T>) -> RefPtr<T> {
        RefPtr {
            ptr: NonNull::new_unchecked(ptr),
            _marker: PhantomData,
        }
    }
}

impl<T: ?Sized + Refcounted> Deref for RefPtr<T> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &(*self.ptr.as_ptr()).data }
    }
}

impl<T: ?Sized + Refcounted> Clone for RefPtr<T> {
    fn clone(&self) -> Self {
        unsafe {
            T::Rc::inc_strong(self.ptr.as_ptr());
            RefPtr::from_inner(self.ptr.as_ptr())
        }
    }
}

impl<T: ?Sized + Refcounted> Drop for RefPtr<T> {
    fn drop(&mut self) {
        unsafe { T::Rc::dec_strong(self.ptr.as_ptr()) }
    }
}

unsafe impl<T: ?Sized + Refcounted + Sync + Send> Send for RefPtr<T> {}
unsafe impl<T: ?Sized + Refcounted + Sync + Send> Sync for RefPtr<T> {}

/// Weak reference to a [`Refcounted`] object.
///
/// Weak pointers can only be used on objects which have refcounts supporting
/// weak references.
///
/// # Example
///
/// ```
/// # use refptr::*;
/// # use std::thread;
/// #[refcounted(atomic, weak)]
/// struct HeapInt { i: i32 }
///
/// let here = make_refptr!(HeapInt { i: 10 });
/// let weak = WeakPtr::new(&*here);
/// assert_eq!(weak.upgrade().unwrap().i, 10);
/// drop(here);
/// assert!(weak.upgrade().is_none());
/// ```
pub struct WeakPtr<T: ?Sized>
where
    T: Refcounted,
    T::Rc: WeakRefcount,
{
    ptr: NonNull<Inner<T>>,
    _marker: PhantomData<T>,
}

impl<T: ?Sized> WeakPtr<T>
where
    T: Refcounted,
    T::Rc: WeakRefcount,
{
    /// Obtain a new weak reference to a refcounted object.
    pub fn new(val: &T) -> WeakPtr<T> {
        unsafe {
            let ptr = Inner::cast(val as *const T as *mut T);
            T::Rc::inc_weak(ptr);
            WeakPtr::from_inner(ptr)
        }
    }

    /// Attempt to upgrade this weak reference into a strong reference,
    /// returning it.
    pub fn upgrade(&self) -> Option<RefPtr<T>> {
        unsafe {
            if T::Rc::upgrade(self.ptr.as_ptr()) {
                Some(RefPtr::from_inner(self.ptr.as_ptr()))
            } else {
                None
            }
        }
    }

    /// Gets the number of strong references to this allocation.
    pub fn strong_count(&self) -> usize {
        unsafe { T::Rc::strong_count(self.ptr.as_ptr()) }
    }

    /// Gets the number of weak references to this allocation.
    ///
    /// If there are no remaining strong references, this will
    /// return `0`.
    pub fn weak_count(&self) -> usize {
        unsafe { T::Rc::weak_count(self.ptr.as_ptr()) }
    }

    unsafe fn from_inner(ptr: *mut Inner<T>) -> WeakPtr<T> {
        WeakPtr {
            ptr: NonNull::new_unchecked(ptr),
            _marker: PhantomData,
        }
    }
}

impl<T: ?Sized> Clone for WeakPtr<T>
where
    T: Refcounted,
    T::Rc: WeakRefcount,
{
    fn clone(&self) -> Self {
        unsafe {
            T::Rc::inc_weak(self.ptr.as_ptr());
            WeakPtr::from_inner(self.ptr.as_ptr())
        }
    }
}

impl<T: ?Sized> Drop for WeakPtr<T>
where
    T: Refcounted,
    T::Rc: WeakRefcount,
{
    fn drop(&mut self) {
        unsafe { T::Rc::dec_weak(self.ptr.as_ptr()) }
    }
}

impl<T: ?Sized> fmt::Debug for WeakPtr<T>
where
    T: Refcounted,
    T::Rc: WeakRefcount,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "(WeakPtr)")
    }
}

unsafe impl<T: ?Sized> Send for WeakPtr<T>
where
    T: Refcounted + Send + Sync,
    T::Rc: WeakRefcount,
{
}
unsafe impl<T: ?Sized> Sync for WeakPtr<T>
where
    T: Refcounted + Send + Sync,
    T::Rc: WeakRefcount,
{
}

// Trait impls for `RefPtr<T>`

impl<T: ?Sized + Refcounted + PartialEq> PartialEq for RefPtr<T> {
    fn eq(&self, other: &RefPtr<T>) -> bool {
        **self == **other
    }
}

impl<T: ?Sized + Refcounted + PartialOrd> PartialOrd for RefPtr<T> {
    fn partial_cmp(&self, other: &RefPtr<T>) -> Option<Ordering> {
        (**self).partial_cmp(&**other)
    }
    fn lt(&self, other: &RefPtr<T>) -> bool {
        *(*self) < *(*other)
    }
    fn le(&self, other: &RefPtr<T>) -> bool {
        *(*self) <= *(*other)
    }
    fn gt(&self, other: &RefPtr<T>) -> bool {
        *(*self) > *(*other)
    }
    fn ge(&self, other: &RefPtr<T>) -> bool {
        *(*self) >= *(*other)
    }
}

impl<T: ?Sized + Refcounted + Ord> Ord for RefPtr<T> {
    fn cmp(&self, other: &RefPtr<T>) -> Ordering {
        (**self).cmp(&**other)
    }
}

impl<T: ?Sized + Refcounted + Eq> Eq for RefPtr<T> {}

impl<T: ?Sized + Refcounted + fmt::Display> fmt::Display for RefPtr<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&**self, f)
    }
}

impl<T: ?Sized + Refcounted + fmt::Debug> fmt::Debug for RefPtr<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&**self, f)
    }
}

impl<T: ?Sized + Refcounted> fmt::Pointer for RefPtr<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Pointer::fmt(&(&**self as *const T), f)
    }
}

impl<T: ?Sized + Refcounted + Hash> Hash for RefPtr<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        (**self).hash(state)
    }
}

impl<T: ?Sized + Refcounted> From<&T> for RefPtr<T> {
    fn from(v: &T) -> Self {
        RefPtr::new(v)
    }
}

/// Allocate a new instance of a [`Refcounted`] struct using a struct literal.
///
/// Returns a `RefPtr<T>` strong reference to the newly allocated struct.
///
/// # Example
///
/// ```
/// # use refptr::*;
/// #[refcounted(local)]
/// struct HeapInt { value: i32 }
///
/// let ptr = make_refptr!(HeapInt { value: 10 });
/// ```
#[macro_export]
macro_rules! make_refptr {
    ($($seg:ident $(::<$($t:ty),*>)?)::+ { $($f:tt)* }) => {
        {
            let value = $crate::__rt::ManuallyDrop::new($($seg $(::<$($t),*>)?)::+ {
                _refcnt_marker: unsafe { $crate::__rt::PhantomRefcnt::new() },
                $($f)*
            });

            unsafe { $crate::__rt::alloc(value) }
        }
    }
}