replace QueueStateGuard enum with generic associated lifetime

bonzini · bonzini · commit 57ca35bf8624 · 2021-10-26T11:22:56.000+02:00
Use the same trick that is used in vm-memory for iterators, so that
QueueStateT::lock's returned guard type has a generic associate lifetime.

This makes the QueueStateT trait independent of the QueueState
struct that implements it.

Signed-off-by: Paolo Bonzini &lt;pbonzini@redhat.com&gt;
diff --git a/crates/virtio-queue/src/lib.rs b/crates/virtio-queue/src/lib.rs
@@ -419,39 +419,23 @@ impl VirtqUsedElem {
 
 unsafe impl ByteValued for VirtqUsedElem {}
 
-/// Struct to hold an exclusive reference to the underlying `QueueState` object.
-pub enum QueueStateGuard<'a> {
-    /// A reference to a `QueueState` object.
-    StateObject(&'a mut QueueState),
-    /// A `MutexGuard` for a `QueueState` object.
-    MutexGuard(MutexGuard<'a, QueueState>),
-}
-
-impl<'a> Deref for QueueStateGuard<'a> {
-    type Target = QueueState;
-
-    fn deref(&self) -> &Self::Target {
-        match self {
-            QueueStateGuard::StateObject(v) => v,
-            QueueStateGuard::MutexGuard(v) => v.deref(),
-        }
-    }
-}
-
-impl<'a> DerefMut for QueueStateGuard<'a> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        match self {
-            QueueStateGuard::StateObject(v) => v,
-            QueueStateGuard::MutexGuard(v) => v.deref_mut(),
-        }
-    }
+/// Lifetime-generic guard associated to a QueueStateT. In practice,
+/// instead of having a `QueueStateT::Guard<'a>` generic associated type,
+/// you have to write `<QueueStateT::Guard as QueueStateGuard<'a>>::Out`.
+pub trait QueueStateGuard<'a> {
+    /// Type of the guard returned by the `lock` method.
+    type Out: DerefMut<Target = QueueState>;
 }
 
 /// Trait to access and manipulate a virtio queue.
 ///
 /// To optimize for performance, different implementations of the `QueueStateT` trait may be
 /// provided for single-threaded context and multi-threaded context.
 pub trait QueueStateT {
+    /// Lifetime-generic guard. Usually this is just `Self`, and implementors
+    /// of `QueueStateT` also implement `QueueStateGuard`.
+    type Guard: for<'a> QueueStateGuard<'a>;
+
     /// Construct an empty virtio queue state object with the given `max_size`.
     fn new(max_size: u16) -> Self;
 
@@ -465,7 +449,7 @@ pub trait QueueStateT {
     ///
     /// Logically this method will acquire the underlying lock protecting the `QueueState` Object.
     /// The lock will be released when the returned object gets dropped.
-    fn lock(&mut self) -> QueueStateGuard;
+    fn lock(&mut self) -> <Self::Guard as QueueStateGuard>::Out;
 
     /// Get the maximum size of the virtio queue.
     fn max_size(&self) -> u16;
@@ -664,7 +648,13 @@ impl QueueState {
     }
 }
 
+impl<'a> QueueStateGuard<'a> for QueueState {
+    type Out = &'a mut QueueState;
+}
+
 impl QueueStateT for QueueState {
+    type Guard = Self;
+
     fn new(max_size: u16) -> Self {
         QueueState {
             max_size,
@@ -751,8 +741,8 @@ impl QueueStateT for QueueState {
         self.event_idx_enabled = false;
     }
 
-    fn lock(&mut self) -> QueueStateGuard {
-        QueueStateGuard::StateObject(self)
+    fn lock(&mut self) -> &mut Self {
+        self
     }
 
     fn max_size(&self) -> u16 {
@@ -919,7 +909,13 @@ pub struct QueueStateSync {
     state: Arc<Mutex<QueueState>>,
 }
 
+impl<'a> QueueStateGuard<'a> for QueueStateSync {
+    type Out = MutexGuard<'a, QueueState>;
+}
+
 impl QueueStateT for QueueStateSync {
+    type Guard = Self;
+
     fn new(max_size: u16) -> Self {
         QueueStateSync {
             state: Arc::new(Mutex::new(QueueState::new(max_size))),
@@ -934,8 +930,8 @@ impl QueueStateT for QueueStateSync {
         self.state.lock().unwrap().reset();
     }
 
-    fn lock(&mut self) -> QueueStateGuard {
-        QueueStateGuard::MutexGuard(self.state.lock().unwrap())
+    fn lock(&mut self) -> MutexGuard<QueueState> {
+        self.state.lock().unwrap()
     }
 
     fn max_size(&self) -> u16 {
@@ -1040,7 +1036,7 @@ impl<M: GuestAddressSpace, S: QueueStateT> Queue<M, S> {
     ///
     /// Logically this method will acquire the underlying lock protecting the `QueueState` Object.
     /// The lock will be released when the returned object gets dropped.
-    pub fn lock(&mut self) -> QueueStateGuard {
+    pub fn lock(&mut self) -> <S::Guard as QueueStateGuard>::Out {
         self.state.lock()
     }
 
