1#[cfg(feature = "alloc")]
14use alloc::vec::Vec;
15#[cfg(feature = "std")]
16use core::sync::atomic::{AtomicUsize, Ordering};
17#[cfg(feature = "std")]
18use std::cell::Cell;
19
20#[cfg(feature = "std")]
21std::thread_local! {
22 static CORE_TICKET: Cell<Option<usize>> = const { Cell::new(None) };
27}
28
29#[cfg(feature = "std")]
30static NEXT_CORE: AtomicUsize = AtomicUsize::new(0);
31
32#[cfg(feature = "std")]
34fn local_index(count: usize) -> usize {
35 let ticket = CORE_TICKET.with(|cell| {
36 cell.get().unwrap_or_else(|| {
37 let assigned = NEXT_CORE.fetch_add(1, Ordering::Relaxed);
38 cell.set(Some(assigned));
39 assigned
40 })
41 });
42 ticket % count
43}
44
45#[cfg(feature = "alloc")]
48pub struct PerCore<T> {
49 slots: Vec<T>,
50 count: usize,
51}
52
53#[cfg(feature = "alloc")]
54impl<T> PerCore<T> {
55 pub fn new_with(count: usize, factory: impl FnMut(usize) -> T) -> Self {
56 let slots = (0..count).map(factory).collect();
57 Self { slots, count }
58 }
59
60 pub fn from_vec(slots: Vec<T>) -> Self {
61 let count = slots.len();
62 Self { slots, count }
63 }
64
65 #[cfg(feature = "std")]
69 pub fn local(&self) -> &T {
70 &self.slots[local_index(self.count)]
71 }
72
73 pub fn slot(&self, index: usize) -> &T {
74 &self.slots[index]
75 }
76
77 #[must_use]
78 pub fn count(&self) -> usize {
79 self.count
80 }
81}
82
83pub struct StaticPerCore<T, const N: usize> {
86 slots: [T; N],
87}
88
89impl<T, const N: usize> StaticPerCore<T, N> {
90 pub fn new_with(factory: impl FnMut(usize) -> T) -> Self {
91 Self {
92 slots: core::array::from_fn(factory),
93 }
94 }
95
96 #[cfg(feature = "std")]
100 pub fn local(&self) -> &T {
101 &self.slots[local_index(N)]
102 }
103
104 pub fn slot(&self, index: usize) -> &T {
105 &self.slots[index]
106 }
107
108 #[must_use]
109 pub const fn count(&self) -> usize {
110 N
111 }
112}
113
114#[cfg(test)]
115#[allow(clippy::unwrap_used, clippy::expect_used)]
116mod tests {
117 use super::*;
118
119 #[test]
120 fn static_per_core_routes_by_index() {
121 let cores = StaticPerCore::<usize, 4>::new_with(|index| index * 10);
122 assert_eq!(cores.count(), 4);
123 assert_eq!(*cores.slot(0), 0);
124 assert_eq!(*cores.slot(3), 30);
125 }
126
127 #[cfg(feature = "alloc")]
128 #[test]
129 fn per_core_new_with_and_from_vec() {
130 let built = PerCore::new_with(3, |index| index + 1);
131 assert_eq!(built.count(), 3);
132 assert_eq!(*built.slot(2), 3);
133
134 let from = PerCore::from_vec(alloc::vec![10usize, 20]);
135 assert_eq!(from.count(), 2);
136 assert_eq!(*from.slot(1), 20);
137 }
138
139 #[cfg(feature = "std")]
140 #[test]
141 fn local_is_sticky_within_a_thread() {
142 let cores = StaticPerCore::<usize, 8>::new_with(|index| index);
143 let first = *cores.local();
146 assert_eq!(*cores.local(), first);
147 assert!(first < cores.count());
148 }
149}