move hot reloading utilities to dioxus repository

src/server/mod.rs

@@ -1,3 +1,4 @@
+use crate::{builder, plugin::PluginManager, serve::Serve, BuildResult, CrateConfig, Result};
 use axum::{
     body::{Full, HttpBody},
     extract::{ws::Message, Extension, TypedHeader, WebSocketUpgrade},
@@ -9,18 +10,19 @@ use axum::{
 use cargo_metadata::diagnostic::Diagnostic;
 use colored::Colorize;
 use dioxus_core::Template;
+use dioxus_html::HtmlCtx;
+use dioxus_rsx::hot_reload::*;
 use notify::{RecommendedWatcher, Watcher};
-
-use std::{net::UdpSocket, path::PathBuf, process::Command, sync::Arc};
+use std::{
+    net::UdpSocket,
+    path::PathBuf,
+    process::Command,
+    sync::{Arc, Mutex},
+};
+use tokio::sync::broadcast;
 use tower::ServiceBuilder;
 use tower_http::services::fs::{ServeDir, ServeFileSystemResponseBody};
 
-use crate::{builder, plugin::PluginManager, serve::Serve, BuildResult, CrateConfig, Result};
-use tokio::sync::broadcast;
-
-mod hot_reload;
-use hot_reload::*;
-
 pub struct BuildManager {
     config: CrateConfig,
     reload_tx: broadcast::Sender<()>,
@@ -67,6 +69,62 @@ pub async fn startup(port: u16, config: CrateConfig) -> Result<()> {
     Ok(())
 }
 
+pub struct HotReloadState {
+    pub messages: broadcast::Sender<Template<'static>>,
+    pub build_manager: Arc<BuildManager>,
+    pub file_map: Arc<Mutex<FileMap<HtmlCtx>>>,
+    pub watcher_config: CrateConfig,
+}
+
+pub async fn hot_reload_handler(
+    ws: WebSocketUpgrade,
+    _: Option<TypedHeader<headers::UserAgent>>,
+    Extension(state): Extension<Arc<HotReloadState>>,
+) -> impl IntoResponse {
+    ws.on_upgrade(|mut socket| async move {
+        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: Vec<_> = {
+                    state
+                        .file_map
+                        .lock()
+                        .unwrap()
+                        .map
+                        .values()
+                        .filter_map(|(_, template_slot)| *template_slot)
+                        .collect()
+                };
+                for template in templates {
+                    if socket
+                        .send(Message::Text(serde_json::to_string(&template).unwrap()))
+                        .await
+                        .is_err()
+                    {
+                        return;
+                    }
+                }
+            }
+            log::info!("finished");
+        }
+
+        let mut rx = state.messages.subscribe();
+        loop {
+            if let Ok(rsx) = rx.recv().await {
+                if socket
+                    .send(Message::Text(serde_json::to_string(&rsx).unwrap()))
+                    .await
+                    .is_err()
+                {
+                    break;
+                };
+            }
+        }
+    })
+}
+
 #[allow(unused_assignments)]
 pub async fn startup_hot_reload(port: u16, config: CrateConfig) -> Result<()> {
     let first_build_result = crate::builder::build(&config, false)?;
@@ -77,7 +135,9 @@ pub async fn startup_hot_reload(port: u16, config: CrateConfig) -> Result<()> {
 
     let dist_path = config.out_dir.clone();
     let (reload_tx, _) = broadcast::channel(100);
-    let file_map = Arc::new(Mutex::new(FileMap::new(config.crate_dir.clone())));
+    let file_map = Arc::new(Mutex::new(FileMap::<HtmlCtx>::new(
+        config.crate_dir.clone(),
+    )));
     let build_manager = Arc::new(BuildManager {
         config: config.clone(),
         reload_tx: reload_tx.clone(),
@@ -119,13 +179,10 @@ pub async fn startup_hot_reload(port: u16, config: CrateConfig) -> Result<()> {
                         if path.extension().and_then(|p| p.to_str()) != Some("rs") {
                             continue;
                         }
-                        let mut file = File::open(path.clone()).unwrap();
-                        let mut src = String::new();
-                        file.read_to_string(&mut src).expect("Unable to read file");
                         // find changes to the rsx in the file
                         let mut map = file_map.lock().unwrap();
 
-                        match update_rsx(&path, &crate_dir, src, &mut map) {
+                        match map.update_rsx(&path, &crate_dir) {
                             UpdateResult::UpdatedRsx(msgs) => {
                                 messages.extend(msgs);
                             }