Merge pull request #56 from terrarier2111/master

src/lib.rs

@@ -94,13 +94,14 @@ const POINTER_WIDTH: u8 = 32;
 const POINTER_WIDTH: u8 = 64;
 
 /// The total number of buckets stored in each thread local.
-const BUCKETS: usize = (POINTER_WIDTH + 1) as usize;
+/// All buckets combined can hold up to `usize::MAX - 1` entries.
+const BUCKETS: usize = (POINTER_WIDTH - 1) as usize;
 
 /// Thread-local variable wrapper
 ///
 /// See the [module-level documentation](index.html) for more.
 pub struct ThreadLocal<T: Send> {
-    /// The buckets in the thread local. The nth bucket contains `2^(n-1)`
+    /// The buckets in the thread local. The nth bucket contains `2^n`
     /// elements. Each bucket is lazily allocated.
     buckets: [AtomicPtr<Entry<T>>; BUCKETS],
 
@@ -135,16 +136,11 @@ impl<T: Send> Default for ThreadLocal<T> {
 
 impl<T: Send> Drop for ThreadLocal<T> {
     fn drop(&mut self) {
-        let mut bucket_size = 1;
-
         // Free each non-null bucket
         for (i, bucket) in self.buckets.iter_mut().enumerate() {
             let bucket_ptr = *bucket.get_mut();
 
-            let this_bucket_size = bucket_size;
-            if i != 0 {
-                bucket_size <<= 1;
-            }
+            let this_bucket_size = 1 << i;
 
             if bucket_ptr.is_null() {
                 break;
@@ -165,22 +161,14 @@ impl<T: Send> ThreadLocal<T> {
     /// access the thread local it will never reallocate. The capacity may be rounded up to the
     /// nearest power of two.
     pub fn with_capacity(capacity: usize) -> ThreadLocal<T> {
-        let allocated_buckets = capacity
-            .checked_sub(1)
-            .map(|c| usize::from(POINTER_WIDTH) - (c.leading_zeros() as usize) + 1)
-            .unwrap_or(0);
+        let allocated_buckets = usize::from(POINTER_WIDTH) - (capacity.leading_zeros() as usize);
 
         let mut buckets = [ptr::null_mut(); BUCKETS];
-        let mut bucket_size = 1;
         for (i, bucket) in buckets[..allocated_buckets].iter_mut().enumerate() {
-            *bucket = allocate_bucket::<T>(bucket_size);
-
-            if i != 0 {
-                bucket_size <<= 1;
-            }
+            *bucket = allocate_bucket::<T>(1 << i);
         }
 
-        ThreadLocal {
+        Self {
             // Safety: AtomicPtr has the same representation as a pointer and arrays have the same
             // representation as a sequence of their inner type.
             buckets: unsafe { mem::transmute(buckets) },
@@ -432,9 +420,7 @@ impl RawIter {
 
     #[inline]
     fn next_bucket(&mut self) {
-        if self.bucket != 0 {
-            self.bucket_size <<= 1;
-        }
+        self.bucket_size <<= 1;
         self.bucket += 1;
         self.index = 0;
     }

src/thread_id.rs

@@ -11,7 +11,6 @@ use std::cell::Cell;
 use std::cmp::Reverse;
 use std::collections::BinaryHeap;
 use std::sync::Mutex;
-use std::usize;
 
 /// Thread ID manager which allocates thread IDs. It attempts to aggressively
 /// reuse thread IDs where possible to avoid cases where a ThreadLocal grows
@@ -21,8 +20,8 @@ struct ThreadIdManager {
     free_list: BinaryHeap<Reverse<usize>>,
 }
 impl ThreadIdManager {
-    fn new() -> ThreadIdManager {
-        ThreadIdManager {
+    fn new() -> Self {
+        Self {
             free_from: 0,
             free_list: BinaryHeap::new(),
         }
@@ -31,11 +30,11 @@ impl ThreadIdManager {
         if let Some(id) = self.free_list.pop() {
             id.0
         } else {
+            // `free_from` can't overflow as each thread takes up at least 2 bytes of memory and
+            // thus we can't even have `usize::MAX / 2 + 1` threads.
+
             let id = self.free_from;
-            self.free_from = self
-                .free_from
-                .checked_add(1)
-                .expect("Ran out of thread IDs");
+            self.free_from += 1;
             id
         }
     }
@@ -60,12 +59,12 @@ pub(crate) struct Thread {
     pub(crate) index: usize,
 }
 impl Thread {
-    fn new(id: usize) -> Thread {
-        let bucket = usize::from(POINTER_WIDTH) - id.leading_zeros() as usize;
-        let bucket_size = 1 << bucket.saturating_sub(1);
-        let index = if id != 0 { id ^ bucket_size } else { 0 };
+    fn new(id: usize) -> Self {
+        let bucket = usize::from(POINTER_WIDTH) - ((id + 1).leading_zeros() as usize) - 1;
+        let bucket_size = 1 << bucket;
+        let index = id - (bucket_size - 1);
 
-        Thread {
+        Self {
             id,
             bucket,
             bucket_size,
@@ -184,24 +183,24 @@ fn test_thread() {
     let thread = Thread::new(1);
     assert_eq!(thread.id, 1);
     assert_eq!(thread.bucket, 1);
-    assert_eq!(thread.bucket_size, 1);
+    assert_eq!(thread.bucket_size, 2);
     assert_eq!(thread.index, 0);
 
     let thread = Thread::new(2);
     assert_eq!(thread.id, 2);
-    assert_eq!(thread.bucket, 2);
+    assert_eq!(thread.bucket, 1);
     assert_eq!(thread.bucket_size, 2);
-    assert_eq!(thread.index, 0);
+    assert_eq!(thread.index, 1);
 
     let thread = Thread::new(3);
     assert_eq!(thread.id, 3);
     assert_eq!(thread.bucket, 2);
-    assert_eq!(thread.bucket_size, 2);
-    assert_eq!(thread.index, 1);
+    assert_eq!(thread.bucket_size, 4);
+    assert_eq!(thread.index, 0);
 
     let thread = Thread::new(19);
     assert_eq!(thread.id, 19);
-    assert_eq!(thread.bucket, 5);
+    assert_eq!(thread.bucket, 4);
     assert_eq!(thread.bucket_size, 16);
-    assert_eq!(thread.index, 3);
+    assert_eq!(thread.index, 4);
 }