Skip to content

Commit 8d8878a

Browse files
EFanZhEFanZh
committed
Add RefCounts and RcLayout types
1 parent 523d399 commit 8d8878a

File tree

3 files changed

+257
-0
lines changed

3 files changed

+257
-0
lines changed

‎library/alloc/src/lib.rs‎

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

208+
#[cfg(not(no_rc))]
209+
mod raw_rc;
208210
mod raw_vec;
209211

210212
// 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: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
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+
#[inline]
45+
fn inner(value_layout: Layout, length: usize) -> RcLayout {
46+
// We can use `repeat_packed` here because the outer function passes `T::LAYOUT` as the
47+
// `value_layout`, which is already padded to a multiple of its alignment.
48+
value_layout.repeat_packed(length).and_then(RcLayout::try_from_value_layout).unwrap()
49+
}
50+
51+
inner(T::LAYOUT, length)
52+
}
53+
54+
/// Returns an `Layout` object that describes the reference-counted allocation.
55+
pub(crate) fn get(&self) -> Layout {
56+
self.0
57+
}
58+
59+
/// Returns the byte offset of the value stored in a reference-counted allocation that is
60+
/// described by `self`.
61+
#[inline]
62+
pub(crate) fn value_offset(&self) -> usize {
63+
// SAFETY:
64+
//
65+
// This essentially calculates `size_of::<RefCounts>().next_multiple_of(self.align())`.
66+
//
67+
// See comments in `Layout::size_rounded_up_to_custom_align` for detailed explanation.
68+
unsafe {
69+
let align_m1 = self.0.align().unchecked_sub(1);
70+
71+
size_of::<RefCounts>().unchecked_add(align_m1) & !align_m1
72+
}
73+
}
74+
75+
/// Returns the byte size of the value stored in a reference-counted allocation that is
76+
/// described by `self`.
77+
#[cfg(not(no_global_oom_handling))]
78+
#[inline]
79+
pub(crate) fn value_size(&self) -> usize {
80+
unsafe { self.0.size().unchecked_sub(self.value_offset()) }
81+
}
82+
83+
/// Creates an `RcLayout` for storing a value that is pointed to by `value_ptr`.
84+
///
85+
/// # Safety
86+
///
87+
/// `value_ptr` has correct metadata of `T`.
88+
#[cfg(not(no_global_oom_handling))]
89+
pub(crate) unsafe fn from_value_ptr<T>(value_ptr: NonNull<T>) -> Self
90+
where
91+
T: ?Sized,
92+
{
93+
/// A helper trait for computing `RcLayout` to store a `Self` object. If `Self` is
94+
/// `Sized`, the `RcLayout` value is computed at compile time.
95+
trait SpecRcLayout {
96+
unsafe fn spec_rc_layout(value_ptr: NonNull<Self>) -> RcLayout;
97+
}
98+
99+
impl<T> SpecRcLayout for T
100+
where
101+
T: ?Sized,
102+
{
103+
#[inline]
104+
default unsafe fn spec_rc_layout(value_ptr: NonNull<Self>) -> RcLayout {
105+
RcLayout::from_value_layout(unsafe { Layout::for_value_raw(value_ptr.as_ptr()) })
106+
}
107+
}
108+
109+
impl<T> SpecRcLayout for T {
110+
#[inline]
111+
unsafe fn spec_rc_layout(_: NonNull<Self>) -> RcLayout {
112+
Self::RC_LAYOUT
113+
}
114+
}
115+
116+
unsafe { T::spec_rc_layout(value_ptr) }
117+
}
118+
119+
/// Creates an `RcLayout` for storing a value that is pointed to by `value_ptr`, assuming the
120+
/// value is small enough to fit inside a reference-counted allocation.
121+
///
122+
/// # Safety
123+
///
124+
/// - `value_ptr` has correct metadata for a `T` object.
125+
/// - It is known that the memory layout described by `value_ptr` can be used to create an
126+
/// `RcLayout` successfully.
127+
pub(crate) unsafe fn from_value_ptr_unchecked<T>(value_ptr: NonNull<T>) -> Self
128+
where
129+
T: ?Sized,
130+
{
131+
/// A helper trait for computing `RcLayout` to store a `Self` object. If `Self` is
132+
/// `Sized`, the `RcLayout` value is computed at compile time.
133+
trait SpecRcLayoutUnchecked {
134+
unsafe fn spec_rc_layout_unchecked(value_ptr: NonNull<Self>) -> RcLayout;
135+
}
136+
137+
impl<T> SpecRcLayoutUnchecked for T
138+
where
139+
T: ?Sized,
140+
{
141+
#[inline]
142+
default unsafe fn spec_rc_layout_unchecked(value_ptr: NonNull<Self>) -> RcLayout {
143+
unsafe {
144+
RcLayout::from_value_layout_unchecked(Layout::for_value_raw(value_ptr.as_ptr()))
145+
}
146+
}
147+
}
148+
149+
impl<T> SpecRcLayoutUnchecked for T {
150+
#[inline]
151+
unsafe fn spec_rc_layout_unchecked(_: NonNull<Self>) -> RcLayout {
152+
Self::RC_LAYOUT
153+
}
154+
}
155+
156+
unsafe { T::spec_rc_layout_unchecked(value_ptr) }
157+
}
158+
}
159+
160+
pub(crate) trait RcLayoutExt {
161+
/// Computes `RcLayout` at compile time if `Self` is `Sized`.
162+
const RC_LAYOUT: RcLayout;
163+
}
164+
165+
impl<T> RcLayoutExt for T {
166+
const RC_LAYOUT: RcLayout = match RcLayout::try_from_value_layout(T::LAYOUT) {
167+
Ok(rc_layout) => rc_layout,
168+
Err(_) => panic!("value is too big to store in a reference-counted allocation"),
169+
};
170+
}

0 commit comments

Comments
 (0)