From 3d699e7a7552c23d3b40efffc3e5fb14e2f2e9e3 Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Wed, 13 Mar 2024 15:10:44 -0700
Subject: [PATCH] Match on more places for asset reloading

---
 packages/cli/src/server/mod.rs                | 16 +++++++++++-
 packages/cli/src/server/web/hot_reload.rs     | 26 ++++++-------------
 packages/desktop/src/app.rs                   |  3 +++
 packages/dioxus-tui/src/lib.rs                |  1 +
 packages/fullstack/src/hooks/server_future.rs |  1 +
 packages/fullstack/src/hot_reload.rs          |  1 +
 packages/hot-reload/src/lib.rs                |  2 +-
 packages/liveview/src/pool.rs                 |  2 ++
 .../src/hot_reload/hot_reloading_file_map.rs  |  2 --
 9 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/packages/cli/src/server/mod.rs b/packages/cli/src/server/mod.rs
index 7dfb24ac0..c624e2770 100644
--- a/packages/cli/src/server/mod.rs
+++ b/packages/cli/src/server/mod.rs
@@ -2,6 +2,7 @@ use crate::{cfg::ConfigOptsServe, BuildResult, Result};
 use dioxus_cli_config::CrateConfig;
 
 use cargo_metadata::diagnostic::Diagnostic;
+use dioxus_core::Template;
 use dioxus_hot_reload::HotReloadMsg;
 use dioxus_html::HtmlCtx;
 use dioxus_rsx::hot_reload::*;
@@ -25,7 +26,20 @@ pub struct HotReloadState {
     pub messages: broadcast::Sender<HotReloadMsg>,
 
     /// The file map that tracks the state of the projecta
-    pub file_map: Arc<Mutex<FileMap<HtmlCtx>>>,
+    pub file_map: SharedFileMap,
+}
+type SharedFileMap = Arc<Mutex<FileMap<HtmlCtx>>>;
+
+impl HotReloadState {
+    pub fn all_templates(&self) -> Vec<Template> {
+        self.file_map
+            .lock()
+            .unwrap()
+            .map
+            .values()
+            .flat_map(|v| v.templates.values().copied())
+            .collect()
+    }
 }
 
 /// Sets up a file watcher.
diff --git a/packages/cli/src/server/web/hot_reload.rs b/packages/cli/src/server/web/hot_reload.rs
index e72a7cc62..d57bfb330 100644
--- a/packages/cli/src/server/web/hot_reload.rs
+++ b/packages/cli/src/server/web/hot_reload.rs
@@ -27,20 +27,9 @@ async fn hotreload_loop(mut socket: WebSocket, state: HotReloadState) -> anyhow:
     log::info!("🔥 Hot Reload WebSocket connected");
 
     // update any rsx calls that changed before the websocket connected.
-    log::info!("🔮 Finding updates since last compile...");
-
-    let templates = state
-        .file_map
-        .lock()
-        .unwrap()
-        .map
-        .values()
-        .flat_map(|v| v.templates.values().copied())
-        .collect::<Vec<_>>();
-
-    println!("previously changed: {:?}", templates);
-
-    for template in templates {
+    // These templates will be sent down immediately so the page is in sync with the hotreloaded version
+    // The compiled version will be different from the one we actually want to present
+    for template in state.all_templates() {
         socket
             .send(Message::Text(serde_json::to_string(&template).unwrap()))
             .await?;
@@ -65,16 +54,17 @@ async fn hotreload_loop(mut socket: WebSocket, state: HotReloadState) -> anyhow:
 
             let Ok(msg) = msg else { break };
 
-            println!("msg: {:?}", msg);
-
             match msg {
                 HotReloadMsg::UpdateTemplate(template) => {
                     Message::Text(serde_json::to_string(&template).unwrap())
                 }
                 HotReloadMsg::UpdateAsset(asset) => {
-                    Message::Text(format!("asset: {}", asset.display()))
+                    Message::Text(format!("reload-asset: {}", asset.display()))
+                }
+                HotReloadMsg::Shutdown => {
+                    log::info!("🔥 Hot Reload WebSocket shutting down");
+                    break;
                 }
-                HotReloadMsg::Shutdown => todo!(),
             }
         };
 
diff --git a/packages/desktop/src/app.rs b/packages/desktop/src/app.rs
index 31c644c0a..6f5abd279 100644
--- a/packages/desktop/src/app.rs
+++ b/packages/desktop/src/app.rs
@@ -277,6 +277,9 @@ impl App {
             dioxus_hot_reload::HotReloadMsg::Shutdown => {
                 self.control_flow = ControlFlow::Exit;
             }
+
+            // todo: enable asset hotreloading in desktop
+            dioxus_hot_reload::HotReloadMsg::UpdateAsset(_) => {}
         }
     }
 
diff --git a/packages/dioxus-tui/src/lib.rs b/packages/dioxus-tui/src/lib.rs
index 294c0300b..45354b785 100644
--- a/packages/dioxus-tui/src/lib.rs
+++ b/packages/dioxus-tui/src/lib.rs
@@ -184,6 +184,7 @@ impl Driver for DioxusRenderer {
                     dioxus_hot_reload::HotReloadMsg::Shutdown => {
                         std::process::exit(0);
                     }
+                    _ => {}
                 }
             }
         });
diff --git a/packages/fullstack/src/hooks/server_future.rs b/packages/fullstack/src/hooks/server_future.rs
index a79d0ed0c..c988c59b6 100644
--- a/packages/fullstack/src/hooks/server_future.rs
+++ b/packages/fullstack/src/hooks/server_future.rs
@@ -63,6 +63,7 @@ where
     }
 }
 
+#[cfg(feature = "web")]
 #[inline]
 fn kick_future<F, T>(user_fut: F)
 where
diff --git a/packages/fullstack/src/hot_reload.rs b/packages/fullstack/src/hot_reload.rs
index 001894031..812d423b8 100644
--- a/packages/fullstack/src/hot_reload.rs
+++ b/packages/fullstack/src/hot_reload.rs
@@ -35,6 +35,7 @@ impl Default for HotReloadState {
                 dioxus_hot_reload::HotReloadMsg::Shutdown => {
                     std::process::exit(0);
                 }
+                _ => {}
             }
         });
 
diff --git a/packages/hot-reload/src/lib.rs b/packages/hot-reload/src/lib.rs
index c416f239a..2f8c119b3 100644
--- a/packages/hot-reload/src/lib.rs
+++ b/packages/hot-reload/src/lib.rs
@@ -21,7 +21,7 @@ pub enum HotReloadMsg {
     /// A template has been updated
     UpdateTemplate(Template),
 
-    /// A template has been updated
+    /// An asset discovered by rsx! has been updated
     UpdateAsset(PathBuf),
 
     /// The program needs to be recompiled, and the client should shut down
diff --git a/packages/liveview/src/pool.rs b/packages/liveview/src/pool.rs
index 614143841..1c69f6440 100644
--- a/packages/liveview/src/pool.rs
+++ b/packages/liveview/src/pool.rs
@@ -218,6 +218,8 @@ pub async fn run(mut vdom: VirtualDom, ws: impl LiveViewSocket) -> Result<(), Li
                     dioxus_hot_reload::HotReloadMsg::UpdateTemplate(new_template) => {
                         vdom.replace_template(new_template);
                     }
+                    // todo: enable hotreloading in liveview
+                    dioxus_hot_reload::HotReloadMsg::UpdateAsset(_) => {}
                     dioxus_hot_reload::HotReloadMsg::Shutdown => {
                         std::process::exit(0);
                     },
diff --git a/packages/rsx/src/hot_reload/hot_reloading_file_map.rs b/packages/rsx/src/hot_reload/hot_reloading_file_map.rs
index bf1701420..055bff682 100644
--- a/packages/rsx/src/hot_reload/hot_reloading_file_map.rs
+++ b/packages/rsx/src/hot_reload/hot_reloading_file_map.rs
@@ -243,8 +243,6 @@ impl<Ctx: HotReloadingContext> FileMap<Ctx> {
             collect_assetlike_attrs(node, &mut asset_urls);
         }
 
-        println!("asset urls: {:?}", asset_urls);
-
         asset_urls
     }