hyper_scripter/
state.rs

1use std::cell::UnsafeCell;
2#[cfg(not(feature = "no-state-check"))]
3use std::sync::atomic::{AtomicU8, Ordering::SeqCst};
4
5const UNINITIALIZED: u8 = 0;
6const INITIALIZING: u8 = 1;
7const INITIALIZED: u8 = 2;
8
9pub struct State<T> {
10    data: UnsafeCell<Option<T>>,
11    #[cfg(not(feature = "no-state-check"))]
12    status: AtomicU8,
13}
14unsafe impl<T> Sync for State<T> {}
15
16#[macro_export]
17macro_rules! set_once {
18    ($state:expr, $f:expr) => {{
19        use std::sync::Once;
20        static ONCE: Once = Once::new();
21        ONCE.call_once(|| {
22            $state.set($f());
23        });
24    }};
25}
26
27#[macro_export]
28macro_rules! local_global_state {
29    ($mod_name:ident, $type_name:ident, $test_default:expr) => {
30        mod $mod_name {
31            use super::*;
32            static GLOBAL: $crate::state::State<$type_name> = $crate::state::State::new();
33
34            #[cfg(not(feature = "no-state-check"))]
35            thread_local! {
36                static LOCAL: std::cell::Cell<Option<&'static $type_name>> = std::cell::Cell::new(None);
37            }
38
39            pub fn get() -> &'static $type_name {
40                #[cfg(test)]
41                $crate::set_once!(GLOBAL, $test_default);
42
43                #[cfg(not(feature = "no-state-check"))]
44                if let Some(local) = LOCAL.get() {
45                    return local;
46                }
47
48                GLOBAL.get()
49            }
50
51            pub fn set(data: $type_name) {
52                #[cfg(test)]
53                log::info!("測試中，不設定狀態");
54                #[cfg(not(test))]
55                GLOBAL.set(data);
56            }
57
58            #[cfg(not(feature = "no-state-check"))]
59            pub fn set_local(data: &'static $type_name) {
60                LOCAL.set(Some(data));
61            }
62        }
63    };
64}
65
66impl<T: Sized> State<T> {
67    pub const fn new() -> State<T> {
68        State {
69            data: UnsafeCell::new(None),
70            #[cfg(not(feature = "no-state-check"))]
71            status: AtomicU8::new(UNINITIALIZED),
72        }
73    }
74
75    pub fn set(&self, data: T) {
76        #[cfg(not(feature = "no-state-check"))]
77        {
78            let status = self
79                .status
80                .compare_exchange(UNINITIALIZED, INITIALIZING, SeqCst, SeqCst);
81            log::debug!("設定前的狀態：{:?}", status);
82            if status.is_err() {
83                panic!("多次設定狀態");
84            }
85        }
86        let ptr = self.data.get();
87        unsafe {
88            *ptr = Some(data);
89        }
90        #[cfg(not(feature = "no-state-check"))]
91        self.status.store(INITIALIZED, SeqCst);
92    }
93    pub fn get(&self) -> &T {
94        #[cfg(not(feature = "no-state-check"))]
95        match self.status.load(SeqCst) {
96            UNINITIALIZED => {
97                panic!("還沒設定就取狀態");
98            }
99            INITIALIZING => {
100                while self.status.load(SeqCst) == INITIALIZING {
101                    std::hint::spin_loop();
102                }
103            }
104            _ => (),
105        }
106        let ptr = self.data.get();
107        unsafe { (&*ptr).as_ref().unwrap() }
108    }
109}
hyper_scripter/state.rs

hyper_scripter/
state.rs
