Skip to content
This repository was archived by the owner on Nov 15, 2023. It is now read-only.
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
275 changes: 149 additions & 126 deletions environmental/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@

//! Safe global references to stack variables.
//!
//! Set up a global reference with declare_simple! macro giving it a name and type.
//! Set up a global reference with environmental! macro giving it a name and type.
//! Use the `using` function scoped under its name to name a reference and call a function that
//! takes no parameters yet can access said reference through the similarly placed `with` function.
//!
Expand All @@ -25,7 +25,7 @@
//! ```
//! #[macro_use] extern crate environmental;
//! // create a place for the global reference to exist.
//! declare_simple!(counter: u32);
//! environmental!(counter: u32);
//! fn stuff() {
//! // do some stuff, accessing the named reference as desired.
//! counter::with(|i| *i += 1);
Expand All @@ -40,14 +40,13 @@
//! }
//! ```

#[doc(hidden)]
pub use std::cell::RefCell;
use std::cell::RefCell;
use std::thread::LocalKey;

#[doc(hidden)]
pub fn using<'a, T: 'a, R, S, F: FnOnce() -> R>(
global: &'static LocalKey<RefCell<*mut S>>,
protected: &'a mut T,
pub fn using<T: ?Sized, R, F: FnOnce() -> R>(
global: &'static LocalKey<RefCell<Option<*mut T>>>,
protected: &mut T,
f: F
) -> R {
// store the `protected` reference as a pointer so we can provide it to logic running within
Expand All @@ -57,62 +56,69 @@ pub fn using<'a, T: 'a, R, S, F: FnOnce() -> R>(
// - that no other thread will use it; and
// - that we do not use the original mutating reference while the pointer.
// exists.
let original = global.with(|r| {
let mut b = r.borrow_mut();
let o = *b;
*b = protected as *mut T as *mut S;
o
});
let r = f();
global.with(|r| *r.borrow_mut() = original);
r
}

#[doc(hidden)]
pub fn with<'r, R, S, T: 'r, F: FnOnce(&'r mut T) -> R>(
global: &'static LocalKey<RefCell<*mut S>>,
mutator: F,
) -> Option<R> {
global.with(|r| {
let br = r.borrow_mut();
if *br != 0 as *mut S {
// safe because it's only non-zero when it's being called from using, which
// is holding on to the underlying reference (and not using it itself) safely.
unsafe {
Some(mutator(&mut *(*br as *mut S as *mut T)))
let original = {
let mut global = r.borrow_mut();
::std::mem::replace(&mut *global, Some(protected as _))
};

// even if `f` panics the original will be replaced.
struct ReplaceOriginal<'a, T: 'a + ?Sized> {
original: Option<*mut T>,
global: &'a RefCell<Option<*mut T>>,
}

impl<'a, T: 'a + ?Sized> Drop for ReplaceOriginal<'a, T> {
fn drop(&mut self) {
*self.global.borrow_mut() = self.original.take();
}
} else {
None
}

let _guard = ReplaceOriginal {
original,
global: r,
};

f()
})
}

#[doc(hidden)]
#[macro_export]
macro_rules! decl {
($name:ident : $t:ty) => {
thread_local! {
static $name: $crate::RefCell<*mut $t> = $crate::RefCell::new(0 as *mut $t);
pub fn with<T: ?Sized, R, F: FnOnce(&mut T) -> R>(
global: &'static LocalKey<RefCell<Option<*mut T>>>,
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

worth noting that Option<*mut X> has a zero-value optimization which means it takes the same size as a *mut X and the None simply occupies NULL (platform-dependent) in the pointer space.

mutator: F,
) -> Option<R> {
global.with(|r| unsafe {
let ptr = r.borrow_mut();
match *ptr {
Some(ptr) => {
// safe because it's only non-zero when it's being called from using, which
// is holding on to the underlying reference (and not using it itself) safely.
Some(mutator(&mut *ptr))
}
None => None,
}
}
})
}

/// Declare a new global reference module whose underlying value does not contain references.
///
/// Will create a module of a given name that contains two functions:
/// * `pub fn using<'a: 'b, 'b, R, F: FnOnce() -> R, T: 'a>(protected: &'b mut T, f: F) -> R`
/// * `pub fn using<R, F: FnOnce() -> R>(protected: &mut $t, f: F) -> R`
/// This executes `f`, returning its value. During the call, the module's reference is set to
/// be equal to `protected`.
/// * `pub fn with<R, F: for<'r, 't: 'r> FnOnce(&'r mut $t<'t>) -> R>(f: F) -> Option<R>`
/// * `pub fn with<R, F: FnOnce(&mut $t) -> R>(f: F) -> Option<R>`
/// This executes `f`, returning `Some` of its value if called from code that is being executed
/// as part of a `using` call. If not, it returns `None`. `f` is provided with one argument: the
/// same reference as provided to the most recent `using` call.
///
/// # Examples
///
/// ```
/// Initializing the global context with a given value.
///
/// ```rust
/// #[macro_use] extern crate environmental;
/// declare_simple!(counter: u32);
/// environmental!(counter: u32);
/// fn main() {
/// let mut counter_value = 41u32;
/// counter::using(&mut counter_value, || {
Expand All @@ -128,122 +134,139 @@ macro_rules! decl {
/// println!("The answer is {:?}", counter_value); // 42
/// }
/// ```
#[macro_export]
macro_rules! declare_simple {
($name:ident : $t:tt) => {
mod $name {
#[allow(unused_imports)]
use super::*;

decl!(GLOBAL: $t);

pub fn using<'a: 'b, 'b, R, F: FnOnce() -> R, T: 'a>(
protected: &'b mut T,
f: F
) -> R {
$crate::using(&GLOBAL, protected, f)
}

pub fn with<R, F: for<'r> FnOnce(&'r mut $t) -> R>(
f: F
) -> Option<R> {
let dummy = ();
with_closed(f, &dummy)
}

fn with_closed<'d: 'r, 'r, R, F: FnOnce(&'r mut $t) -> R>(
f: F,
_dummy: &'d (),
) -> Option<R> {
$crate::with(&GLOBAL, f)
}
}
}
}

/// Declare a new global reference module whose underlying value is generic over a reference.
///
/// Will create a module of a given name that contains two functions:
/// Roughly the same, but with a trait object:
///
/// * `pub fn using<'a: 'b, 'b, R, F: FnOnce() -> R, T: 'a>(protected: &'b mut T, f: F) -> R`
/// This executes `f`, returning its value. During the call, the module's reference is set to
/// be equal to `protected`.
/// * `pub fn with<R, F: for<'r, 't: 'r> FnOnce(&'r mut $t<'t>) -> R>(f: F) -> Option<R>`
/// This executes `f`, returning `Some` of its value if called from code that is being executed
/// as part of a `using` call. If not, it returns `None`. `f` is provided with one argument: the
/// same reference as provided to the most recent `using` call.
/// ```rust
/// #[macro_use] extern crate environmental;
///
/// # Examples
/// trait Increment { fn increment(&mut self); }
///
/// ```
/// #[macro_use] extern crate environmental;
/// // a type that we want to reference from a temp global; it has a reference in it.
/// struct WithReference<'a> { answer_ref: &'a mut u32, }
/// // create a place for the global reference to exist.
/// declare_generic!(counter: WithReference);
/// fn stuff() {
/// // do some stuff, accessing the named reference as desired.
/// counter::with(|i| *i.answer_ref += 1);
/// impl Increment for i32 {
/// fn increment(&mut self) { *self += 1 }
/// }
///
/// environmental!(val: Increment + 'static);
///
/// fn main() {
/// // declare a stack variable of the same type as our global declaration.
/// let mut answer = 41u32;
/// {
/// let mut ref_struct = WithReference { answer_ref: &mut answer, };
/// counter::using(&mut ref_struct, stuff);
/// }
/// println!("The answer is {:?}", answer); // will print 42!
/// let mut local = 0i32;
/// val::using(&mut local, || {
/// val::with(|v| for _ in 0..5 { v.increment() });
/// });
///
/// assert_eq!(local, 5);
/// }
/// ```
#[macro_export]
macro_rules! declare_generic {
($name:ident : $t:tt) => {
mod $name {
#[allow(unused_imports)]
use super::*;
macro_rules! environmental {
($name:ident : $t:ty) => {
#[allow(non_camel_case_types)]
struct $name { __private_field: () }

thread_local!(static GLOBAL: ::std::cell::RefCell<Option<*mut $t>>
= ::std::cell::RefCell::new(None));

decl!(GLOBAL: $t<'static> );
impl $name {
#[allow(unused_imports)]

pub fn using<'a: 'b, 'b, R, F: FnOnce() -> R, T: 'a>(
protected: &'b mut T,
pub fn using<R, F: FnOnce() -> R>(
protected: &mut $t,
f: F
) -> R {
$crate::using(&GLOBAL, protected, f)
}

pub fn with<R, F: for<'r, 't: 'r> FnOnce(&'r mut $t<'t>) -> R>(
pub fn with<R, F: FnOnce(&mut $t) -> R>(
f: F
) -> Option<R> {
let dummy = ();
with_closed(f, &dummy)
}

fn with_closed<'d: 't, 't: 'r, 'r, R, F: FnOnce(&'r mut $t<'t>) -> R>(
f: F,
_dummy: &'d (),
) -> Option<R> {
$crate::with(&GLOBAL, f)
$crate::with(&GLOBAL, |x| f(x))
}
}
}
};
}

#[cfg(test)]
mod tests {
declare_simple!(counter: u32);

fn stuff() {
// do some stuff, accessing the named reference as desired.
counter::with(|value| *value += 1);
}

#[test]
fn simple_works() {
environmental!(counter: u32);

fn stuff() { counter::with(|value| *value += 1); };

// declare a stack variable of the same type as our global declaration.
let mut local = 41u32;

// call stuff, setting up our `counter` environment as a refrence to our local counter var.
counter::using(&mut local, stuff);
println!("The answer is {:?}", local); // will print 42!
assert_eq!(local, 42);
stuff(); // safe! doesn't do anything.
}

#[test]
fn overwrite_with_lesser_lifetime() {
environmental!(items: Vec<u8>);

let mut local_items = vec![1, 2, 3];
items::using(&mut local_items, || {
let dies_at_end = vec![4, 5, 6];
items::with(|items| *items = dies_at_end);
});

assert_eq!(local_items, vec![4, 5, 6]);
}

#[test]
fn declare_with_trait_object() {
trait Foo {
fn get(&self) -> i32;
fn set(&mut self, x: i32);
}

impl Foo for i32 {
fn get(&self) -> i32 { *self }
fn set(&mut self, x: i32) { *self = x }
}

environmental!(foo: Foo + 'static);

fn stuff() {
foo::with(|value| {
let new_val = value.get() + 1;
value.set(new_val);
});
}

let mut local = 41i32;
foo::using(&mut local, stuff);

assert_eq!(local, 42);

stuff(); // doesn't do anything.
}

#[test]
fn unwind_recursive() {
use std::panic;

environmental!(items: Vec<u8>);

let panicked = panic::catch_unwind(|| {
let mut local_outer = vec![1, 2, 3];

items::using(&mut local_outer, || {
let mut local_inner = vec![4, 5, 6];
items::using(&mut local_inner, || {
panic!("are you unsafe?");
})
});
}).is_err();

assert!(panicked);

let mut was_cleared = true;
items::with(|_items| was_cleared = false);

assert!(was_cleared);
}
}
Loading