From 948a62401ef717eb484c2215713291749f0f35ed Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Tue, 16 Jul 2013 17:20:43 -0700 Subject: [PATCH] Add a `get_mut` method for TLS Simulates borrow checks for '@mut' boxes, or at least it's the same idea. --- src/libstd/local_data.rs | 77 ++++++++++++++++++++++++++- src/libstd/task/local_data_priv.rs | 84 ++++++++++++++++++++++++------ 2 files changed, 145 insertions(+), 16 deletions(-) diff --git a/src/libstd/local_data.rs b/src/libstd/local_data.rs index be170cce07e..6c1640e683e 100644 --- a/src/libstd/local_data.rs +++ b/src/libstd/local_data.rs @@ -41,7 +41,7 @@ magic. use prelude::*; -use task::local_data_priv::{local_get, local_pop, local_set, Handle}; +use task::local_data_priv::*; #[cfg(test)] use task; @@ -95,6 +95,13 @@ pub fn get(key: Key<@T>, f: &fn(Option<&@T>) -> U) -> U { pub fn get(key: Key, f: &fn(Option<&T>) -> U) -> U { unsafe { local_get(Handle::new(), key, f) } } +/** + * Retrieve a mutable borrowed pointer to a task-local data value. + */ +#[cfg(not(stage0))] +pub fn get_mut(key: Key, f: &fn(Option<&mut T>) -> U) -> U { + unsafe { local_get_mut(Handle::new(), key, f) } +} /** * Store a value in task-local data. If this key already has a value, * that value is overwritten (and its destructor is run). @@ -262,6 +269,34 @@ fn test_static_pointer() { fn test_owned() { static key: Key<~int> = &Key; set(key, ~1); + + do get(key) |v| { + do get(key) |v| { + do get(key) |v| { + assert_eq!(**v.unwrap(), 1); + } + assert_eq!(**v.unwrap(), 1); + } + assert_eq!(**v.unwrap(), 1); + } + set(key, ~2); + do get(key) |v| { + assert_eq!(**v.unwrap(), 2); + } +} + +#[test] +fn test_get_mut() { + static key: Key = &Key; + set(key, 1); + + do get_mut(key) |v| { + *v.unwrap() = 2; + } + + do get(key) |v| { + assert_eq!(*v.unwrap(), 2); + } } #[test] @@ -283,3 +318,43 @@ fn test_same_key_type() { get(key4, |x| assert_eq!(*x.unwrap(), 4)); get(key5, |x| assert_eq!(*x.unwrap(), 5)); } + +#[test] +#[should_fail] +fn test_nested_get_set1() { + static key: Key = &Key; + set(key, 4); + do get(key) |_| { + set(key, 4); + } +} + +#[test] +#[should_fail] +fn test_nested_get_mut2() { + static key: Key = &Key; + set(key, 4); + do get(key) |_| { + get_mut(key, |_| {}) + } +} + +#[test] +#[should_fail] +fn test_nested_get_mut3() { + static key: Key = &Key; + set(key, 4); + do get_mut(key) |_| { + get(key, |_| {}) + } +} + +#[test] +#[should_fail] +fn test_nested_get_mut4() { + static key: Key = &Key; + set(key, 4); + do get_mut(key) |_| { + get_mut(key, |_| {}) + } +} diff --git a/src/libstd/task/local_data_priv.rs b/src/libstd/task/local_data_priv.rs index d46e5707f14..75fd6eacc1b 100644 --- a/src/libstd/task/local_data_priv.rs +++ b/src/libstd/task/local_data_priv.rs @@ -44,6 +44,21 @@ impl Handle { } } +#[deriving(Eq)] +enum LoanState { + NoLoan, ImmLoan, MutLoan +} + +impl LoanState { + fn describe(&self) -> &'static str { + match *self { + NoLoan => "no loan", + ImmLoan => "immutable", + MutLoan => "mutable" + } + } +} + trait LocalData {} impl LocalData for T {} @@ -77,7 +92,7 @@ impl LocalData for T {} // // n.b. If TLS is used heavily in future, this could be made more efficient with // a proper map. -type TaskLocalMap = ~[Option<(*libc::c_void, TLSValue, uint)>]; +type TaskLocalMap = ~[Option<(*libc::c_void, TLSValue, LoanState)>]; type TLSValue = ~LocalData:; fn cleanup_task_local_map(map_ptr: *libc::c_void) { @@ -152,9 +167,10 @@ pub unsafe fn local_pop(handle: Handle, for map.mut_iter().advance |entry| { match *entry { - Some((k, _, loans)) if k == key_value => { - if loans != 0 { - fail!("TLS value has been loaned via get already"); + Some((k, _, loan)) if k == key_value => { + if loan != NoLoan { + fail!("TLS value cannot be removed because it is already \ + borrowed as %s", loan.describe()); } // Move the data out of the `entry` slot via util::replace. This // is guaranteed to succeed because we already matched on `Some` @@ -192,6 +208,29 @@ pub unsafe fn local_pop(handle: Handle, pub unsafe fn local_get(handle: Handle, key: local_data::Key, f: &fn(Option<&T>) -> U) -> U { + local_get_with(handle, key, ImmLoan, f) +} + +pub unsafe fn local_get_mut(handle: Handle, + key: local_data::Key, + f: &fn(Option<&mut T>) -> U) -> U { + do local_get_with(handle, key, MutLoan) |x| { + match x { + None => f(None), + // We're violating a lot of compiler guarantees with this + // invocation of `transmute_mut`, but we're doing runtime checks to + // ensure that it's always valid (only one at a time). + // + // there is no need to be upset! + Some(x) => { f(Some(cast::transmute_mut(x))) } + } + } +} + +unsafe fn local_get_with(handle: Handle, + key: local_data::Key, + state: LoanState, + f: &fn(Option<&T>) -> U) -> U { // This function must be extremely careful. Because TLS can store owned // values, and we must have some form of `get` function other than `pop`, // this function has to give a `&` reference back to the caller. @@ -218,12 +257,24 @@ pub unsafe fn local_get(handle: Handle, None => { return f(None); } Some(i) => { let ret; + let mut return_loan = false; match map[i] { - Some((_, ref data, ref mut loans)) => { - *loans = *loans + 1; + Some((_, ref data, ref mut loan)) => { + match (state, *loan) { + (_, NoLoan) => { + *loan = state; + return_loan = true; + } + (ImmLoan, ImmLoan) => {} + (want, cur) => { + fail!("TLS slot cannot be borrowed as %s because \ + it is already borrowed as %s", + want.describe(), cur.describe()); + } + } // data was created with `~~T as ~LocalData`, so we extract // pointer part of the trait, (as ~~T), and then use - // compiler coercions to achieve a '&' pointer + // compiler coercions to achieve a '&' pointer. match *cast::transmute::<&TLSValue, &(uint, ~~T)>(data) { (_vtable, ref box) => { let value: &T = **box; @@ -238,9 +289,11 @@ pub unsafe fn local_get(handle: Handle, // 'f' returned because `f` could have appended more TLS items which // in turn relocated the vector. Hence we do another lookup here to // fixup the loans. - match map[i] { - Some((_, _, ref mut loans)) => { *loans = *loans - 1; } - None => { libc::abort(); } + if return_loan { + match map[i] { + Some((_, _, ref mut loan)) => { *loan = NoLoan; } + None => { libc::abort(); } + } } return ret; } @@ -269,9 +322,10 @@ pub unsafe fn local_set(handle: Handle, // First see if the map contains this key already let curspot = map.iter().position(|entry| { match *entry { - Some((ekey, _, loans)) if key == ekey => { - if loans != 0 { - fail!("TLS value has been loaned via get already"); + Some((ekey, _, loan)) if key == ekey => { + if loan != NoLoan { + fail!("TLS value cannot be overwritten because it is + already borrowed as %s", loan.describe()) } true } @@ -286,7 +340,7 @@ pub unsafe fn local_set(handle: Handle, } match insertion_position(map, keyval) { - Some(i) => { map[i] = Some((keyval, data, 0)); } - None => { map.push(Some((keyval, data, 0))); } + Some(i) => { map[i] = Some((keyval, data, NoLoan)); } + None => { map.push(Some((keyval, data, NoLoan))); } } }