1use crate::{
2 cfg::{self, CfgPrivate},
3 clear::Clear,
4 page,
5 sync::{
6 alloc,
7 atomic::{
8 AtomicPtr, AtomicUsize,
9 Ordering::{self, *},
10 },
11 },
12 tid::Tid,
13 Pack,
14};
15
16use std::{fmt, ptr, slice};
17
18// ┌─────────────┐ ┌────────┐
19// │ page 1 │ │ │
20// ├─────────────┤ ┌───▶│ next──┼─┐
21// │ page 2 │ │ ├────────┤ │
22// │ │ │ │XXXXXXXX│ │
23// │ local_free──┼─┘ ├────────┤ │
24// │ global_free─┼─┐ │ │◀┘
25// ├─────────────┤ └───▶│ next──┼─┐
26// │ page 3 │ ├────────┤ │
27// └─────────────┘ │XXXXXXXX│ │
28// ... ├────────┤ │
29// ┌─────────────┐ │XXXXXXXX│ │
30// │ page n │ ├────────┤ │
31// └─────────────┘ │ │◀┘
32// │ next──┼───▶
33// ├────────┤
34// │XXXXXXXX│
35// └────────┘
36// ...
37pub(crate) struct Shard<T, C: cfg::Config> {
38 /// The shard's parent thread ID.
39 pub(crate) tid: usize,
40 /// The local free list for each page.
41 ///
42 /// These are only ever accessed from this shard's thread, so they are
43 /// stored separately from the shared state for the page that can be
44 /// accessed concurrently, to minimize false sharing.
45 local: Box<[page::Local]>,
46 /// The shared state for each page in this shard.
47 ///
48 /// This consists of the page's metadata (size, previous size), remote free
49 /// list, and a pointer to the actual array backing that page.
50 shared: Box<[page::Shared<T, C>]>,
51}
52
53pub(crate) struct Array<T, C: cfg::Config> {
54 shards: Box<[Ptr<T, C>]>,
55 max: AtomicUsize,
56}
57
58#[derive(Debug)]
59struct Ptr<T, C: cfg::Config>(AtomicPtr<alloc::Track<Shard<T, C>>>);
60
61#[derive(Debug)]
62pub(crate) struct IterMut<'a, T: 'a, C: cfg::Config + 'a>(slice::IterMut<'a, Ptr<T, C>>);
63
64// === impl Shard ===
65
66impl<T, C> Shard<T, C>
67where
68 C: cfg::Config,
69{
70 #[inline(always)]
71 pub(crate) fn with_slot<'a, U>(
72 &'a self,
73 idx: usize,
74 f: impl FnOnce(&'a page::Slot<T, C>) -> Option<U>,
75 ) -> Option<U> {
76 debug_assert_eq_in_drop!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
77 let (addr, page_index) = page::indices::<C>(idx);
78
79 test_println!("-> {:?}", addr);
80 if page_index >= self.shared.len() {
81 return None;
82 }
83
84 self.shared[page_index].with_slot(addr, f)
85 }
86
87 pub(crate) fn new(tid: usize) -> Self {
88 let mut total_sz = 0;
89 let shared = (0..C::MAX_PAGES)
90 .map(|page_num| {
91 let sz = C::page_size(page_num);
92 let prev_sz = total_sz;
93 total_sz += sz;
94 page::Shared::new(sz, prev_sz)
95 })
96 .collect();
97 let local = (0..C::MAX_PAGES).map(|_| page::Local::new()).collect();
98 Self { tid, local, shared }
99 }
100}
101
102impl<T, C> Shard<Option<T>, C>
103where
104 C: cfg::Config,
105{
106 /// Remove an item on the shard's local thread.
107 pub(crate) fn take_local(&self, idx: usize) -> Option<T> {
108 debug_assert_eq_in_drop!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
109 let (addr, page_index) = page::indices::<C>(idx);
110
111 test_println!("-> remove_local {:?}", addr);
112
113 self.shared
114 .get(page_index)?
115 .take(addr, C::unpack_gen(idx), self.local(page_index))
116 }
117
118 /// Remove an item, while on a different thread from the shard's local thread.
119 pub(crate) fn take_remote(&self, idx: usize) -> Option<T> {
120 debug_assert_eq_in_drop!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
121 debug_assert!(Tid::<C>::current().as_usize() != self.tid);
122
123 let (addr, page_index) = page::indices::<C>(idx);
124
125 test_println!("-> take_remote {:?}; page {:?}", addr, page_index);
126
127 let shared = self.shared.get(page_index)?;
128 shared.take(addr, C::unpack_gen(idx), shared.free_list())
129 }
130
131 pub(crate) fn remove_local(&self, idx: usize) -> bool {
132 debug_assert_eq_in_drop!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
133 let (addr, page_index) = page::indices::<C>(idx);
134
135 if page_index >= self.shared.len() {
136 return false;
137 }
138
139 self.shared[page_index].remove(addr, C::unpack_gen(idx), self.local(page_index))
140 }
141
142 pub(crate) fn remove_remote(&self, idx: usize) -> bool {
143 debug_assert_eq_in_drop!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
144 let (addr, page_index) = page::indices::<C>(idx);
145
146 if page_index >= self.shared.len() {
147 return false;
148 }
149
150 let shared = &self.shared[page_index];
151 shared.remove(addr, C::unpack_gen(idx), shared.free_list())
152 }
153
154 pub(crate) fn iter(&self) -> std::slice::Iter<'_, page::Shared<Option<T>, C>> {
155 self.shared.iter()
156 }
157}
158
159impl<T, C> Shard<T, C>
160where
161 T: Clear + Default,
162 C: cfg::Config,
163{
164 pub(crate) fn init_with<U>(
165 &self,
166 mut init: impl FnMut(usize, &page::Slot<T, C>) -> Option<U>,
167 ) -> Option<U> {
168 // Can we fit the value into an exist`ing page?
169 for (page_idx, page) in self.shared.iter().enumerate() {
170 let local = self.local(page_idx);
171
172 test_println!("-> page {}; {:?}; {:?}", page_idx, local, page);
173
174 if let Some(res) = page.init_with(local, &mut init) {
175 return Some(res);
176 }
177 }
178
179 None
180 }
181
182 pub(crate) fn mark_clear_local(&self, idx: usize) -> bool {
183 debug_assert_eq_in_drop!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
184 let (addr, page_index) = page::indices::<C>(idx);
185
186 if page_index >= self.shared.len() {
187 return false;
188 }
189
190 self.shared[page_index].mark_clear(addr, C::unpack_gen(idx), self.local(page_index))
191 }
192
193 pub(crate) fn mark_clear_remote(&self, idx: usize) -> bool {
194 debug_assert_eq_in_drop!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
195 let (addr, page_index) = page::indices::<C>(idx);
196
197 if page_index >= self.shared.len() {
198 return false;
199 }
200
201 let shared = &self.shared[page_index];
202 shared.mark_clear(addr, C::unpack_gen(idx), shared.free_list())
203 }
204
205 pub(crate) fn clear_after_release(&self, idx: usize) {
206 crate::sync::atomic::fence(crate::sync::atomic::Ordering::Acquire);
207 let tid = Tid::<C>::current().as_usize();
208 test_println!(
209 "-> clear_after_release; self.tid={:?}; current.tid={:?};",
210 tid,
211 self.tid
212 );
213 if tid == self.tid {
214 self.clear_local(idx);
215 } else {
216 self.clear_remote(idx);
217 }
218 }
219
220 fn clear_local(&self, idx: usize) -> bool {
221 debug_assert_eq_in_drop!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
222 let (addr, page_index) = page::indices::<C>(idx);
223
224 if page_index >= self.shared.len() {
225 return false;
226 }
227
228 self.shared[page_index].clear(addr, C::unpack_gen(idx), self.local(page_index))
229 }
230
231 fn clear_remote(&self, idx: usize) -> bool {
232 debug_assert_eq_in_drop!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
233 let (addr, page_index) = page::indices::<C>(idx);
234
235 if page_index >= self.shared.len() {
236 return false;
237 }
238
239 let shared = &self.shared[page_index];
240 shared.clear(addr, C::unpack_gen(idx), shared.free_list())
241 }
242
243 #[inline(always)]
244 fn local(&self, i: usize) -> &page::Local {
245 #[cfg(debug_assertions)]
246 debug_assert_eq_in_drop!(
247 Tid::<C>::current().as_usize(),
248 self.tid,
249 "tried to access local data from another thread!"
250 );
251
252 &self.local[i]
253 }
254}
255
256impl<T: fmt::Debug, C: cfg::Config> fmt::Debug for Shard<T, C> {
257 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
258 let mut d = f.debug_struct("Shard");
259
260 #[cfg(debug_assertions)]
261 d.field("tid", &self.tid);
262 d.field("shared", &self.shared).finish()
263 }
264}
265
266// === impl Array ===
267
268impl<T, C> Array<T, C>
269where
270 C: cfg::Config,
271{
272 pub(crate) fn new() -> Self {
273 let mut shards = Vec::with_capacity(C::MAX_SHARDS);
274 for _ in 0..C::MAX_SHARDS {
275 // XXX(eliza): T_T this could be avoided with maybeuninit or something...
276 shards.push(Ptr::null());
277 }
278 Self {
279 shards: shards.into(),
280 max: AtomicUsize::new(0),
281 }
282 }
283
284 #[inline]
285 pub(crate) fn get(&self, idx: usize) -> Option<&Shard<T, C>> {
286 test_println!("-> get shard={}", idx);
287 self.shards.get(idx)?.load(Acquire)
288 }
289
290 #[inline]
291 pub(crate) fn current(&self) -> (Tid<C>, &Shard<T, C>) {
292 let tid = Tid::<C>::current();
293 test_println!("current: {:?}", tid);
294 let idx = tid.as_usize();
295 assert!(
296 idx < self.shards.len(),
297 "Thread count overflowed the configured max count. \
298 Thread index = {}, max threads = {}.",
299 idx,
300 C::MAX_SHARDS,
301 );
302 // It's okay for this to be relaxed. The value is only ever stored by
303 // the thread that corresponds to the index, and we are that thread.
304 let shard = self.shards[idx].load(Relaxed).unwrap_or_else(|| {
305 let ptr = Box::into_raw(Box::new(alloc::Track::new(Shard::new(idx))));
306 test_println!("-> allocated new shard for index {} at {:p}", idx, ptr);
307 self.shards[idx].set(ptr);
308 let mut max = self.max.load(Acquire);
309 while max < idx {
310 match self.max.compare_exchange(max, idx, AcqRel, Acquire) {
311 Ok(_) => break,
312 Err(actual) => max = actual,
313 }
314 }
315 test_println!("-> highest index={}, prev={}", std::cmp::max(max, idx), max);
316 unsafe {
317 // Safety: we just put it there!
318 &*ptr
319 }
320 .get_ref()
321 });
322 (tid, shard)
323 }
324
325 pub(crate) fn iter_mut(&mut self) -> IterMut<'_, T, C> {
326 test_println!("Array::iter_mut");
327 let max = self.max.load(Acquire);
328 test_println!("-> highest index={}", max);
329 IterMut(self.shards[0..=max].iter_mut())
330 }
331}
332
333impl<T, C: cfg::Config> Drop for Array<T, C> {
334 fn drop(&mut self) {
335 // XXX(eliza): this could be `with_mut` if we wanted to impl a wrapper for std atomics to change `get_mut` to `with_mut`...
336 let max = self.max.load(Acquire);
337 for shard in &self.shards[0..=max] {
338 // XXX(eliza): this could be `with_mut` if we wanted to impl a wrapper for std atomics to change `get_mut` to `with_mut`...
339 let ptr = shard.0.load(Acquire);
340 if ptr.is_null() {
341 continue;
342 }
343 let shard = unsafe {
344 // Safety: this is the only place where these boxes are
345 // deallocated, and we have exclusive access to the shard array,
346 // because...we are dropping it...
347 Box::from_raw(ptr)
348 };
349 drop(shard)
350 }
351 }
352}
353
354impl<T: fmt::Debug, C: cfg::Config> fmt::Debug for Array<T, C> {
355 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
356 let max = self.max.load(Acquire);
357 let mut set = f.debug_map();
358 for shard in &self.shards[0..=max] {
359 let ptr = shard.0.load(Acquire);
360 if let Some(shard) = ptr::NonNull::new(ptr) {
361 set.entry(&format_args!("{:p}", ptr), unsafe { shard.as_ref() });
362 } else {
363 set.entry(&format_args!("{:p}", ptr), &());
364 }
365 }
366 set.finish()
367 }
368}
369
370// === impl Ptr ===
371
372impl<T, C: cfg::Config> Ptr<T, C> {
373 #[inline]
374 fn null() -> Self {
375 Self(AtomicPtr::new(ptr::null_mut()))
376 }
377
378 #[inline]
379 fn load(&self, order: Ordering) -> Option<&Shard<T, C>> {
380 let ptr = self.0.load(order);
381 test_println!("---> loaded={:p} (order={:?})", ptr, order);
382 if ptr.is_null() {
383 test_println!("---> null");
384 return None;
385 }
386 let track = unsafe {
387 // Safety: The returned reference will have the same lifetime as the
388 // reference to the shard pointer, which (morally, if not actually)
389 // owns the shard. The shard is only deallocated when the shard
390 // array is dropped, and it won't be dropped while this pointer is
391 // borrowed --- and the returned reference has the same lifetime.
392 //
393 // We know that the pointer is not null, because we just
394 // null-checked it immediately prior.
395 &*ptr
396 };
397
398 Some(track.get_ref())
399 }
400
401 #[inline]
402 fn set(&self, new: *mut alloc::Track<Shard<T, C>>) {
403 self.0
404 .compare_exchange(ptr::null_mut(), new, AcqRel, Acquire)
405 .expect("a shard can only be inserted by the thread that owns it, this is a bug!");
406 }
407}
408
409// === Iterators ===
410
411impl<'a, T, C> Iterator for IterMut<'a, T, C>
412where
413 T: 'a,
414 C: cfg::Config + 'a,
415{
416 type Item = &'a Shard<T, C>;
417 fn next(&mut self) -> Option<Self::Item> {
418 test_println!("IterMut::next");
419 loop {
420 // Skip over empty indices if they are less than the highest
421 // allocated shard. Some threads may have accessed the slab
422 // (generating a thread ID) but never actually inserted data, so
423 // they may have never allocated a shard.
424 let next = self.0.next();
425 test_println!("-> next.is_some={}", next.is_some());
426 if let Some(shard) = next?.load(Acquire) {
427 test_println!("-> done");
428 return Some(shard);
429 }
430 }
431 }
432}
433