Implement DerefMut for Lazy

matklad · Jan 14, 2020 · fa5cf7c · fa5cf7c
1 parent 3353532
commit fa5cf7c
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 3 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -237,7 +237,7 @@ pub mod unsync {
     use core::{
         cell::{Cell, UnsafeCell},
         fmt,
-        ops::Deref,
+        ops::{Deref, DerefMut},
     };
 
     #[cfg(feature = "std")]
@@ -550,6 +550,13 @@ pub mod unsync {
         }
     }
 
+    impl<T, F: FnOnce() -> T> DerefMut for Lazy<T, F> {
+        fn deref_mut(&mut self) -> &mut T {
+            Lazy::force(self);
+            self.cell.get_mut().unwrap_or_else(|| unreachable!())
+        }
+    }
+
     impl<T: Default> Default for Lazy<T> {
         /// Creates a new lazy value using `Default` as the initializing function.
         fn default() -> Lazy<T> {
@@ -560,7 +567,13 @@ pub mod unsync {
 
 #[cfg(feature = "std")]
 pub mod sync {
-    use std::{cell::Cell, fmt, hint::unreachable_unchecked, panic::RefUnwindSafe};
+    use std::{
+        cell::Cell,
+        fmt,
+        hint::unreachable_unchecked,
+        ops::{Deref, DerefMut},
+        panic::RefUnwindSafe,
+    };
 
     use crate::imp::OnceCell as Imp;
 
@@ -900,13 +913,20 @@ pub mod sync {
         }
     }
 
-    impl<T, F: FnOnce() -> T> ::std::ops::Deref for Lazy<T, F> {
+    impl<T, F: FnOnce() -> T> Deref for Lazy<T, F> {
         type Target = T;
         fn deref(&self) -> &T {
             Lazy::force(self)
         }
     }
 
+    impl<T, F: FnOnce() -> T> DerefMut for Lazy<T, F> {
+        fn deref_mut(&mut self) -> &mut T {
+            Lazy::force(self);
+            self.cell.get_mut().unwrap_or_else(|| unreachable!())
+        }
+    }
+
     impl<T: Default> Default for Lazy<T> {
         /// Creates a new lazy value using `Default` as the initializing function.
         fn default() -> Lazy<T> {

diff --git a/tests/test.rs b/tests/test.rs
@@ -111,6 +111,25 @@ mod unsync {
         assert_eq!(called.get(), 1);
     }
 
+    #[test]
+    fn lazy_deref_mut() {
+        let called = Cell::new(0);
+        let mut x = Lazy::new(|| {
+            called.set(called.get() + 1);
+            92
+        });
+
+        assert_eq!(called.get(), 0);
+
+        let y = *x - 30;
+        assert_eq!(y, 62);
+        assert_eq!(called.get(), 1);
+
+        *x /= 2;
+        assert_eq!(*x, 46);
+        assert_eq!(called.get(), 1);
+    }
+
     #[test]
     fn lazy_default() {
         static CALLED: AtomicUsize = AtomicUsize::new(0);
@@ -384,6 +403,25 @@ mod sync {
         assert_eq!(called.load(SeqCst), 1);
     }
 
+    #[test]
+    fn lazy_deref_mut() {
+        let called = AtomicUsize::new(0);
+        let mut x = Lazy::new(|| {
+            called.fetch_add(1, SeqCst);
+            92
+        });
+
+        assert_eq!(called.load(SeqCst), 0);
+
+        let y = *x - 30;
+        assert_eq!(y, 62);
+        assert_eq!(called.load(SeqCst), 1);
+
+        *x /= 2;
+        assert_eq!(*x, 46);
+        assert_eq!(called.load(SeqCst), 1);
+    }
+
     #[test]
     fn lazy_default() {
         static CALLED: AtomicUsize = AtomicUsize::new(0);