[PATCH yew-widget-toolkit] state: (tree)store: fix mutable borrow during listener notification

[Dominik Csapak <d.csapak@proxmox.com>]

In the (tree)store, when a write guard was dropped, the listener were
notified while there was still a mutable borrow on it. This would lead
to a panic when trying to access the tree or store inside the listener.

To avoid that, use a similar pattern as we have e.g. for SharedState or
Selection, and use ManuallyDrop. Then we can drop the borrow before
notifying the listeners.

Signed-off-by: Dominik Csapak <d.csapak@proxmox.com>
---
inspired by the recent changes to the loader/asyncabortguard by thomas.

 src/state/store.rs                      | 49 +++++++++++++++++++++----
 src/state/tree_store/keyed_slab_tree.rs |  7 ++--
 src/state/tree_store/mod.rs             | 43 ++++++++++++++++++++--
 3 files changed, 83 insertions(+), 16 deletions(-)

diff --git a/src/state/store.rs b/src/state/store.rs
index ec34f50..1afe0a4 100644
--- a/src/state/store.rs
+++ b/src/state/store.rs
@@ -1,4 +1,5 @@
 use std::cell::{Ref, RefCell, RefMut};
+use std::mem::ManuallyDrop;
 use std::ops::Range;
 use std::ops::{Deref, DerefMut};
 use std::rc::Rc;
@@ -185,7 +186,7 @@ impl<T: 'static> Store<T> {
     /// being in place.
     pub fn try_write(&self) -> Result<StoreWriteGuard<'_, T>, Error> {
         Ok(StoreWriteGuard {
-            state: self.inner.try_borrow_mut()?,
+            state: ManuallyDrop::new(self.inner.try_borrow_mut()?),
             update: true,
         })
     }
@@ -277,7 +278,7 @@ impl<T: 'static> Store<T> {
 ///
 /// Notifies store listeners when dropped, except when using [StoreWriteGuard::skip_update]
 pub struct StoreWriteGuard<'a, T: 'static> {
-    state: RefMut<'a, StoreState<T>>,
+    state: ManuallyDrop<RefMut<'a, StoreState<T>>>,
     update: bool,
 }
 
@@ -306,9 +307,22 @@ impl<T> DerefMut for StoreWriteGuard<'_, T> {
 
 impl<T: 'static> Drop for StoreWriteGuard<'_, T> {
     fn drop(&mut self) {
-        if self.update {
+        let listeners = if self.update {
             self.version += 1;
-            self.state.notify_listeners();
+            Some(self.state.clone_listeners())
+        } else {
+            None
+        };
+
+        // SAFETY: Only called once (here), and the value is not used afterwards.
+        unsafe {
+            ManuallyDrop::drop(&mut self.state);
+        }
+
+        if let Some(listeners) = listeners {
+            for (_key, listener) in listeners.iter() {
+                listener.emit(());
+            }
         }
     }
 }
@@ -495,10 +509,9 @@ impl<T: 'static> StoreState<T> {
         self.listeners.remove(key);
     }
 
-    pub(crate) fn notify_listeners(&self) {
-        for (_key, listener) in self.listeners.iter() {
-            listener.emit(());
-        }
+    /// Snapshot of the listeners, so they can be notified after releasing the borrow.
+    pub(crate) fn clone_listeners(&self) -> Slab<Callback<()>> {
+        self.listeners.clone()
     }
 
     fn set_sorter(&mut self, sorter: impl IntoSorterFn<T>) {
@@ -643,3 +656,23 @@ where
         Some((pos, node))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::cell::Cell;
+
+    use super::*;
+
+    #[test]
+    fn listener_can_read_store_during_notification() {
+        let store = Store::with_extract_key(ExtractKeyFn::new(|v: &u32| Key::from(*v as u64)));
+        let seen = Rc::new(Cell::new(0));
+        let _observer = store.add_listener({
+            let store = store.clone();
+            let seen = Rc::clone(&seen);
+            move |()| seen.set(store.read().data().len())
+        });
+        store.set_data(vec![1, 2, 3]);
+        assert_eq!(seen.get(), 3);
+    }
+}
diff --git a/src/state/tree_store/keyed_slab_tree.rs b/src/state/tree_store/keyed_slab_tree.rs
index 1ff0d84..11e93d8 100644
--- a/src/state/tree_store/keyed_slab_tree.rs
+++ b/src/state/tree_store/keyed_slab_tree.rs
@@ -279,10 +279,9 @@ impl<T> KeyedSlabTree<T> {
         self.listeners.remove(key);
     }
 
-    pub(crate) fn notify_listeners(&self) {
-        for (_key, listener) in self.listeners.iter() {
-            listener.emit(());
-        }
+    /// Snapshot of the listeners, so they can be notified after releasing the borrow.
+    pub(crate) fn clone_listeners(&self) -> Slab<Callback<()>> {
+        self.listeners.clone()
     }
 
     pub(crate) fn set_sorter(&mut self, sorter: impl IntoSorterFn<T>) {
diff --git a/src/state/tree_store/mod.rs b/src/state/tree_store/mod.rs
index e45e1f9..5e84b8f 100644
--- a/src/state/tree_store/mod.rs
+++ b/src/state/tree_store/mod.rs
@@ -14,6 +14,7 @@ pub use keyed_slab_tree::{
 mod slab_tree_serde;
 
 use std::cell::{Ref, RefCell, RefMut};
+use std::mem::ManuallyDrop;
 use std::ops::{Deref, DerefMut, Range};
 use std::rc::Rc;
 
@@ -143,7 +144,7 @@ impl<T: 'static> TreeStore<T> {
     ///
     /// Panics if the store is already locked.
     pub fn write(&self) -> TreeStoreWriteGuard<'_, T> {
-        let tree = self.inner.borrow_mut();
+        let tree = ManuallyDrop::new(self.inner.borrow_mut());
 
         TreeStoreWriteGuard {
             initial_version: tree.version(),
@@ -316,7 +317,7 @@ impl<T> Deref for TreeStoreReadGuard<'_, T> {
 
 /// A wrapper type for a mutably borrowed [TreeStore]
 pub struct TreeStoreWriteGuard<'a, T> {
-    tree: RefMut<'a, KeyedSlabTree<T>>,
+    tree: ManuallyDrop<RefMut<'a, KeyedSlabTree<T>>>,
     initial_version: usize,
 }
 
@@ -336,8 +337,22 @@ impl<T> DerefMut for TreeStoreWriteGuard<'_, T> {
 
 impl<T> Drop for TreeStoreWriteGuard<'_, T> {
     fn drop(&mut self) {
-        if self.tree.version() != self.initial_version {
-            self.tree.notify_listeners();
+        let changed = self.tree.version() != self.initial_version;
+        let listeners = if changed {
+            Some(self.tree.clone_listeners())
+        } else {
+            None
+        };
+
+        // SAFETY: Only called once (here) and value is not used afterwards.
+        unsafe {
+            ManuallyDrop::drop(&mut self.tree);
+        }
+
+        if let Some(listeners) = listeners {
+            for (_key, listener) in listeners.iter() {
+                listener.emit(());
+            }
         }
     }
 }
@@ -462,3 +477,23 @@ where
         Some((pos, node))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::cell::Cell;
+
+    use super::*;
+
+    #[test]
+    fn listener_can_read_store_during_notification() {
+        let store = TreeStore::with_extract_key(ExtractKeyFn::new(|v: &u32| Key::from(*v as u64)));
+        let seen = Rc::new(Cell::new(false));
+        let _observer = store.add_listener({
+            let store = store.clone();
+            let seen = Rc::clone(&seen);
+            move |()| seen.set(store.read().root().is_some())
+        });
+        store.write().set_root(1);
+        assert!(seen.get());
+    }
+}
-- 
2.47.3