Use MaybeUninit on rustc >= 1.36.0

Author: Kestrer

build.rs

@@ -0,0 +1,48 @@
+use std::error::Error;
+use std::ffi::OsString;
+use std::process::{exit, Command};
+
+fn main() {
+    let rustc = std::env::var_os("RUSTC").unwrap_or_else(|| OsString::from("rustc"));
+    let command = format!("`{} --version`", rustc.to_string_lossy());
+
+    let output = Command::new(&rustc)
+        .arg("--version")
+        .output()
+        .unwrap_or_else(|e| {
+            eprintln!("Error: failed to run {}: {}", command, e);
+            exit(1)
+        });
+
+    let supports_maybe_uninit = parse(&output.stdout).unwrap_or_else(|e| {
+        eprintln!("Error: failed to parse output of {}: {}", command, e);
+        exit(1)
+    });
+
+    if supports_maybe_uninit {
+        println!("cargo:rustc-cfg=supports_maybe_uninit");
+    }
+}
+
+fn parse(output: &[u8]) -> Result<bool, Box<dyn Error>> {
+    let s = std::str::from_utf8(output)?;
+    let last_line = s.lines().last().unwrap_or(s);
+    let mut words = last_line.trim().split(' ');
+    if words.next() != Some("rustc") {
+        return Err("version does not start with 'rustc'".into());
+    }
+    let mut triplet = words
+        .next()
+        .ok_or("output does not contain version triplet")?
+        .split('.');
+    if triplet.next() != Some("1") {
+        return Err("rustc major version is not 1".into());
+    }
+    let minor: u32 = triplet
+        .next()
+        .ok_or("rustc version does not contain minor version")?
+        .parse()
+        .map_err(|e| format!("failed to parse minor version: {}", e))?;
+
+    Ok(minor >= 36)
+}

src/lib.rs

@@ -70,12 +70,14 @@
 extern crate lazy_static;
 
 mod cached;
+mod maybe_uninit;
 mod thread_id;
 mod unreachable;
 
 #[allow(deprecated)]
 pub use cached::{CachedIntoIter, CachedIterMut, CachedThreadLocal};
 
+use maybe_uninit::MaybeUninit;
 use std::cell::UnsafeCell;
 use std::fmt;
 use std::iter::FusedIterator;
@@ -86,7 +88,7 @@ use std::ptr;
 use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering};
 use std::sync::Mutex;
 use thread_id::Thread;
-use unreachable::{UncheckedOptionExt, UncheckedResultExt};
+use unreachable::UncheckedResultExt;
 
 // Use usize::BITS once it has stabilized and the MSRV has been bumped.
 #[cfg(target_pointer_width = "16")]
@@ -119,8 +121,7 @@ pub struct ThreadLocal<T: Send> {
 
 struct Entry<T> {
     present: AtomicBool,
-    // Use MaybeUninit once the MSRV has been bumped.
-    value: UnsafeCell<Option<T>>,
+    value: UnsafeCell<MaybeUninit<T>>,
 }
 
 // ThreadLocal is always Sync, even if T isn't
@@ -226,7 +227,15 @@ impl<T: Send> ThreadLocal<T> {
         if bucket_ptr.is_null() {
             return None;
         }
-        unsafe { (&*(*bucket_ptr.add(thread.index)).value.get()).as_ref() }
+        unsafe {
+            let entry = &*bucket_ptr.add(thread.index);
+            // Read without atomic operations as only this thread can set the value.
+            if (&entry.present as *const _ as *const bool).read() {
+                Some(&*(&*entry.value.get()).as_ptr())
+            } else {
+                None
+            }
+        }
     }
 
     #[cold]
@@ -251,12 +260,12 @@ impl<T: Send> ThreadLocal<T> {
         // Insert the new element into the bucket
         let entry = unsafe { &*bucket_ptr.add(thread.index) };
         let value_ptr = entry.value.get();
-        unsafe { value_ptr.write(Some(data)) };
+        unsafe { value_ptr.write(MaybeUninit::new(data)) };
         entry.present.store(true, Ordering::Release);
 
         self.values.fetch_add(1, Ordering::Release);
 
-        unsafe { (&*value_ptr).as_ref().unchecked_unwrap() }
+        unsafe { &*(&*value_ptr).as_ptr() }
     }
 
     /// Returns an iterator over the local values of all threads in unspecified
@@ -370,7 +379,7 @@ impl<'a, T: Send + Sync> Iterator for Iter<'a, T> {
                     self.index += 1;
                     if entry.present.load(Ordering::Acquire) {
                         self.yielded += 1;
-                        return Some(unsafe { (&*entry.value.get()).as_ref().unchecked_unwrap() });
+                        return Some(unsafe { &*(&*entry.value.get()).as_ptr() });
                     }
                 }
             }
@@ -401,7 +410,7 @@ struct RawIterMut<T: Send> {
 }
 
 impl<T: Send> Iterator for RawIterMut<T> {
-    type Item = *mut Option<T>;
+    type Item = *mut MaybeUninit<T>;
 
     fn next(&mut self) -> Option<Self::Item> {
         if self.remaining == 0 {
@@ -448,7 +457,7 @@ impl<'a, T: Send + 'a> Iterator for IterMut<'a, T> {
     fn next(&mut self) -> Option<&'a mut T> {
         self.raw
             .next()
-            .map(|x| unsafe { &mut *(*x).as_mut().unchecked_unwrap() })
+            .map(|x| unsafe { &mut *(&mut *x).as_mut_ptr() })
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
@@ -470,7 +479,7 @@ impl<T: Send> Iterator for IntoIter<T> {
     fn next(&mut self) -> Option<T> {
         self.raw
             .next()
-            .map(|x| unsafe { (*x).take().unchecked_unwrap() })
+            .map(|x| unsafe { std::mem::replace(&mut *x, MaybeUninit::uninit()).assume_init() })
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
@@ -485,7 +494,7 @@ fn allocate_bucket<T>(size: usize) -> *mut Entry<T> {
         (0..size)
             .map(|_| Entry::<T> {
                 present: AtomicBool::new(false),
-                value: UnsafeCell::new(None),
+                value: UnsafeCell::new(MaybeUninit::uninit()),
             })
             .collect::<Vec<_>>()
             .into_boxed_slice(),

src/maybe_uninit.rs

@@ -0,0 +1,39 @@
+#[cfg(supports_maybe_uninit)]
+pub(crate) use std::mem::MaybeUninit;
+
+#[cfg(not(supports_maybe_uninit))]
+mod polyfill {
+    use std::mem::ManuallyDrop;
+    use std::ptr;
+
+    use unreachable::UncheckedOptionExt;
+
+    /// A simple `Option`-based implementation of `MaybeUninit` for compiler versions < 1.36.0.
+    pub struct MaybeUninit<T>(Option<ManuallyDrop<T>>);
+
+    impl<T> MaybeUninit<T> {
+        pub fn new(val: T) -> Self {
+            MaybeUninit(Some(ManuallyDrop::new(val)))
+        }
+        pub fn uninit() -> Self {
+            MaybeUninit(None)
+        }
+        pub fn as_ptr(&self) -> *const T {
+            self.0
+                .as_ref()
+                .map(|v| &**v as *const _)
+                .unwrap_or_else(ptr::null)
+        }
+        pub fn as_mut_ptr(&mut self) -> *mut T {
+            self.0
+                .as_mut()
+                .map(|v| &mut **v as *mut _)
+                .unwrap_or_else(ptr::null_mut)
+        }
+        pub unsafe fn assume_init(self) -> T {
+            ManuallyDrop::into_inner(self.0.unchecked_unwrap())
+        }
+    }
+}
+#[cfg(not(supports_maybe_uninit))]
+pub(crate) use self::polyfill::MaybeUninit;