stack_tokens/
lib.rs

1//! This library implements stack tokens which can be used to safely borrow
2//! values with stack-local lifetimes.
3//!
4//! [`StackToken`]s are zero sized objects that can be placed on the call
5//! stack with the [`stack_token!`] macro which then can be used to safely
6//! borrow from places such as thread local storage with reduced lifetimes.
7//!
8//! Without stack tokens the only safe API for such constructs are callback
9//! based such as the [`LocalKey::with`](std::thread::LocalKey::with) API.
10//! This problem however is not always restricted to thread local storage
11//! directly as some APIs are internally constrained by similar challenges.
12//!
13//! The problem usually appears when a proxy object wants to lend out some
14//! memory but it does not have a better lifetime than itself to constrain
15//! the value, but it does not directly own the value it's trying to lend
16//! out.  As a Rust programmer one is enticed to try to constrain it by
17//! the lifetime of `&self` but thanks to [`Box::leak`] that lifetime can
18//! become `&'static`.
19//!
20//! For more information see the [the blog post describing the concept](https://lucumr.pocoo.org/2022/11/23/stack-tokens/).
21//!
22//! # Ref Cells
23//!
24//! This example shows how stack tokens can be used with the
25//! [`RefCellLocalKeyExt`] extension trait to directly borrow into [`RefCell`]s
26//! in a thread local.
27//!
28//! ```
29//! use stack_tokens::{stack_token, RefCellLocalKeyExt};
30//! use std::cell::RefCell;
31//!
32//! thread_local! {
33//!     static VEC: RefCell<Vec<i32>> = RefCell::default();
34//! }
35//!
36//! // places a token on the stack.
37//! stack_token!(scope);
38//!
39//! // you can now directly deref the thread local without closures
40//! VEC.as_mut(scope).push(42);
41//! assert_eq!(VEC.as_ref(scope).len(), 1);
42//! ```
43//!
44//! # Basic Use
45//!
46//! This example shows how stack tokens can be used with the [`LocalKeyExt`]
47//! extension trait to get stack local borrows from a standard library thread
48//! local.
49//!
50//! ```
51//! use stack_tokens::{stack_token, LocalKeyExt};
52//! use std::sync::atomic::{AtomicUsize, Ordering};
53//!
54//! thread_local! {
55//!     static COUNTER: AtomicUsize = AtomicUsize::new(0);
56//! }
57//!
58//! // places a token on the stack
59//! stack_token!(scope);
60//!
61//! // borrow can be used to get a stack local reference
62//! COUNTER.borrow(scope).fetch_add(1, Ordering::Acquire);
63//! assert_eq!(COUNTER.borrow(scope).load(Ordering::Acquire), 1);
64//! ```
65//!
66//! # Implementing Stack Local APIs
67//!
68//! To implement your own methods that use stack tokens introduce a new lifetime
69//! (eg: `'stack`) and constrain both `&self` and the passed token with it:
70//!
71//! ```
72//! use stack_tokens::StackToken;
73//! use std::marker::PhantomData;
74//!
75//! struct MyTls<T>(PhantomData::<T>);
76//!
77//! impl<T> MyTls<T> {
78//!     pub fn get<'stack>(&'stack self, token: &'stack StackToken) -> &'stack T {
79//!         let _ = token;
80//!         todo!()
81//!     }
82//! }
83//! ```
84use std::cell::{Ref, RefCell, RefMut};
85use std::marker::PhantomData;
86use std::mem::transmute;
87use std::thread::LocalKey;
88
89/// A token to bind lifetimes to a specific stack.
90///
91/// For more information see [`stack_token`].
92pub struct StackToken {
93    _marker: PhantomData<*const ()>,
94}
95
96impl StackToken {
97    #[doc(hidden)]
98    pub unsafe fn __private_new() -> StackToken {
99        StackToken {
100            _marker: PhantomData,
101        }
102    }
103}
104
105/// Creates a new [`StackToken`] with a given name on the stack.
106#[macro_export]
107macro_rules! stack_token {
108    ($name:ident) => {
109        #[allow(unsafe_code)]
110        let $name = &unsafe { $crate::StackToken::__private_new() };
111    };
112}
113
114/// Adds [`StackToken`] support to the standard library's [`LocalKey`].
115pub trait LocalKeyExt<T> {
116    /// Borrows the value from the TLS with a [`StackToken`].
117    fn borrow<'stack>(&'static self, token: &'stack StackToken) -> &'stack T;
118}
119
120impl<T: 'static> LocalKeyExt<T> for LocalKey<T> {
121    fn borrow<'stack>(&'static self, token: &'stack StackToken) -> &'stack T {
122        let _ = token;
123        self.with(|value| unsafe { transmute::<&T, &'stack T>(value) })
124    }
125}
126
127/// Additional utility methods to [`LocalKey`]s holding [`RefCell`] values.
128///
129/// This extension traits provides the two methods [`as_ref`](Self::as_ref)
130/// and [`as_mut`](Self::as_mut) that let you directly borrow into the
131/// contained [`RefCell`] with a [`StackToken`].
132pub trait RefCellLocalKeyExt<T> {
133    /// Acquires a reference to the contained value.
134    fn as_ref<'stack>(&'static self, token: &'stack StackToken) -> Ref<'stack, T>;
135
136    /// Acquires a mutable reference to the contained value.
137    fn as_mut<'stack>(&'static self, token: &'stack StackToken) -> RefMut<'stack, T>;
138}
139
140impl<T: 'static> RefCellLocalKeyExt<T> for LocalKey<RefCell<T>> {
141    fn as_ref<'stack>(&'static self, token: &'stack StackToken) -> Ref<'stack, T> {
142        self.borrow(token).borrow()
143    }
144
145    fn as_mut<'stack>(&'static self, token: &'stack StackToken) -> RefMut<'stack, T> {
146        self.borrow(token).borrow_mut()
147    }
148}
149
150#[test]
151fn test_tls_basic() {
152    use crate::stack_token;
153    use std::cell::RefCell;
154
155    thread_local! { static FOO: RefCell<u32> = RefCell::default(); }
156
157    stack_token!(scope);
158    *FOO.borrow(scope).borrow_mut() += 1;
159    assert_eq!(*FOO.borrow(scope).borrow(), 1);
160}
161
162#[test]
163fn test_tls_ref_cell() {
164    use crate::stack_token;
165    use std::cell::RefCell;
166
167    thread_local! { static FOO: RefCell<u32> = RefCell::default(); }
168
169    stack_token!(scope);
170    *FOO.as_mut(scope) += 1;
171    assert_eq!(*FOO.as_ref(scope), 1);
172}
stack_tokens/lib.rs

stack_tokens/
lib.rs
