make eval copy

packages/liveview/src/query.rs

@@ -1,5 +1,7 @@
 use std::{cell::RefCell, rc::Rc};
 
+use futures_util::FutureExt;
+use generational_box::{Owner, UnsyncStorage};
 use serde::{de::DeserializeOwned, Deserialize};
 use serde_json::Value;
 use slab::Slab;
@@ -42,8 +44,8 @@ let dioxus = {
 
 /// Tracks what query ids are currently active
 
-struct SharedSlab<T = ()> {
-    slab: Rc<RefCell<Slab<T>>>,
+pub(crate) struct SharedSlab<T = ()> {
+    pub(crate) slab: Rc<RefCell<Slab<T>>>,
 }
 
 impl<T> Clone for SharedSlab<T> {
@@ -62,9 +64,10 @@ impl<T> Default for SharedSlab<T> {
     }
 }
 
-struct QueryEntry {
+pub(crate) struct QueryEntry {
     channel_sender: tokio::sync::mpsc::UnboundedSender<Value>,
     return_sender: Option<tokio::sync::oneshot::Sender<Value>>,
+    pub(crate) owner: Option<Owner<UnsyncStorage>>,
 }
 
 const QUEUE_NAME: &str = "__msg_queues";
@@ -72,7 +75,7 @@ const QUEUE_NAME: &str = "__msg_queues";
 /// Handles sending and receiving arbitrary queries from the webview. Queries can be resolved non-sequentially, so we use ids to track them.
 #[derive(Clone)]
 pub(crate) struct QueryEngine {
-    active_requests: SharedSlab<QueryEntry>,
+    pub(crate) active_requests: SharedSlab<QueryEntry>,
     query_tx: tokio::sync::mpsc::UnboundedSender<String>,
 }
 
@@ -91,6 +94,7 @@ impl QueryEngine {
         let request_id = self.active_requests.slab.borrow_mut().insert(QueryEntry {
             channel_sender: tx,
             return_sender: Some(return_tx),
+            owner: None,
         });
 
         // start the query
@@ -162,7 +166,7 @@ pub(crate) struct Query<V: DeserializeOwned> {
     query_engine: QueryEngine,
     pub receiver: tokio::sync::mpsc::UnboundedReceiver<Value>,
     pub return_receiver: Option<tokio::sync::oneshot::Receiver<Value>>,
-    id: usize,
+    pub id: usize,
     phantom: std::marker::PhantomData<V>,
 }
 
@@ -198,12 +202,14 @@ impl<V: DeserializeOwned> Query<V> {
         Ok(())
     }
 
-    /// Receive a message from the query
-    pub async fn recv(&mut self) -> Result<Value, QueryError> {
+    /// Poll the query for a message
+    pub fn poll_recv(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<Value, QueryError>> {
         self.receiver
-            .recv()
-            .await
-            .ok_or(QueryError::Recv(RecvError::Closed))
+            .poll_recv(cx)
+            .map(|result| result.ok_or(QueryError::Recv(RecvError::Closed)))
     }
 
     /// Receive the result of the query
@@ -215,6 +221,19 @@ impl<V: DeserializeOwned> Query<V> {
             None => Err(QueryError::Finished),
         }
     }
+
+    /// Poll the query for a result
+    pub fn poll_result(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<Value, QueryError>> {
+        match self.return_receiver.as_mut() {
+            Some(receiver) => receiver
+                .poll_unpin(cx)
+                .map_err(|_| QueryError::Recv(RecvError::Closed)),
+            None => std::task::Poll::Ready(Err(QueryError::Finished)),
+        }
+    }
 }
 
 impl<V: DeserializeOwned> Drop for Query<V> {