Revision control
Copy as Markdown
Other Tools
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
use std::{
fmt::Debug,
ops::{Deref, DerefMut},
};
use crate::util::sync::Mutex;
// Note: this is meant to be used only in low-contention
// scenarios.
pub struct PerThreadStorage<T> {
storage: Mutex<Vec<T>>,
init: fn() -> T,
}
impl<T> Debug for PerThreadStorage<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "<thread storage>")
}
}
pub struct PerThreadStorageRef<'a, T> {
r: &'a PerThreadStorage<T>,
val: Option<T>,
}
impl<T> PerThreadStorage<T> {
pub fn new(init: fn() -> T) -> Self {
Self {
storage: Mutex::new(vec![]),
init,
}
}
pub fn get(&self) -> PerThreadStorageRef<'_, T> {
let t = {
let mut a = self.storage.lock().unwrap();
if let Some(x) = a.pop() {
x
} else {
drop(a);
(self.init)()
}
};
PerThreadStorageRef {
r: self,
val: Some(t),
}
}
}
impl<'a, T> Drop for PerThreadStorageRef<'a, T> {
fn drop(&mut self) {
let v = self.val.take().unwrap();
let mut a = self.r.storage.lock().unwrap();
a.push(v);
}
}
impl<'a, T> Deref for PerThreadStorageRef<'a, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
self.val.as_ref().unwrap()
}
}
impl<'a, T> DerefMut for PerThreadStorageRef<'a, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
self.val.as_mut().unwrap()
}
}