From 8ff50df3a5dc14b90b15fd84122a12436326144a Mon Sep 17 00:00:00 2001
From: Evan Almloff <ealmlof1@stumail.jccc.edu>
Date: Wed, 15 Jun 2022 13:29:35 -0500
Subject: [PATCH] preserve whitespace in rsx

---
 src/server/hot_reload.rs | 37 ++++++++++++++++---------------------
 src/server/mod.rs        | 29 ++++++++++++++++++++++-------
 2 files changed, 38 insertions(+), 28 deletions(-)

diff --git a/src/server/hot_reload.rs b/src/server/hot_reload.rs
index a9f51abe7..f79a7a50e 100644
--- a/src/server/hot_reload.rs
+++ b/src/server/hot_reload.rs
@@ -15,6 +15,7 @@ pub use std::time::SystemTime;
 pub use std::{fs, io, path::Path};
 pub use std::{fs::File, io::Read};
 pub use syn::__private::ToTokens;
+use syn::spanned::Spanned;
 use tokio::sync::broadcast;
 
 pub struct HotReloadState {
@@ -91,17 +92,16 @@ pub async fn hot_reload_handler(
                         if let Ok(old) = syn::parse_file(&v) {
                             if let DiffResult::RsxChanged(changed) = find_rsx(&new, &old) {
                                 for (old, new) in changed.into_iter() {
-                                    if let Some(hr) = get_min_location(
+                                    let hr = get_location(
                                         k.strip_prefix(&state.watcher_config.crate_dir).unwrap(),
                                         old.to_token_stream(),
-                                    ) {
-                                        let rsx = new.to_string();
-                                        let msg = SetRsxMessage {
-                                            location: hr,
-                                            new_text: rsx,
-                                        };
-                                        messages.push(msg);
-                                    }
+                                    );
+                                    let rsx = new.to_string();
+                                    let msg = SetRsxMessage {
+                                        location: hr,
+                                        new_text: rsx,
+                                    };
+                                    messages.push(msg);
                                 }
                             }
                         }
@@ -159,16 +159,11 @@ pub async fn hot_reload_handler(
     })
 }
 
-pub fn get_min_location(path: &Path, ts: TokenStream) -> Option<CodeLocation> {
-    ts.into_iter()
-        .map(|tree| {
-            let location = tree.span();
-            let start = location.start();
-            CodeLocation {
-                file: path.display().to_string(),
-                line: start.line as u32,
-                column: start.column as u32 + 1,
-            }
-        })
-        .min_by(|cl1, cl2| cl1.line.cmp(&cl2.line).then(cl1.column.cmp(&cl2.column)))
+pub fn get_location(path: &Path, ts: TokenStream) -> CodeLocation {
+    let span = ts.span().start();
+    CodeLocation {
+        file: path.display().to_string(),
+        line: span.line as u32,
+        column: span.column as u32 + 1,
+    }
 }
diff --git a/src/server/mod.rs b/src/server/mod.rs
index 8a98fd9d9..3397beac6 100644
--- a/src/server/mod.rs
+++ b/src/server/mod.rs
@@ -7,6 +7,7 @@ use axum::{
     Router,
 };
 use notify::{RecommendedWatcher, Watcher};
+use syn::spanned::Spanned;
 
 use std::{path::PathBuf, sync::Arc};
 use tower::ServiceBuilder;
@@ -93,16 +94,30 @@ pub async fn startup_hot_reload(config: CrateConfig) -> Result<()> {
                                 DiffResult::RsxChanged(changed) => {
                                     log::info!("reloading rsx");
                                     for (old, new) in changed.into_iter() {
-                                        if let Some(hr) = get_min_location(
+                                        let hr = get_location(
                                             &path.strip_prefix(&crate_dir).unwrap().to_path_buf(),
                                             old.to_token_stream(),
-                                        ) {
-                                            let rsx = new.to_string();
-                                            let _ = hot_reload_tx.send(SetRsxMessage {
-                                                location: hr,
-                                                new_text: rsx,
-                                            });
+                                        );
+                                        // get the original source code to preserve whitespace
+                                        let span = new.span();
+                                        let start = span.start();
+                                        let end = span.end();
+                                        let mut lines: Vec<_> = src
+                                            .lines()
+                                            .skip(start.line - 1)
+                                            .take(end.line - start.line + 1)
+                                            .collect();
+                                        if let Some(first) = lines.first_mut() {
+                                            *first = first.split_at(start.column).1;
                                         }
+                                        if let Some(last) = lines.last_mut() {
+                                            *last = last.split_at(end.column).0;
+                                        }
+                                        let rsx = lines.join("\n");
+                                        let _ = hot_reload_tx.send(SetRsxMessage {
+                                            location: hr,
+                                            new_text: rsx,
+                                        });
                                     }
                                 }
                             }