From 2a6a8b42a6b010ce23dc84096d07646f7ac492a4 Mon Sep 17 00:00:00 2001
From: JohnnyMorganz <johnnymorganz@outlook.com>
Date: Mon, 1 Aug 2022 09:07:07 +0100
Subject: [PATCH] Add `--watch` to sourcemap generation (#602)

* Implement watch argument

* Add forget call

* Clippy fixes

* Update changelog
---
 CHANGELOG.md         |  3 ++
 src/cli/sourcemap.rs | 65 ++++++++++++++++++++++++++++++++++----------
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 56999e1b..566c4af3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,9 @@
 # Rojo Changelog
 
 ## Unreleased Changes
+* Added `--watch` flag to `rojo sourcemap` ([#602])
+
+[#602]: https://github.com/rojo-rbx/rojo/pull/602
 
 ## [7.2.1] - July 8, 2022
 * Fixed notification sound by changing it to a generic sound. ([#566])
diff --git a/src/cli/sourcemap.rs b/src/cli/sourcemap.rs
index 9d761e71..353c6020 100644
--- a/src/cli/sourcemap.rs
+++ b/src/cli/sourcemap.rs
@@ -1,5 +1,6 @@
 use std::{
     io::{BufWriter, Write},
+    mem::forget,
     path::{Path, PathBuf},
 };
 
@@ -8,6 +9,7 @@ use fs_err::File;
 use memofs::Vfs;
 use rbx_dom_weak::types::Ref;
 use serde::Serialize;
+use tokio::runtime::Runtime;
 
 use crate::{
     serve_session::ServeSession,
@@ -50,6 +52,10 @@ pub struct SourcemapCommand {
     /// If non-script files should be included or not. Defaults to false.
     #[clap(long)]
     pub include_non_scripts: bool,
+
+    /// Whether to automatically recreate a snapshot when any input files change.
+    #[clap(long)]
+    pub watch: bool,
 }
 
 impl SourcemapCommand {
@@ -58,9 +64,10 @@ impl SourcemapCommand {
 
         log::trace!("Constructing in-memory filesystem");
         let vfs = Vfs::new_default();
+        vfs.set_watch_enabled(self.watch);
 
         let session = ServeSession::new(vfs, &project_path)?;
-        let tree = session.tree();
+        let mut cursor = session.message_queue().cursor();
 
         let filter = if self.include_non_scripts {
             filter_nothing
@@ -68,19 +75,24 @@ impl SourcemapCommand {
             filter_non_scripts
         };
 
-        let root_node = recurse_create_node(&tree, tree.get_root_id(), session.root_dir(), filter);
+        write_sourcemap(&session, self.output.as_deref(), filter)?;
 
-        if let Some(output_path) = self.output {
-            let mut file = BufWriter::new(File::create(&output_path)?);
-            serde_json::to_writer(&mut file, &root_node)?;
-            file.flush()?;
+        if self.watch {
+            let rt = Runtime::new().unwrap();
 
-            println!("Created sourcemap at {}", output_path.display());
-        } else {
-            let output = serde_json::to_string(&root_node)?;
-            println!("{}", output);
+            loop {
+                let receiver = session.message_queue().subscribe(cursor);
+                let (new_cursor, _patch_set) = rt.block_on(receiver).unwrap();
+                cursor = new_cursor;
+
+                write_sourcemap(&session, self.output.as_deref(), filter)?;
+            }
         }
 
+        // Avoid dropping ServeSession: it's potentially VERY expensive to drop
+        // and we're about to exit anyways.
+        forget(session);
+
         Ok(())
     }
 }
@@ -90,10 +102,10 @@ fn filter_nothing(_instance: &InstanceWithMeta) -> bool {
 }
 
 fn filter_non_scripts(instance: &InstanceWithMeta) -> bool {
-    match instance.class_name() {
-        "Script" | "LocalScript" | "ModuleScript" => true,
-        _ => false,
-    }
+    matches!(
+        instance.class_name(),
+        "Script" | "LocalScript" | "ModuleScript"
+    )
 }
 
 fn recurse_create_node(
@@ -106,7 +118,7 @@ fn recurse_create_node(
 
     let mut children = Vec::new();
     for &child_id in instance.children() {
-        if let Some(child_node) = recurse_create_node(tree, child_id, &project_dir, filter) {
+        if let Some(child_node) = recurse_create_node(tree, child_id, project_dir, filter) {
             children.push(child_node);
         }
     }
@@ -134,3 +146,26 @@ fn recurse_create_node(
         children,
     })
 }
+
+fn write_sourcemap(
+    session: &ServeSession,
+    output: Option<&Path>,
+    filter: fn(&InstanceWithMeta) -> bool,
+) -> anyhow::Result<()> {
+    let tree = session.tree();
+
+    let root_node = recurse_create_node(&tree, tree.get_root_id(), session.root_dir(), filter);
+
+    if let Some(output_path) = output {
+        let mut file = BufWriter::new(File::create(&output_path)?);
+        serde_json::to_writer(&mut file, &root_node)?;
+        file.flush()?;
+
+        println!("Created sourcemap at {}", output_path.display());
+    } else {
+        let output = serde_json::to_string(&root_node)?;
+        println!("{}", output);
+    }
+
+    Ok(())
+}