From 876233b9ba6ca8ed76301b7da7199bc1568fe25c Mon Sep 17 00:00:00 2001
From: Wedson Almeida Filho <wedsonaf@google.com>
Date: Fri, 15 Oct 2021 17:36:45 +0100
Subject: [PATCH] rust: improve `ScopeGuard`.

This allows us to use a scope guard with mutable data by wrapping it
(the mutable data). Capturing variables mutably in the closure would
render them inaccessbile in the function body.

Signed-off-by: Wedson Almeida Filho <wedsonaf@google.com>
---
 rust/kernel/types.rs | 83 ++++++++++++++++++++++++++++++++++++--------
 1 file changed, 68 insertions(+), 15 deletions(-)

diff --git a/rust/kernel/types.rs b/rust/kernel/types.rs
index 63e6e7a639a032..5916dff3a45a3b 100644
--- a/rust/kernel/types.rs
+++ b/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)
         }
     }
 }