1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

use crate::sync::RwLock;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::{
    cell::{Cell, UnsafeCell},
    fmt,
    marker::PhantomData,
};
pub(crate) struct Local<T> {
    // TODO(eliza): this once used a `crossbeam_util::ShardedRwLock`. We may
    // eventually wish to replace it with a sharded lock implementation on top
    // of our internal `RwLock` wrapper type. If possible, we should profile
    // this first to determine if it's necessary.
    inner: RwLock<Inner<T>>,
}

type Inner<T> = Vec<Option<UnsafeCell<T>>>;

/// Uniquely identifies a thread.
#[repr(transparent)]
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub(crate) struct Id {
    id: usize,
    _not_send: PhantomData<UnsafeCell<()>>,
}

// === impl Local ===

impl<T> Local<T> {
    pub(crate) fn new() -> Self {
        let len = Id::current().as_usize();
        // Preallocate up to the current thread ID, so we don't have to inside
        // the lock.
        let mut data = Vec::with_capacity(len);
        data.resize_with(len, || None);
        Local {
            inner: RwLock::new(data),
        }
    }

    pub(crate) fn with_or_else<O>(
        &self,
        new: impl FnOnce() -> T,
        f: impl FnOnce(&mut T) -> O,
    ) -> Option<O> {
        let i = Id::current().as_usize();
        let mut f = Some(f);
        self.try_with_index(i, |item| f.take().expect("called twice")(item))
            .or_else(move || {
                self.new_thread(i, new);
                self.try_with_index(i, |item| f.take().expect("called twice")(item))
            })
    }

    fn try_with_index<O>(&self, i: usize, f: impl FnOnce(&mut T) -> O) -> Option<O> {
        let lock = try_lock!(self.inner.read(), else return None);
        let slot = lock.get(i)?.as_ref()?;
        let item = unsafe { &mut *slot.get() };
        Some(f(item))
    }

    #[cold]
    fn new_thread(&self, i: usize, new: impl FnOnce() -> T) {
        let mut lock = try_lock!(self.inner.write());
        let this = &mut *lock;
        this.resize_with(i + 1, || None);
        this[i] = Some(UnsafeCell::new(new()));
    }
}

impl<T: Default> Local<T> {
    #[inline]
    pub(crate) fn with<O>(&self, f: impl FnOnce(&mut T) -> O) -> Option<O> {
        self.with_or_else(T::default, f)
    }
}

unsafe impl<T> Sync for Local<T> {}

impl<T: fmt::Debug> fmt::Debug for Local<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let id = Id::current();
        self.try_with_index(id.as_usize(), |local| {
            f.debug_struct("Local")
                .field("thread", &id)
                .field("local", &*local)
                .finish()
        })
        .unwrap_or_else(|| {
            f.debug_struct("Local")
                .field("thread", &id)
                .field("local", &format_args!("<uninitialized>"))
                .finish()
        })
    }
}

// === impl Id ===

impl Id {
    pub(crate) fn current() -> Self {
        thread_local! {
            static MY_ID: Cell<Option<Id>> = Cell::new(None);
        }

        MY_ID
            .try_with(|my_id| my_id.get().unwrap_or_else(|| Self::new_thread(my_id)))
            .unwrap_or_else(|_| Self::poisoned())
    }

    pub(crate) fn as_usize(self) -> usize {
        self.id
    }

    #[cold]
    fn new_thread(local: &Cell<Option<Id>>) -> Self {
        static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
        let id = NEXT_ID.fetch_add(1, Ordering::AcqRel);
        let tid = Self {
            id,
            _not_send: PhantomData,
        };
        local.set(Some(tid));
        tid
    }

    #[cold]
    fn poisoned() -> Self {
        Self {
            id: std::usize::MAX,
            _not_send: PhantomData,
        }
    }

    /// Returns true if the local thread ID was accessed while unwinding.
    pub(crate) fn is_poisoned(self) -> bool {
        self.id == std::usize::MAX
    }
}

impl fmt::Debug for Id {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.is_poisoned() {
            f.debug_tuple("Id")
                .field(&format_args!("<poisoned>"))
                .finish()
        } else {
            f.debug_tuple("Id").field(&self.id).finish()
        }
    }
}