lazy_init.rs 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156
  1. // Copyright (C) DragonOS Community longjin
  2. // This program is free software; you can redistribute it and/or
  3. // modify it under the terms of the GNU General Public License
  4. // as published by the Free Software Foundation; either version 2
  5. // of the License, or (at your option) any later version.
  6. // This program is distributed in the hope that it will be useful,
  7. // but WITHOUT ANY WARRANTY; without even the implied warranty of
  8. // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  9. // GNU General Public License for more details.
  10. // You should have received a copy of the GNU General Public License
  11. // along with this program; if not, write to the Free Software
  12. // Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  13. // Or you can visit https://www.gnu.org/licenses/gpl-2.0.html
  14. #![allow(dead_code)]
  15. use std::cell::UnsafeCell;
  16. use std::fmt::Debug;
  17. use std::mem::MaybeUninit;
  18. use std::ops::{Deref, DerefMut};
  19. use std::sync::atomic::{AtomicBool, Ordering};
  20. use std::sync::Mutex;
  21. /// A wrapper around a value that is initialized lazily.
  22. pub struct Lazy<T> {
  23. /// The lock that is used to ensure that only one thread calls the init function at the same time.
  24. init_lock: Mutex<()>,
  25. /// The value that is initialized lazily.
  26. value: UnsafeCell<MaybeUninit<T>>,
  27. /// Whether the value has been initialized.
  28. initialized: AtomicBool,
  29. }
  30. impl<T> Lazy<T> {
  31. /// Creates a new `Lazy` value that will be initialized with the
  32. /// result of the closure `init`.
  33. pub const fn new() -> Lazy<T> {
  34. Lazy {
  35. value: UnsafeCell::new(MaybeUninit::uninit()),
  36. initialized: AtomicBool::new(false),
  37. init_lock: Mutex::new(()),
  38. }
  39. }
  40. /// Returns true if the value has been initialized.
  41. #[inline(always)]
  42. pub fn initialized(&self) -> bool {
  43. let initialized = self.initialized.load(Ordering::Acquire);
  44. if initialized {
  45. return true;
  46. }
  47. return false;
  48. }
  49. /// Ensures that this lazy value is initialized. If the value has not
  50. /// yet been initialized, this will raise a panic.
  51. #[inline(always)]
  52. pub fn ensure(&self) {
  53. if self.initialized() {
  54. return;
  55. }
  56. panic!("Lazy value was not initialized");
  57. }
  58. pub fn init(&self, value: T) {
  59. assert!(!self.initialized());
  60. // We need this lock to ensure that only one thread calls the init function at the same time.
  61. let _init_guard = self.init_lock.lock();
  62. // Check again, in case another thread initialized it while we were waiting for the lock.
  63. let initialized = self.initialized();
  64. if initialized {
  65. return;
  66. }
  67. unsafe {
  68. (*self.value.get()).as_mut_ptr().write(value);
  69. }
  70. self.initialized.store(true, Ordering::Release);
  71. }
  72. /// Forces the evaluation of this lazy value and returns a reference to
  73. /// the result. This is equivalent to the `Deref` impl, but is explicit.
  74. /// This will initialize the value if it has not yet been initialized.
  75. pub fn get(&self) -> &T {
  76. self.ensure();
  77. return unsafe { self.get_unchecked() };
  78. }
  79. /// Returns a reference to the value if it has been initialized.
  80. /// Otherwise, returns `None`.
  81. pub fn try_get(&self) -> Option<&T> {
  82. if self.initialized() {
  83. return Some(unsafe { self.get_unchecked() });
  84. }
  85. return None;
  86. }
  87. /// Forces the evaluation of this lazy value and returns a mutable
  88. /// reference to the result. This is equivalent to the `DerefMut` impl,
  89. /// but is explicit. This will initialize the value if it has not yet
  90. /// been initialized.
  91. pub fn get_mut(&mut self) -> &mut T {
  92. self.ensure();
  93. return unsafe { self.get_mut_unchecked() };
  94. }
  95. #[inline(always)]
  96. pub unsafe fn get_unchecked(&self) -> &T {
  97. return &*(*self.value.get()).as_ptr();
  98. }
  99. #[inline(always)]
  100. pub unsafe fn get_mut_unchecked(&mut self) -> &mut T {
  101. return &mut *(*self.value.get()).as_mut_ptr();
  102. }
  103. }
  104. impl<T> Deref for Lazy<T> {
  105. type Target = T;
  106. #[inline(always)]
  107. fn deref(&self) -> &T {
  108. return self.get();
  109. }
  110. }
  111. impl<T> DerefMut for Lazy<T> {
  112. #[inline(always)]
  113. fn deref_mut(&mut self) -> &mut T {
  114. return self.get_mut();
  115. }
  116. }
  117. impl<T: Debug> Debug for Lazy<T> {
  118. fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
  119. if let Some(value) = self.try_get() {
  120. return write!(f, "Lazy({:?})", value);
  121. } else {
  122. return write!(f, "Lazy(uninitialized)");
  123. }
  124. }
  125. }
  126. impl<T> Drop for Lazy<T> {
  127. fn drop(&mut self) {
  128. if self.initialized() {
  129. unsafe {
  130. (*self.value.get()).as_mut_ptr().drop_in_place();
  131. }
  132. }
  133. }
  134. }
  135. unsafe impl<T: Send + Sync> Sync for Lazy<T> {}
  136. unsafe impl<T: Send> Send for Lazy<T> {}