proc_state/
lib.rs

1#![doc = include_str!("README.md")]
2
3use core::mem::MaybeUninit;
4use std::hash::{DefaultHasher, Hash, Hasher};
5
6#[doc(hidden)]
7pub mod _impl {
8    pub use proc_state_macro::random;
9}
10
11/// Generate [`Global<T>`] at the context. No arguments are supported.
12///
13/// ```
14/// # use std::sync::Mutex;
15/// const COUNTER: proc_state::Global<Mutex<usize>> = proc_state::new!();
16/// ```
17#[macro_export]
18macro_rules! new {
19    () => {
20        unsafe {
21            $crate::Global::_new_inner(
22                ::core::file!(),
23                ::core::line!(),
24                ::core::module_path!(),
25                $crate::_impl::random!(),
26            )
27        }
28    };
29}
30
31/// Owning global state, shared between macro call.
32pub struct Global<T> {
33    file: &'static str,
34    line: u32,
35    module_path: &'static str,
36    random: usize,
37    phantom: core::marker::PhantomData<T>,
38}
39
40impl<T: Send + Sync> Global<T> {
41    #[doc(hidden)]
42    pub const unsafe fn _new_inner(
43        file: &'static str,
44        line: u32,
45        module_path: &'static str,
46        random: usize,
47    ) -> Self {
48        Global {
49            file,
50            line,
51            module_path,
52            random,
53            phantom: core::marker::PhantomData,
54        }
55    }
56
57    fn get_ident_name(&self) -> String {
58        let mut hasher = DefaultHasher::new();
59        std::process::id().hash(&mut hasher);
60        self.file.hash(&mut hasher);
61        self.line.hash(&mut hasher);
62        self.module_path.hash(&mut hasher);
63        self.random.hash(&mut hasher);
64        format!("PROC_STATE_PTR_{}", hasher.finish())
65    }
66
67    /// Get a reference to the inner value, or insert new value if empty.
68    pub fn or_insert(&self, slot: T) -> &'static T {
69        self.or_insert_with(|| slot)
70    }
71
72    /// Get a reference to the inner value.
73    pub fn get(&self) -> Option<&'static T> {
74        let name = self.get_ident_name();
75        if let Ok(var) = std::env::var(&name) {
76            let pointer: Result<usize, _> = var.parse();
77            if let Ok(pointer) = pointer {
78                unsafe {
79                    return Some(std::mem::transmute(pointer));
80                }
81            }
82        }
83        None
84    }
85
86    /// Get a reference to the inner value, or insert new value if empty.
87    pub fn or_insert_with(&self, f: impl FnOnce() -> T) -> &'static T {
88        if let Some(got) = self.get() {
89            return got;
90        }
91        let name = self.get_ident_name();
92        let mut allocated = Box::new(MaybeUninit::<T>::uninit());
93        unsafe { std::env::set_var(&name, (allocated.as_ref() as *const _ as usize).to_string()) }
94        if let Some(got) = self.get() {
95            if got as *const _ as usize == allocated.as_ref() as *const _ as usize {
96                // Panicking in f() has no effect on this function
97                let generated = f();
98                allocated.write(generated);
99                let r = allocated.as_ref() as *const _;
100                std::mem::forget(allocated);
101                unsafe { std::mem::transmute(r) }
102            } else {
103                got
104            }
105        } else {
106            panic!("Cannot set environment variable")
107        }
108    }
109}