src/libstd/thread_local/scoped.rs

   1 // Copyright 2014 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 //! Scoped thread-local storage
  12 //!
  13 //! This module provides the ability to generate *scoped* thread-local
  14 //! variables. In this sense, scoped indicates that thread local storage
  15 //! actually stores a reference to a value, and this reference is only placed
  16 //! in storage for a scoped amount of time.
  17 //!
  18 //! There are no restrictions on what types can be placed into a scoped
  19 //! variable, but all scoped variables are initialized to the equivalent of
  20 //! null. Scoped thread local stor is useful when a value is present for a known
  21 //! period of time and it is not required to relinquish ownership of the
  22 //! contents.
  23 //!
  24 //! # Example
  25 //!
  26 //! ```
  27 //! scoped_thread_local!(static FOO: uint)
  28 //!
  29 //! // Initially each scoped slot is empty.
  30 //! assert!(!FOO.is_set());
  31 //!
  32 //! // When inserting a value, the value is only in place for the duration
  33 //! // of the closure specified.
  34 //! FOO.set(&1, || {
  35 //!     FOO.with(|slot| {
  36 //!         assert_eq!(*slot, 1);
  37 //!     });
  38 //! });
  39 //! ```
  40
  41 #![macro_escape]
  42
  43 use prelude::*;
  44
  45 // macro hygiene sure would be nice, wouldn't it?
  46 #[doc(hidden)] pub use self::imp::KeyInner;
  47 #[doc(hidden)] pub use sys_common::thread_local::INIT as OS_INIT;
  48
  49 /// Type representing a thread local storage key corresponding to a reference
  50 /// to the type parameter `T`.
  51 ///
  52 /// Keys are statically allocated and can contain a reference to an instance of
  53 /// type `T` scoped to a particular lifetime. Keys provides two methods, `set`
  54 /// and `with`, both of which currently use closures to control the scope of
  55 /// their contents.
  56 pub struct Key<T> { #[doc(hidden)] pub inner: KeyInner<T> }
  57
  58 /// Declare a new scoped thread local storage key.
  59 ///
  60 /// This macro declares a `static` item on which methods are used to get and
  61 /// set the value stored within.
  62 #[macro_export]
  63 macro_rules! scoped_thread_local(
  64     (static $name:ident: $t:ty) => (
  65         __scoped_thread_local_inner!(static $name: $t)
  66     );
  67     (pub static $name:ident: $t:ty) => (
  68         __scoped_thread_local_inner!(pub static $name: $t)
  69     );
  70 )
  71
  72 #[macro_export]
  73 #[doc(hidden)]
  74 macro_rules! __scoped_thread_local_inner(
  75     (static $name:ident: $t:ty) => (
  76         #[cfg_attr(not(any(windows, target_os = "android", target_os = "ios")),
  77                    thread_local)]
  78         static $name: ::std::thread_local::scoped::Key<$t> =
  79             __scoped_thread_local_inner!($t);
  80     );
  81     (pub static $name:ident: $t:ty) => (
  82         #[cfg_attr(not(any(windows, target_os = "android", target_os = "ios")),
  83                    thread_local)]
  84         pub static $name: ::std::thread_local::scoped::Key<$t> =
  85             __scoped_thread_local_inner!($t);
  86     );
  87     ($t:ty) => ({
  88         use std::thread_local::scoped::Key as __Key;
  89
  90         #[cfg(not(any(windows, target_os = "android", target_os = "ios")))]
  91         const INIT: __Key<$t> = __Key {
  92             inner: ::std::thread_local::scoped::KeyInner {
  93                 inner: ::std::cell::UnsafeCell { value: 0 as *mut _ },
  94             }
  95         };
  96
  97         #[cfg(any(windows, target_os = "android", target_os = "ios"))]
  98         const INIT: __Key<$t> = __Key {
  99             inner: ::std::thread_local::scoped::KeyInner {
 100                 inner: ::std::thread_local::scoped::OS_INIT,
 101                 marker: ::std::kinds::marker::InvariantType,
 102             }
 103         };
 104
 105         INIT
 106     })
 107 )
 108
 109 impl<T> Key<T> {
 110     /// Insert a value into this scoped thread local storage slot for a
 111     /// duration of a closure.
 112     ///
 113     /// While `cb` is running, the value `t` will be returned by `get` unless
 114     /// this function is called recursively inside of `cb`.
 115     ///
 116     /// Upon return, this function will restore the previous value, if any
 117     /// was available.
 118     ///
 119     /// # Example
 120     ///
 121     /// ```
 122     /// scoped_thread_local!(static FOO: uint)
 123     ///
 124     /// FOO.set(&100, || {
 125     ///     let val = FOO.with(|v| *v);
 126     ///     assert_eq!(val, 100);
 127     ///
 128     ///     // set can be called recursively
 129     ///     FOO.set(&101, || {
 130     ///         // ...
 131     ///     });
 132     ///
 133     ///     // Recursive calls restore the previous value.
 134     ///     let val = FOO.with(|v| *v);
 135     ///     assert_eq!(val, 100);
 136     /// });
 137     /// ```
 138     pub fn set<R>(&'static self, t: &T, cb: || -> R) -> R {
 139         struct Reset<'a, T: 'a> {
 140             key: &'a KeyInner<T>,
 141             val: *mut T,
 142         }
 143         #[unsafe_destructor]
 144         impl<'a, T> Drop for Reset<'a, T> {
 145             fn drop(&mut self) {
 146                 unsafe { self.key.set(self.val) }
 147             }
 148         }
 149
 150         let prev = unsafe {
 151             let prev = self.inner.get();
 152             self.inner.set(t as *const T as *mut T);
 153             prev
 154         };
 155
 156         let _reset = Reset { key: &self.inner, val: prev };
 157         cb()
 158     }
 159
 160     /// Get a value out of this scoped variable.
 161     ///
 162     /// This function takes a closure which receives the value of this
 163     /// variable.
 164     ///
 165     /// # Panics
 166     ///
 167     /// This function will panic if `set` has not previously been called.
 168     ///
 169     /// # Example
 170     ///
 171     /// ```no_run
 172     /// scoped_thread_local!(static FOO: uint)
 173     ///
 174     /// FOO.with(|slot| {
 175     ///     // work with `slot`
 176     /// });
 177     /// ```
 178     pub fn with<R>(&'static self, cb: |&T| -> R) -> R {
 179         unsafe {
 180             let ptr = self.inner.get();
 181             assert!(!ptr.is_null(), "cannot access a scoped thread local \
 182                                      variable without calling `set` first");
 183             cb(&*ptr)
 184         }
 185     }
 186
 187     /// Test whether this TLS key has been `set` for the current thread.
 188     pub fn is_set(&'static self) -> bool {
 189         unsafe { !self.inner.get().is_null() }
 190     }
 191 }
 192
 193 #[cfg(not(any(windows, target_os = "android", target_os = "ios")))]
 194 mod imp {
 195     use std::cell::UnsafeCell;
 196
 197     // FIXME: Should be a `Cell`, but that's not `Sync`
 198     #[doc(hidden)]
 199     pub struct KeyInner<T> { pub inner: UnsafeCell<*mut T> }
 200
 201     #[doc(hidden)]
 202     impl<T> KeyInner<T> {
 203         #[doc(hidden)]
 204         pub unsafe fn set(&self, ptr: *mut T) { *self.inner.get() = ptr; }
 205         #[doc(hidden)]
 206         pub unsafe fn get(&self) -> *mut T { *self.inner.get() }
 207     }
 208 }
 209
 210 #[cfg(any(windows, target_os = "android", target_os = "ios"))]
 211 mod imp {
 212     use kinds::marker;
 213     use sys_common::thread_local::StaticKey as OsStaticKey;
 214
 215     #[doc(hidden)]
 216     pub struct KeyInner<T> {
 217         pub inner: OsStaticKey,
 218         pub marker: marker::InvariantType<T>,
 219     }
 220
 221     #[doc(hidden)]
 222     impl<T> KeyInner<T> {
 223         #[doc(hidden)]
 224         pub unsafe fn set(&self, ptr: *mut T) { self.inner.set(ptr as *mut _) }
 225         #[doc(hidden)]
 226         pub unsafe fn get(&self) -> *mut T { self.inner.get() as *mut _ }
 227     }
 228 }
 229
 230
 231 #[cfg(test)]
 232 mod tests {
 233     use cell::Cell;
 234     use prelude::*;
 235
 236     #[test]
 237     fn smoke() {
 238         scoped_thread_local!(static BAR: uint)
 239
 240         assert!(!BAR.is_set());
 241         BAR.set(&1, || {
 242             assert!(BAR.is_set());
 243             BAR.with(|slot| {
 244                 assert_eq!(*slot, 1);
 245             });
 246         });
 247         assert!(!BAR.is_set());
 248     }
 249
 250     #[test]
 251     fn cell_allowed() {
 252         scoped_thread_local!(static BAR: Cell<uint>)
 253
 254         BAR.set(&Cell::new(1), || {
 255             BAR.with(|slot| {
 256                 assert_eq!(slot.get(), 1);
 257             });
 258         });
 259     }
 260 }
 261