From c954a116eb87a41ff5a5cc935afe798bc433a10a Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Wed, 19 Jul 2023 14:23:34 -0700
Subject: [PATCH] Fix race condition in use_future

---
 packages/hooks/src/usefuture.rs | 115 ++++++--------------------------
 1 file changed, 21 insertions(+), 94 deletions(-)

diff --git a/packages/hooks/src/usefuture.rs b/packages/hooks/src/usefuture.rs
index 18bf50310..4d90adbd4 100644
--- a/packages/hooks/src/usefuture.rs
+++ b/packages/hooks/src/usefuture.rs
@@ -1,12 +1,8 @@
 #![allow(missing_docs)]
 use dioxus_core::{ScopeState, TaskId};
-use std::{
-    any::Any,
-    cell::{Cell, RefCell},
-    future::{Future, IntoFuture},
-    rc::Rc,
-    sync::Arc,
-};
+use std::{any::Any, cell::Cell, future::Future, rc::Rc, sync::Arc};
+
+use crate::{use_state, UseState};
 
 /// A future that resolves to a value.
 ///
@@ -31,43 +27,30 @@ where
     F: Future<Output = T> + 'static,
     D: UseFutureDep,
 {
+    let val = use_state(cx, || None);
+
     let state = cx.use_hook(move || UseFuture {
         update: cx.schedule_update(),
         needs_regen: Cell::new(true),
-        values: Default::default(),
-        task: Cell::new(None),
+        state: val.clone(),
+        task: Default::default(),
         dependencies: Vec::new(),
-        waker: Default::default(),
     });
 
-    *state.waker.borrow_mut() = None;
-
     if dependencies.clone().apply(&mut state.dependencies) || state.needs_regen.get() {
-        // We don't need regen anymore
-        state.needs_regen.set(false);
+        // kill the old one, if it exists
+        if let Some(task) = state.task.take() {
+            cx.remove_future(task);
+        }
 
         // Create the new future
         let fut = future(dependencies.out());
-
-        // Clone in our cells
-        let values = state.values.clone();
-        let schedule_update = state.update.clone();
-        let waker = state.waker.clone();
-
-        // Cancel the current future
-        if let Some(current) = state.task.take() {
-            cx.remove_future(current);
-        }
+        let val = val.clone();
+        let task = state.task.clone();
 
         state.task.set(Some(cx.push_future(async move {
-            let res = fut.await;
-            values.borrow_mut().push(Box::leak(Box::new(res)));
-
-            // if there's a waker, we dont re-render the component. Instead we just progress that future
-            match waker.borrow().as_ref() {
-                Some(waker) => waker.wake_by_ref(),
-                None => schedule_update(),
-            }
+            val.set(Some(fut.await));
+            task.take();
         })));
     }
 
@@ -80,21 +63,12 @@ pub enum FutureState<'a, T> {
     Regenerating(&'a T), // the old value
 }
 
-pub struct UseFuture<T> {
+pub struct UseFuture<T: 'static> {
     update: Arc<dyn Fn()>,
     needs_regen: Cell<bool>,
-    task: Cell<Option<TaskId>>,
+    task: Rc<Cell<Option<TaskId>>>,
     dependencies: Vec<Box<dyn Any>>,
-    waker: Rc<RefCell<Option<std::task::Waker>>>,
-    values: Rc<RefCell<Vec<*mut T>>>,
-}
-
-impl<T> Drop for UseFuture<T> {
-    fn drop(&mut self) {
-        for value in self.values.take().into_iter() {
-            drop(unsafe { Box::from_raw(value) })
-        }
-    }
+    state: UseState<Option<T>>,
 }
 
 pub enum UseFutureState<'a, T> {
@@ -120,30 +94,16 @@ impl<T> UseFuture<T> {
         }
     }
 
-    // clears the value in the future slot without starting the future over
-    pub fn clear(&self) -> Option<T> {
-        todo!()
-        // (self.update)();
-        // self.slot.replace(None)
-    }
-
     // Manually set the value in the future slot without starting the future over
-    pub fn set(&self, _new_value: T) {
-        // self.slot.set(Some(new_value));
-        // self.needs_regen.set(true);
-        // (self.update)();
-        todo!()
+    pub fn set(&self, new_value: T) {
+        self.state.set(Some(new_value));
     }
 
     /// Return any value, even old values if the future has not yet resolved.
     ///
     /// If the future has never completed, the returned value will be `None`.
     pub fn value(&self) -> Option<&T> {
-        self.values
-            .borrow_mut()
-            .last()
-            .cloned()
-            .map(|x| unsafe { &*x })
+        self.state.current_val.as_ref().as_ref()
     }
 
     /// Get the ID of the future in Dioxus' internal scheduler
@@ -169,35 +129,6 @@ impl<T> UseFuture<T> {
     }
 }
 
-impl<'a, T> IntoFuture for &'a UseFuture<T> {
-    type Output = &'a T;
-    type IntoFuture = UseFutureAwait<'a, T>;
-    fn into_future(self) -> Self::IntoFuture {
-        UseFutureAwait { hook: self }
-    }
-}
-
-pub struct UseFutureAwait<'a, T> {
-    hook: &'a UseFuture<T>,
-}
-
-impl<'a, T> Future for UseFutureAwait<'a, T> {
-    type Output = &'a T;
-
-    fn poll(
-        self: std::pin::Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<Self::Output> {
-        match self.hook.values.borrow_mut().last().cloned() {
-            Some(value) => std::task::Poll::Ready(unsafe { &*value }),
-            None => {
-                self.hook.waker.replace(Some(cx.waker().clone()));
-                std::task::Poll::Pending
-            }
-        }
-    }
-}
-
 pub trait UseFutureDep: Sized + Clone {
     type Out;
     fn out(&self) -> Self::Out;
@@ -343,10 +274,6 @@ mod tests {
                 let blah = "asd";
             });
 
-            let g2 = a.await;
-
-            let g = fut.await;
-
             todo!()
         }
     }