Skip to content

Commit d50cfb8

Browse files
EFanZhEFanZh
committed
Add RefCounts and RcLayout types
1 parent fefce3c commit d50cfb8

File tree

3 files changed

+258
-0
lines changed

3 files changed

+258
-0
lines changed

‎library/alloc/src/lib.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,8 @@
206206
#[macro_use]
207207
mod macros;
208208

209+
#[cfg(not(no_rc))]
210+
mod raw_rc;
209211
mod raw_vec;
210212

211213
// Heaps provided for low-level allocation strategies

‎library/alloc/src/raw_rc/mod.rs‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
//! Base implementation for `rc::{Rc, UniqueRc, Weak}` and `sync::{Arc, UniqueArc, Weak}`.
2+
//!
3+
//! # Allocation Memory Layout
4+
//!
5+
//! The memory layout of a reference-counted allocation is designed so that the memory that stores
6+
//! the reference counts has a fixed offset to the memory that stores the value. In this way,
7+
//! operations that only rely on reference counts can ignore the actual type of the contained value
8+
//! and only care about the address of the contained value, which allows us to share code between
9+
//! reference-counting pointers that have different types of contained values. This can potentially
10+
//! reduce the binary size.
11+
//!
12+
//! Assume the type of the stored value is `T`, the allocation memory layout is designed as follows:
13+
//!
14+
//! - We use a `RefCounts` type to store the reference counts.
15+
//! - The alignment of the allocation is `align_of::<RefCounts>().max(align_of::<T>())`.
16+
//! - The value is stored at offset `size_of::<RefCounts>().next_multiple_of(align_of::<T>())`.
17+
//! - The size of the allocation is
18+
//! `size_of::<RefCounts>().next_multiple_of(align_of::<T>()) + size_of::<T>()`.
19+
//! - The `RefCounts` object is stored at offset
20+
//! `size_of::<RefCounts>().next_multiple_of(align_of::<T>()) - size_of::<RefCounts>()`.
21+
//!
22+
//! Here is a table showing the order and size of each component in an reference counted allocation
23+
//! of a `T` value:
24+
//!
25+
//! | Component | Size |
26+
//! | ----------- | ----------------------------------------------------------------------------------- |
27+
//! | Padding | `size_of::<RefCounts>().next_multiple_of(align_of::<T>()) - size_of::<RefCounts>()` |
28+
//! | `RefCounts` | `size_of::<RefCounts>()` |
29+
//! | `T` | `size_of::<T>()` |
30+
//!
31+
//! This works because:
32+
//!
33+
//! - Both `RefCounts` and the object is stored in the allocation without overlapping.
34+
//! - The `RefCounts` object is stored at offset
35+
//! `size_of::<RefCounts>().next_multiple_of(align_of::<T>()) - size_of::<RefCounts>()`, which
36+
//! has a valid alignment for `RefCounts` because:
37+
//! - If `align_of::<T>() <= align_of::<RefCounts>()`, we have the offset being 0, which has a
38+
//! valid alignment for `RefCounts`.
39+
//! - If `align_of::<T>() > align_of::<RefCounts>()`, we have `align_of::<T>()` being a multiple
40+
//! of `align_of::<RefCounts>()`, since `size_of::<RefCounts>()` is also a multiple of
41+
//! `align_of::<RefCounts>()`, we conclude the offset also has a valid alignment for `RefCounts`.
42+
//! - The value is stored at offset `size_of::<RefCounts>().next_multiple_of(align_of::<T>())`,
43+
//! which trivially satisfies the alignment requirement of `T`.
44+
//! - The distance between the `RefCounts` object and the value is `size_of::<RefCounts>()`, a fixed
45+
//! value.
46+
//!
47+
//! So both the `RefCounts` object and the value object have their alignment and size requirements
48+
//! satisfied. And we get a fixed offset between those two objects.
49+
//!
50+
//! # Reference-counting Pointer Design
51+
//!
52+
//! Both strong and weak reference-counting pointers store a pointer that points to the value
53+
//! object in a reference-counted allocation, instead of a pointer to the beginning of the
54+
//! allocation. This is based on the assumption that users access the contained value more
55+
//! frequently than the reference counters. Also, this possibly allows us to enable some
56+
//! optimizations like:
57+
//!
58+
//! - Making reference-counting pointers have ABI-compatible representation as raw pointers so we
59+
//! can use them directly in FFI interfaces.
60+
//! - Converting `Option<Rc<T>>` to `Option<&T>` without checking for `None` values.
61+
//! - Converting `&[Rc<T>]` to `&[&T]` with zero cost.
62+
63+
#![allow(dead_code)]
64+
65+
use core::cell::UnsafeCell;
66+
67+
mod rc_layout;
68+
69+
/// Stores reference counts.
70+
#[cfg_attr(target_pointer_width = "16", repr(C, align(2)))]
71+
#[cfg_attr(target_pointer_width = "32", repr(C, align(4)))]
72+
#[cfg_attr(target_pointer_width = "64", repr(C, align(8)))]
73+
pub(crate) struct RefCounts {
74+
/// Weak reference count (plus one if there are non-zero strong reference counts).
75+
pub(crate) weak: UnsafeCell<usize>,
76+
/// Strong reference count.
77+
pub(crate) strong: UnsafeCell<usize>,
78+
}
79+
80+
impl RefCounts {
81+
/// Creates a `RefCounts` with weak count of `1` and strong count of `strong_count`.
82+
const fn new(strong_count: usize) -> Self {
83+
Self { weak: UnsafeCell::new(1), strong: UnsafeCell::new(strong_count) }
84+
}
85+
}

‎library/alloc/src/raw_rc/rc_layout.rs‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
use core::alloc::{Layout, LayoutError};
2+
use core::mem::SizedTypeProperties;
3+
use core::ptr::NonNull;
4+
5+
use crate::raw_rc::RefCounts;
6+
7+
/// A `Layout` that describes a reference-counted allocation.
8+
#[derive(Clone, Copy)]
9+
pub(crate) struct RcLayout(Layout);
10+
11+
impl RcLayout {
12+
/// Tries to create an `RcLayout` to store a value with layout `value_layout`. Returns `Err` if
13+
/// `value_layout` is too big to store in a reference-counted allocation.
14+
#[inline]
15+
pub(crate) const fn try_from_value_layout(value_layout: Layout) -> Result<Self, LayoutError> {
16+
match RefCounts::LAYOUT.extend(value_layout) {
17+
Ok((rc_layout, _)) => Ok(Self(rc_layout)),
18+
Err(error) => Err(error),
19+
}
20+
}
21+
22+
/// Creates an `RcLayout` to store a value with layout `value_layout`. Panics if `value_layout`
23+
/// is too big to store in a reference-counted allocation.
24+
#[cfg(not(no_global_oom_handling))]
25+
#[inline]
26+
pub(crate) fn from_value_layout(value_layout: Layout) -> Self {
27+
Self::try_from_value_layout(value_layout).unwrap()
28+
}
29+
30+
/// Creates an `RcLayout` to store a value with layout `value_layout`.
31+
///
32+
/// # Safety
33+
///
34+
/// `RcLayout::try_from_value_layout(value_layout)` must return `Ok`.
35+
#[inline]
36+
pub(crate) unsafe fn from_value_layout_unchecked(value_layout: Layout) -> Self {
37+
unsafe { Self::try_from_value_layout(value_layout).unwrap_unchecked() }
38+
}
39+
40+
/// Creates an `RcLayout` to store an array of `length` elements of type `T`. Panics if the array
41+
/// is too big to store in a reference-counted allocation.
42+
#[cfg(not(no_global_oom_handling))]
43+
pub(crate) fn new_array<T>(length: usize) -> Self {
44+
/// For minimizing monomorphization cost.
45+
#[inline]
46+
fn inner(value_layout: Layout, length: usize) -> RcLayout {
47+
// We can use `repeat_packed` here because the outer function passes `T::LAYOUT` as the
48+
// `value_layout`, which is already padded to a multiple of its alignment.
49+
value_layout.repeat_packed(length).and_then(RcLayout::try_from_value_layout).unwrap()
50+
}
51+
52+
inner(T::LAYOUT, length)
53+
}
54+
55+
/// Returns an `Layout` object that describes the reference-counted allocation.
56+
pub(crate) fn get(&self) -> Layout {
57+
self.0
58+
}
59+
60+
/// Returns the byte offset of the value stored in a reference-counted allocation that is
61+
/// described by `self`.
62+
#[inline]
63+
pub(crate) fn value_offset(&self) -> usize {
64+
// SAFETY:
65+
//
66+
// This essentially calculates `size_of::<RefCounts>().next_multiple_of(self.align())`.
67+
//
68+
// See comments in `Layout::size_rounded_up_to_custom_align` for detailed explanation.
69+
unsafe {
70+
let align_m1 = self.0.align().unchecked_sub(1);
71+
72+
size_of::<RefCounts>().unchecked_add(align_m1) & !align_m1
73+
}
74+
}
75+
76+
/// Returns the byte size of the value stored in a reference-counted allocation that is
77+
/// described by `self`.
78+
#[cfg(not(no_global_oom_handling))]
79+
#[inline]
80+
pub(crate) fn value_size(&self) -> usize {
81+
unsafe { self.0.size().unchecked_sub(self.value_offset()) }
82+
}
83+
84+
/// Creates an `RcLayout` for storing a value that is pointed to by `value_ptr`.
85+
///
86+
/// # Safety
87+
///
88+
/// `value_ptr` has correct metadata of `T`.
89+
#[cfg(not(no_global_oom_handling))]
90+
pub(crate) unsafe fn from_value_ptr<T>(value_ptr: NonNull<T>) -> Self
91+
where
92+
T: ?Sized,
93+
{
94+
/// A helper trait for computing `RcLayout` to store a `Self` object. If `Self` is
95+
/// `Sized`, the `RcLayout` value is computed at compile time.
96+
trait SpecRcLayout {
97+
unsafe fn spec_rc_layout(value_ptr: NonNull<Self>) -> RcLayout;
98+
}
99+
100+
impl<T> SpecRcLayout for T
101+
where
102+
T: ?Sized,
103+
{
104+
#[inline]
105+
default unsafe fn spec_rc_layout(value_ptr: NonNull<Self>) -> RcLayout {
106+
RcLayout::from_value_layout(unsafe { Layout::for_value_raw(value_ptr.as_ptr()) })
107+
}
108+
}
109+
110+
impl<T> SpecRcLayout for T {
111+
#[inline]
112+
unsafe fn spec_rc_layout(_: NonNull<Self>) -> RcLayout {
113+
Self::RC_LAYOUT
114+
}
115+
}
116+
117+
unsafe { T::spec_rc_layout(value_ptr) }
118+
}
119+
120+
/// Creates an `RcLayout` for storing a value that is pointed to by `value_ptr`, assuming the
121+
/// value is small enough to fit inside a reference-counted allocation.
122+
///
123+
/// # Safety
124+
///
125+
/// - `value_ptr` has correct metadata for a `T` object.
126+
/// - It is known that the memory layout described by `value_ptr` can be used to create an
127+
/// `RcLayout` successfully.
128+
pub(crate) unsafe fn from_value_ptr_unchecked<T>(value_ptr: NonNull<T>) -> Self
129+
where
130+
T: ?Sized,
131+
{
132+
/// A helper trait for computing `RcLayout` to store a `Self` object. If `Self` is
133+
/// `Sized`, the `RcLayout` value is computed at compile time.
134+
trait SpecRcLayoutUnchecked {
135+
unsafe fn spec_rc_layout_unchecked(value_ptr: NonNull<Self>) -> RcLayout;
136+
}
137+
138+
impl<T> SpecRcLayoutUnchecked for T
139+
where
140+
T: ?Sized,
141+
{
142+
#[inline]
143+
default unsafe fn spec_rc_layout_unchecked(value_ptr: NonNull<Self>) -> RcLayout {
144+
unsafe {
145+
RcLayout::from_value_layout_unchecked(Layout::for_value_raw(value_ptr.as_ptr()))
146+
}
147+
}
148+
}
149+
150+
impl<T> SpecRcLayoutUnchecked for T {
151+
#[inline]
152+
unsafe fn spec_rc_layout_unchecked(_: NonNull<Self>) -> RcLayout {
153+
Self::RC_LAYOUT
154+
}
155+
}
156+
157+
unsafe { T::spec_rc_layout_unchecked(value_ptr) }
158+
}
159+
}
160+
161+
pub(crate) trait RcLayoutExt {
162+
/// Computes `RcLayout` at compile time if `Self` is `Sized`.
163+
const RC_LAYOUT: RcLayout;
164+
}
165+
166+
impl<T> RcLayoutExt for T {
167+
const RC_LAYOUT: RcLayout = match RcLayout::try_from_value_layout(T::LAYOUT) {
168+
Ok(rc_layout) => rc_layout,
169+
Err(_) => panic!("value is too big to store in a reference-counted allocation"),
170+
};
171+
}

0 commit comments

Comments
 (0)