mmtk/util/rust_util/
mod.rs

1//! This module works around limitations of the Rust programming language, and provides missing
2//! functionalities that we may expect the Rust programming language and its standard libraries
3//! to provide.
4
5pub mod atomic_box;
6pub mod rev_group;
7pub mod zeroed_alloc;
8
9/// Const function for min value of two usize numbers.
10pub const fn min_of_usize(a: usize, b: usize) -> usize {
11    if a > b {
12        b
13    } else {
14        a
15    }
16}
17
18#[cfg(feature = "nightly")]
19pub use core::intrinsics::{likely, unlikely};
20
21// likely() and unlikely() compiler hints in stable Rust
22// [1]: https://github.com/rust-lang/hashbrown/blob/a41bd76de0a53838725b997c6085e024c47a0455/src/raw/mod.rs#L48-L70
23// [2]: https://users.rust-lang.org/t/compiler-hint-for-unlikely-likely-for-if-branches/62102/3
24#[cfg(not(feature = "nightly"))]
25#[inline]
26#[cold]
27fn cold() {}
28
29#[cfg(not(feature = "nightly"))]
30#[inline]
31pub fn likely(b: bool) -> bool {
32    if !b {
33        cold();
34    }
35    b
36}
37#[cfg(not(feature = "nightly"))]
38#[inline]
39pub fn unlikely(b: bool) -> bool {
40    if b {
41        cold();
42    }
43    b
44}
45
46use std::cell::UnsafeCell;
47use std::mem::MaybeUninit;
48use std::sync::Once;
49
50/// InitializeOnce creates an uninitialized value that needs to be manually initialized later. InitializeOnce
51/// guarantees the value is only initialized once. This type is used to allow more efficient reads.
52/// Unlike the `lazy_static!` which checks whether the static is initialized
53/// in every read, InitializeOnce has no extra check for reads.
54pub struct InitializeOnce<T: 'static> {
55    v: UnsafeCell<MaybeUninit<T>>,
56    /// This is used to guarantee `init_fn` is only called once.
57    once: Once,
58}
59
60impl<T> InitializeOnce<T> {
61    pub const fn new() -> Self {
62        InitializeOnce {
63            v: UnsafeCell::new(MaybeUninit::uninit()),
64            once: Once::new(),
65        }
66    }
67
68    /// Initialize the value. This should be called before ever using the struct.
69    /// If this method is called by multiple threads, the first thread will
70    /// initialize the value, and the other threads will be blocked until the
71    /// initialization is done (`Once` returns).
72    pub fn initialize_once(&self, init_fn: &'static dyn Fn() -> T) {
73        self.once.call_once(|| {
74            unsafe { &mut *self.v.get() }.write(init_fn());
75        });
76        debug_assert!(self.once.is_completed());
77    }
78
79    /// Get the value. This should only be used after initialize_once()
80    pub fn get_ref(&self) -> &T {
81        // We only assert in debug builds.
82        debug_assert!(self.once.is_completed());
83        unsafe { (*self.v.get()).assume_init_ref() }
84    }
85
86    /// Get a mutable reference to the value.
87    /// This is currently only used for SFTMap during plan creation (single threaded),
88    /// and before the plan creation is done, the binding cannot use MMTK at all.
89    ///
90    /// # Safety
91    /// The caller needs to make sure there is no race when mutating the value.
92    #[allow(clippy::mut_from_ref)]
93    pub unsafe fn get_mut(&self) -> &mut T {
94        // We only assert in debug builds.
95        debug_assert!(self.once.is_completed());
96        unsafe { (*self.v.get()).assume_init_mut() }
97    }
98}
99
100impl<T> std::ops::Deref for InitializeOnce<T> {
101    type Target = T;
102    fn deref(&self) -> &Self::Target {
103        self.get_ref()
104    }
105}
106
107unsafe impl<T> Sync for InitializeOnce<T> {}
108
109/// Create a formatted string that makes the best effort idenfying the current process and thread.
110pub fn debug_process_thread_id() -> String {
111    use crate::util::os::*;
112    format!(
113        "PID: {}, TID: {}",
114        OS::get_process_id().map_or("(Failed to get PID)".to_string(), |pid| format!("{}", pid)),
115        OS::get_thread_id().map_or("(Failed to get TID)".to_string(), |tid| format!(
116            "{:?}",
117            tid
118        )),
119    )
120}
121
122#[cfg(test)]
123mod initialize_once_tests {
124    use super::*;
125
126    #[test]
127    fn test_threads_compete_initialize() {
128        use std::sync::atomic::AtomicUsize;
129        use std::sync::atomic::Ordering;
130        use std::thread;
131
132        // Create multiple threads to initialize the same `InitializeOnce` value
133        const N_THREADS: usize = 1000;
134        // The test value
135        static I: InitializeOnce<usize> = InitializeOnce::new();
136        // Count how many times the function is called
137        static INITIALIZE_COUNT: AtomicUsize = AtomicUsize::new(0);
138        // The function to create initial value
139        fn initialize_usize() -> usize {
140            INITIALIZE_COUNT.fetch_add(1, Ordering::SeqCst);
141            42
142        }
143
144        let mut threads = vec![];
145        for _ in 1..N_THREADS {
146            threads.push(thread::spawn(|| {
147                I.initialize_once(&initialize_usize);
148                // Every thread should see the value correctly initialized.
149                assert_eq!(*I, 42);
150            }));
151        }
152        threads.into_iter().for_each(|t| t.join().unwrap());
153
154        // The initialize_usize should only be called once
155        assert_eq!(INITIALIZE_COUNT.load(Ordering::SeqCst), 1);
156    }
157}