Skip to content

Commit 36c281d

Browse files
EFanZhEFanZh
committed
Add RefCounts and RcLayout types
1 parent 3ff30e7 commit 36c281d

File tree

3 files changed

+278
-0
lines changed

3 files changed

+278
-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: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
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 from 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+
//! The following table shows the order and size of each component in a 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 are 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 has
36+
//! a valid alignment for `RefCounts` because:
37+
//! - If `align_of::<T>() <= align_of::<RefCounts>()`, then the offset is 0, which has a valid
38+
//! alignment for `RefCounts`.
39+
//! - If `align_of::<T>() > align_of::<RefCounts>()`, then `align_of::<T>()` is a multiple of
40+
//! `align_of::<RefCounts>()`. Since `size_of::<RefCounts>()` is also a multiple of
41+
//! `align_of::<RefCounts>()`, 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+
//! Thus both the `RefCounts` object and the value 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 to the value object inside the
53+
//! reference-counted allocation, instead of a pointer to the start of the allocation. This is based
54+
//! on the assumption that users access the contained value more frequently than the reference
55+
//! counters. Also, this allows us to enable some possible optimizations such as:
56+
//!
57+
//! - Making reference-counting pointers have ABI-compatible representations as raw pointers so we
58+
//! can use them directly in FFI interfaces.
59+
//! - Converting `Option<Rc<T>>` to `Option<&T>` without checking for `None` values.
60+
//! - Converting `&[Rc<T>]` to `&[&T]` with zero cost.
61+
62+
#![allow(dead_code)]
63+
64+
use core::cell::UnsafeCell;
65+
66+
mod rc_layout;
67+
68+
/// Stores reference counts.
69+
#[cfg_attr(target_pointer_width = "16", repr(C, align(2)))]
70+
#[cfg_attr(target_pointer_width = "32", repr(C, align(4)))]
71+
#[cfg_attr(target_pointer_width = "64", repr(C, align(8)))]
72+
pub(crate) struct RefCounts {
73+
/// Weak reference count (plus one if there are non-zero strong reference counts).
74+
pub(crate) weak: UnsafeCell<usize>,
75+
/// Strong reference count.
76+
pub(crate) strong: UnsafeCell<usize>,
77+
}
78+
79+
impl RefCounts {
80+
/// Creates a `RefCounts` with weak count of `1` and strong count of `strong_count`.
81+
const fn new(strong_count: usize) -> Self {
82+
Self { weak: UnsafeCell::new(1), strong: UnsafeCell::new(strong_count) }
83+
}
84+
}

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

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

0 commit comments

Comments
 (0)