From 7cd98bbfd7a374a005eda423eeace5539ec0a066 Mon Sep 17 00:00:00 2001
From: Haibo Huang <hhb@google.com>
Date: Thu, 7 Jan 2021 18:06:50 -0800
Subject: Upgrade rust/crates/thread_local to 1.1.0

Test: make
Change-Id: I0c0348cb22330ebd6c6be633525da259fe14fb0c
---
 src/thread_id.rs | 110 ++++++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 88 insertions(+), 22 deletions(-)

(limited to 'src/thread_id.rs')

diff --git a/src/thread_id.rs b/src/thread_id.rs
index e757948..397f772 100644
--- a/src/thread_id.rs
+++ b/src/thread_id.rs
@@ -5,57 +5,123 @@
 // http://opensource.org/licenses/MIT>, at your option. This file may not be
 // copied, modified, or distributed except according to those terms.
 
+use std::cmp::Reverse;
 use std::collections::BinaryHeap;
 use std::sync::Mutex;
 use std::usize;
+use POINTER_WIDTH;
 
-// Thread ID manager which allocates thread IDs. It attempts to aggressively
-// reuse thread IDs where possible to avoid cases where a ThreadLocal grows
-// indefinitely when it is used by many short-lived threads.
+/// Thread ID manager which allocates thread IDs. It attempts to aggressively
+/// reuse thread IDs where possible to avoid cases where a ThreadLocal grows
+/// indefinitely when it is used by many short-lived threads.
 struct ThreadIdManager {
-    limit: usize,
-    free_list: BinaryHeap<usize>,
+    free_from: usize,
+    free_list: BinaryHeap<Reverse<usize>>,
 }
 impl ThreadIdManager {
     fn new() -> ThreadIdManager {
         ThreadIdManager {
-            limit: usize::MAX,
+            free_from: 0,
             free_list: BinaryHeap::new(),
         }
     }
     fn alloc(&mut self) -> usize {
         if let Some(id) = self.free_list.pop() {
-            id
+            id.0
         } else {
-            let id = self.limit;
-            self.limit = self.limit.checked_sub(1).expect("Ran out of thread IDs");
+            let id = self.free_from;
+            self.free_from = self
+                .free_from
+                .checked_add(1)
+                .expect("Ran out of thread IDs");
             id
         }
     }
     fn free(&mut self, id: usize) {
-        self.free_list.push(id);
+        self.free_list.push(Reverse(id));
     }
 }
 lazy_static! {
     static ref THREAD_ID_MANAGER: Mutex<ThreadIdManager> = Mutex::new(ThreadIdManager::new());
 }
 
-// Non-zero integer which is unique to the current thread while it is running.
-// A thread ID may be reused after a thread exits.
-struct ThreadId(usize);
-impl ThreadId {
-    fn new() -> ThreadId {
-        ThreadId(THREAD_ID_MANAGER.lock().unwrap().alloc())
+/// Data which is unique to the current thread while it is running.
+/// A thread ID may be reused after a thread exits.
+#[derive(Clone, Copy)]
+pub(crate) struct Thread {
+    /// The thread ID obtained from the thread ID manager.
+    pub(crate) id: usize,
+    /// The bucket this thread's local storage will be in.
+    pub(crate) bucket: usize,
+    /// The size of the bucket this thread's local storage will be in.
+    pub(crate) bucket_size: usize,
+    /// The index into the bucket this thread's local storage is in.
+    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 };
+
+        Thread {
+            id,
+            bucket,
+            bucket_size,
+            index,
+        }
+    }
+}
+
+/// Wrapper around `Thread` that allocates and deallocates the ID.
+struct ThreadHolder(Thread);
+impl ThreadHolder {
+    fn new() -> ThreadHolder {
+        ThreadHolder(Thread::new(THREAD_ID_MANAGER.lock().unwrap().alloc()))
     }
 }
-impl Drop for ThreadId {
+impl Drop for ThreadHolder {
     fn drop(&mut self) {
-        THREAD_ID_MANAGER.lock().unwrap().free(self.0);
+        THREAD_ID_MANAGER.lock().unwrap().free(self.0.id);
     }
 }
-thread_local!(static THREAD_ID: ThreadId = ThreadId::new());
 
-/// Returns a non-zero ID for the current thread
-pub fn get() -> usize {
-    THREAD_ID.with(|x| x.0)
+thread_local!(static THREAD_HOLDER: ThreadHolder = ThreadHolder::new());
+
+/// Get the current thread.
+pub(crate) fn get() -> Thread {
+    THREAD_HOLDER.with(|holder| holder.0)
+}
+
+#[test]
+fn test_thread() {
+    let thread = Thread::new(0);
+    assert_eq!(thread.id, 0);
+    assert_eq!(thread.bucket, 0);
+    assert_eq!(thread.bucket_size, 1);
+    assert_eq!(thread.index, 0);
+
+    let thread = Thread::new(1);
+    assert_eq!(thread.id, 1);
+    assert_eq!(thread.bucket, 1);
+    assert_eq!(thread.bucket_size, 1);
+    assert_eq!(thread.index, 0);
+
+    let thread = Thread::new(2);
+    assert_eq!(thread.id, 2);
+    assert_eq!(thread.bucket, 2);
+    assert_eq!(thread.bucket_size, 2);
+    assert_eq!(thread.index, 0);
+
+    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);
+
+    let thread = Thread::new(19);
+    assert_eq!(thread.id, 19);
+    assert_eq!(thread.bucket, 5);
+    assert_eq!(thread.bucket_size, 16);
+    assert_eq!(thread.index, 3);
 }
-- 
cgit v1.2.3