Merge pull request torvalds#541 from wedsonaf/scope-guard

rust: improve `ScopeGuard`.

Author: ojeda

rust/kernel/types.rs

@@ -9,7 +9,13 @@ use crate::{
     sync::{Ref, RefBorrow},
 };
 use alloc::boxed::Box;
-use core::{cell::UnsafeCell, mem::MaybeUninit, ops, ops::Deref, pin::Pin, ptr::NonNull};
+use core::{
+    cell::UnsafeCell,
+    mem::MaybeUninit,
+    ops::{self, Deref, DerefMut},
+    pin::Pin,
+    ptr::NonNull,
+};
 
 /// Permissions.
 ///
@@ -200,29 +206,76 @@ impl<T: PointerWrapper + Deref> PointerWrapper for Pin<T> {
 ///     pr_info!("example2 no early return\n");
 /// }
 /// ```
-pub struct ScopeGuard<T: FnOnce()> {
-    cleanup_func: Option<T>,
+///
+/// In the example below, we need a mutable object (the vector) to be accessible within the log
+/// function, so we wrap it in the [`ScopeGuard`]:
+/// ```
+/// # use kernel::prelude::*;
+/// # use kernel::ScopeGuard;
+/// fn example3(arg: bool) -> Result {
+///     let mut vec =
+///         ScopeGuard::new_with_data(Vec::new(), |v| pr_info!("vec had {} elements\n", v.len()));
+///
+///     vec.try_push(10u8)?;
+///     if arg {
+///         return Ok(());
+///     }
+///     vec.try_push(20u8)?;
+///     Ok(())
+/// }
+/// ```
+///
+/// # Invariants
+///
+/// The value stored in the struct is nearly always `Some(_)`, except between
+/// [`ScopeGuard::dismiss`] and [`ScopeGuard::drop`]: in this case, it will be `None` as the value
+/// will have been returned to the caller. Since  [`ScopeGuard::dismiss`] consumes the guard,
+/// callers won't be able to use it anymore.
+pub struct ScopeGuard<T, F: FnOnce(T)>(Option<(T, F)>);
+
+impl<T, F: FnOnce(T)> ScopeGuard<T, F> {
+    /// Creates a new guarded object wrapping the given data and with the given cleanup function.
+    pub fn new_with_data(data: T, cleanup_func: F) -> Self {
+        // INVARIANT: The struct is being initialised with `Some(_)`.
+        Self(Some((data, cleanup_func)))
+    }
+
+    /// Prevents the cleanup function from running and returns the guarded data.
+    pub fn dismiss(mut self) -> T {
+        // INVARIANT: This is the exception case in the invariant; it is not visible to callers
+        // because this function consumes `self`.
+        self.0.take().unwrap().0
+    }
 }
 
-impl<T: FnOnce()> ScopeGuard<T> {
-    /// Creates a new cleanup object with the given cleanup function.
-    pub fn new(cleanup_func: T) -> Self {
-        Self {
-            cleanup_func: Some(cleanup_func),
-        }
+impl ScopeGuard<(), Box<dyn FnOnce(())>> {
+    /// Creates a new guarded object with the given cleanup function.
+    pub fn new(cleanup: impl FnOnce()) -> ScopeGuard<(), impl FnOnce(())> {
+        ScopeGuard::new_with_data((), move |_| cleanup())
     }
+}
+
+impl<T, F: FnOnce(T)> Deref for ScopeGuard<T, F> {
+    type Target = T;
+
+    fn deref(&self) -> &T {
+        // The type invariants guarantee that `unwrap` will succeed.
+        &self.0.as_ref().unwrap().0
+    }
+}
 
-    /// Prevents the cleanup function from running.
-    pub fn dismiss(mut self) {
-        self.cleanup_func.take();
+impl<T, F: FnOnce(T)> DerefMut for ScopeGuard<T, F> {
+    fn deref_mut(&mut self) -> &mut T {
+        // The type invariants guarantee that `unwrap` will succeed.
+        &mut self.0.as_mut().unwrap().0
     }
 }
 
-impl<T: FnOnce()> Drop for ScopeGuard<T> {
+impl<T, F: FnOnce(T)> Drop for ScopeGuard<T, F> {
     fn drop(&mut self) {
         // Run the cleanup function if one is still present.
-        if let Some(cleanup) = self.cleanup_func.take() {
-            cleanup();
+        if let Some((data, cleanup)) = self.0.take() {
+            cleanup(data)
         }
     }
 }