binius_utils/
thread_local_mut.rs

1// Copyright 2024-2025 Irreducible Inc.
2
3use std::cell::UnsafeCell;
4
5use thread_local::ThreadLocal;
6
7/// Creates a "scratch space" within each thread with mutable access.
8///
9/// This is mainly meant to be used as an optimization to avoid unnecessary allocs/frees within rayon code.
10/// You only pay for allocation of this scratch space once per thread.
11///
12/// Since the space is local to each thread you also don't have to worry about atomicity.
13#[derive(Debug, Default)]
14pub struct ThreadLocalMut<T: Send>(ThreadLocal<UnsafeCell<T>>);
15
16impl<T: Send> ThreadLocalMut<T> {
17	pub fn new() -> Self {
18		Self(ThreadLocal::new())
19	}
20
21	#[inline]
22	pub fn with_mut<U>(&self, init: impl FnOnce() -> T, run_scope: impl FnOnce(&mut T) -> U) -> U {
23		let data = self.0.get_or(|| UnsafeCell::new(init()));
24		run_scope(unsafe { data.get().as_mut().unwrap_unchecked() })
25	}
26}
binius_utils/thread_local_mut.rs

binius_utils/
thread_local_mut.rs
