refactor(core): 💡 QueryHandle store dyn Any

core/src/query.rs

@@ -3,7 +3,7 @@ use std::any::{Any, TypeId};
 use smallvec::SmallVec;
 
 use crate::state::{
-  MapReader, MapWriterAsReader, ReadRef, Reader, StateReader, StateWriter, WriteRef,
+  MapReader, MapWriterAsReader, PartData, ReadRef, Reader, StateReader, StateWriter, WriteRef,
 };
 
 /// A type can composed by many types, this trait help us to query the type and
@@ -35,20 +35,24 @@ pub struct Queryable<T: Any>(pub T);
 pub struct QueryHandle<'a>(InnerHandle<'a>);
 
 /// A reference to a query result of a data, it's similar to `&T`.
-pub struct QueryRef<'a, T> {
+pub struct QueryRef<'a, T: ?Sized> {
   pub(crate) type_ref: &'a T,
-  pub(crate) _data: Option<Box<dyn QueryResult + 'a>>,
+  pub(crate) _data: Option<Box<dyn Any>>,
 }
 
 impl<'a> QueryHandle<'a> {
-  pub fn new(r: &'a dyn Any) -> Self { QueryHandle(InnerHandle::Ref(r)) }
-
   /// Downcast the to type `T` and returns a reference to it,
   /// return `None` if the type not match
   pub fn downcast_ref<T: Any>(&self) -> Option<&T> {
     match self.0 {
       InnerHandle::Ref(r) => r.downcast_ref::<T>(),
-      InnerHandle::Owned(ref o) => query_downcast_ref(o.as_ref()),
+      InnerHandle::Owned(ref o) => o
+        .downcast_ref::<ReadRef<'static, dyn Any>>()
+        .and_then(|r| r.downcast_ref::<T>())
+        .or_else(|| {
+          o.downcast_ref::<WriteRef<'static, dyn Any>>()
+            .and_then(|w| w.downcast_ref::<T>())
+        }),
     }
   }
 
@@ -63,23 +67,42 @@ impl<'a> QueryHandle<'a> {
     let InnerHandle::Owned(ref mut o) = self.0 else {
       return None;
     };
-    (o.query_type() == TypeId::of::<WriteRef<'static, T>>()).then(|| {
-      // SAFETY: the creator guarantees that the query type is `WriteRef<T>`,
-      unsafe { &mut **(o.as_mut() as *mut dyn QueryResult as *mut WriteRef<'a, T>) }
-    })
+
+    o.downcast_mut::<WriteRef<'static, dyn Any>>()?
+      .downcast_mut::<T>()
   }
 
-  pub(crate) fn owned(o: Box<dyn QueryResult + 'a>) -> Self { QueryHandle(InnerHandle::Owned(o)) }
+  pub(crate) fn new(r: &'a dyn Any) -> Self { QueryHandle(InnerHandle::Ref(r)) }
+
+  pub(crate) fn from_read_ref(r: ReadRef<'a, dyn Any>) -> Self {
+    // Safety: The lifetime is maintained in the return handle and will be shortened
+    // once the handle is downcast.
+    let r: ReadRef<'static, dyn Any> = unsafe { std::mem::transmute(r) };
+    QueryHandle(InnerHandle::Owned(Box::new(r)))
+  }
+
+  pub(crate) fn from_write_ref(w: WriteRef<'a, dyn Any>) -> Self {
+    // Safety: The lifetime is maintained in the return handle and will be shortened
+    // once the handle is downcast.
+    let w: WriteRef<'static, dyn Any> = unsafe { std::mem::transmute(w) };
+    QueryHandle(InnerHandle::Owned(Box::new(w)))
+  }
 
   pub fn into_ref<T: Any>(self) -> Option<QueryRef<'a, T>> {
     match self.0 {
       InnerHandle::Ref(r) if r.type_id() == TypeId::of::<T>() => {
         Some(QueryRef { type_ref: r.downcast_ref::<T>().unwrap(), _data: None })
       }
       InnerHandle::Owned(o) => {
-        let r = query_downcast_ref(o.as_ref());
-        // hold the _data to keep the data alive
-        r.map(|r| QueryRef { type_ref: r, _data: Some(o) })
+        let inner = o
+          .downcast_ref::<ReadRef<'static, dyn Any>>()
+          .and_then(|r| r.downcast_ref::<T>())
+          .or_else(|| {
+            o.downcast_ref::<WriteRef<'static, dyn Any>>()
+              .and_then(|w| w.downcast_ref::<T>())
+          })?;
+        let type_ref = unsafe { &*(inner as *const T) };
+        Some(QueryRef { type_ref, _data: Some(o) })
       }
       _ => None,
     }
@@ -89,57 +112,18 @@ impl<'a> QueryHandle<'a> {
     let InnerHandle::Owned(o) = self.0 else {
       return None;
     };
-    (o.query_type() == TypeId::of::<WriteRef<'static, T>>()).then(|| {
-      // SAFETY: the creator guarantees that the query type is `WriteRef<T>`,
-      let w_r = unsafe {
-        let ptr = Box::into_raw(o);
-        Box::from_raw(ptr as *mut WriteRef<'a, T>)
-      };
-      *w_r
-    })
-  }
-}
 
-fn query_downcast_ref<'a, T: Any>(q: &(dyn QueryResult + 'a)) -> Option<&'a T> {
-  let q_type = q.query_type();
-  if q_type == TypeId::of::<T>() {
-    // SAFETY: the creator guarantees that the query type is `T`,
-    let t = unsafe { &*(q as *const dyn QueryResult as *const T) };
-    Some(t)
-  } else if q_type == TypeId::of::<WriteRef<'static, T>>() {
-    // SAFETY: the creator guarantees that the query type is `WriteRef<T>`,
-    let wr = unsafe { &*(q as *const dyn QueryResult as *const WriteRef<'a, T>) };
-    Some(wr)
-  } else if q_type == TypeId::of::<ReadRef<'static, T>>() {
-    // SAFETY: the creator guarantees that the query type is `WriteRef<T>`,
-    let rr = unsafe { &*(q as *const dyn QueryResult as *const ReadRef<'a, T>) };
-    Some(rr)
-  } else {
-    None
+    let w = *o.downcast::<WriteRef<'static, dyn Any>>().ok()?;
+    WriteRef::filter_map(w, |v| v.downcast_mut::<T>().map(PartData::from_ref_mut)).ok()
   }
 }
+
 enum InnerHandle<'a> {
   Ref(&'a dyn Any),
-  Owned(Box<dyn QueryResult + 'a>),
+  Owned(Box<dyn Any>),
 }
 
-pub(crate) trait QueryResult {
-  fn query_type(&self) -> TypeId;
-}
-
-impl<'a> QueryResult for &'a dyn Any {
-  fn query_type(&self) -> TypeId { Any::type_id(*self) }
-}
-
-impl<'a, T: Any> QueryResult for WriteRef<'a, T> {
-  fn query_type(&self) -> TypeId { TypeId::of::<WriteRef<'static, T>>() }
-}
-
-impl<'a, T: Any> QueryResult for ReadRef<'a, T> {
-  fn query_type(&self) -> TypeId { TypeId::of::<ReadRef<'static, T>>() }
-}
-
-impl<'a, T> std::ops::Deref for QueryRef<'a, T> {
+impl<'a, T: ?Sized> std::ops::Deref for QueryRef<'a, T> {
   type Target = T;
 
   fn deref(&self) -> &Self::Target { self.type_ref }
@@ -169,7 +153,8 @@ where
 
   fn query(&self, type_id: TypeId) -> Option<QueryHandle> {
     if type_id == TypeId::of::<T::Value>() {
-      Some(QueryHandle::owned(Box::new(self.write())))
+      let w = WriteRef::map(self.write(), |v| PartData::from_ref(v as &dyn Any));
+      Some(QueryHandle::from_write_ref(w))
     } else if type_id == self.type_id() {
       Some(QueryHandle::new(self))
     } else {
@@ -190,7 +175,8 @@ macro_rules! impl_query_for_reader {
 
     fn query(&self, type_id: TypeId) -> Option<QueryHandle> {
       if type_id == TypeId::of::<V>() {
-        Some(QueryHandle::owned(Box::new(self.read())))
+        let r = ReadRef::map(self.read(), |v| PartData::from_ref(v as &dyn Any));
+        Some(QueryHandle::from_read_ref(r))
       } else if type_id == self.type_id() {
         Some(QueryHandle::new(self))
       } else {
@@ -225,10 +211,7 @@ impl<V: Any> Query for Reader<V> {
 #[cfg(test)]
 mod tests {
   use super::*;
-  use crate::{
-    reset_test_env,
-    state::{PartData, State},
-  };
+  use crate::{reset_test_env, state::State};
 
   #[test]
   fn query_ref() {

core/src/state.rs

@@ -264,6 +264,43 @@ impl<'a, V: ?Sized> WriteRef<'a, V> {
     WriteRef { value, modified: false, modify_scope: orig.modify_scope, info: orig.info }
   }
 
+  /// Makes a new `WriteRef` for an optional component of the borrowed data. The
+  /// original guard is returned as an `Err(..)` if the closure returns
+  /// `None`.
+  ///
+  /// This is an associated function that needs to be used as
+  /// `WriteRef::filter_map(...)`. A method would interfere with methods of the
+  /// same name on `T` used through `Deref`.
+  ///
+  /// # Examples
+  ///
+  /// ```
+  /// use ribir_core::prelude::*;
+  ///
+  /// let c = Stateful::new(vec![1, 2, 3]);
+  /// let b1: WriteRef<'_, Vec<u32>> = c.write();
+  /// let b2: Result<WriteRef<'_, u32>, _> =
+  ///   WriteRef::filter_map(b1, |v| v.get(1).map(PartData::from_ref));
+  /// assert_eq!(*b2.unwrap(), 2);
+  /// ```
+  pub fn filter_map<U: ?Sized, M>(
+    mut orig: WriteRef<'a, V>, part_map: M,
+  ) -> Result<WriteRef<'a, U>, Self>
+  where
+    M: Fn(&mut V) -> Option<PartData<U>>,
+  {
+    match part_map(&mut orig.value) {
+      Some(inner) => {
+        let borrow = orig.value.borrow.clone();
+        let value = ValueMutRef { inner, borrow };
+        let WriteRef { modify_scope, info, .. } = orig;
+
+        Ok(WriteRef { value, modified: false, modify_scope, info })
+      }
+      None => Err(orig),
+    }
+  }
+
   pub fn map_split<U1: ?Sized, U2: ?Sized, F>(
     mut orig: WriteRef<'a, V>, f: F,
   ) -> (WriteRef<'a, U1>, WriteRef<'a, U2>)

core/src/state/state_cell.rs

@@ -257,3 +257,19 @@ impl<'a, T: ?Sized> DerefMut for ValueMutRef<'a, T> {
   #[inline]
   fn deref_mut(&mut self) -> &mut Self::Target { &mut self.inner }
 }
+
+use std::fmt::*;
+
+use super::{QueryRef, WriteRef};
+
+impl<T: ?Sized + Debug> Debug for ReadRef<'_, T> {
+  fn fmt(&self, f: &mut Formatter<'_>) -> Result { Debug::fmt(&**self, f) }
+}
+
+impl<T: ?Sized + Debug> Debug for WriteRef<'_, T> {
+  fn fmt(&self, f: &mut Formatter<'_>) -> Result { Debug::fmt(&**self, f) }
+}
+
+impl<T: ?Sized + Debug> Debug for QueryRef<'_, T> {
+  fn fmt(&self, f: &mut Formatter<'_>) -> Result { Debug::fmt(&**self, f) }
+}