From 732ad73e72863636719a77e08847f92d19519c32 Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Sat, 6 May 2023 10:01:15 +0200
Subject: [PATCH] Sourcemap performance improvements (#668)

This PR brings two performance improvements to the `rojo sourcemap`
command:

- Use `rayon` with a small threadpool to parallelize sourcemap
generation while still keeping startup cost very low
- Remove conversions to owned strings and use lifetimes tied to the dom
instead, which mostly improves performance with the
`--include-non-scripts` flag enabled

From my personal testing on an M1 mac this decreases the sourcemap
generation time of our games by 2x or more, from ~20ms to ~8ms on one
project and ~30ms to ~15ms on another. Generation is pretty fast to
begin with but since sourcemaps are heavily used in interactive tools
(like luau-lsp) a difference of a couple frames can be great for ux.
---
 CHANGELOG.md         |  3 +++
 Cargo.lock           |  2 ++
 Cargo.toml           |  2 ++
 src/cli/sourcemap.rs | 39 +++++++++++++++++++++++----------------
 4 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index deba0561c..b0a7b948a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,9 @@
 # Rojo Changelog
 
 ## Unreleased Changes
+* Significantly improved performance of `rojo sourcemap`. ([#668])
+
+[#668]: https://github.com/rojo-rbx/rojo/pull/668
 
 ## [7.3.0] - April 22, 2023
 * Added `$attributes` to project format. ([#574])
diff --git a/Cargo.lock b/Cargo.lock
index 03c31289b..12e9f8602 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1900,10 +1900,12 @@ dependencies = [
  "maplit",
  "memofs",
  "notify",
+ "num_cpus",
  "opener",
  "paste",
  "pretty_assertions",
  "profiling",
+ "rayon",
  "rbx_binary",
  "rbx_dom_weak",
  "rbx_reflection",
diff --git a/Cargo.toml b/Cargo.toml
index d9b38bf70..bbbbcbf90 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -72,7 +72,9 @@ jod-thread = "0.1.2"
 log = "0.4.14"
 maplit = "1.0.2"
 notify = "4.0.17"
+num_cpus = "1.15.0"
 opener = "0.5.0"
+rayon = "1.7.0"
 reqwest = { version = "0.11.10", features = ["blocking", "json", "native-tls-vendored"] }
 ritz = "0.1.0"
 roblox_install = "1.0.0"
diff --git a/src/cli/sourcemap.rs b/src/cli/sourcemap.rs
index 353c6020e..f7bc1691a 100644
--- a/src/cli/sourcemap.rs
+++ b/src/cli/sourcemap.rs
@@ -7,6 +7,7 @@ use std::{
 use clap::Parser;
 use fs_err::File;
 use memofs::Vfs;
+use rayon::prelude::*;
 use rbx_dom_weak::types::Ref;
 use serde::Serialize;
 use tokio::runtime::Runtime;
@@ -23,15 +24,15 @@ const PATH_STRIP_FAILED_ERR: &str = "Failed to create relative paths for project
 /// Representation of a node in the generated sourcemap tree.
 #[derive(Serialize)]
 #[serde(rename_all = "camelCase")]
-struct SourcemapNode {
-    name: String,
-    class_name: String,
+struct SourcemapNode<'a> {
+    name: &'a str,
+    class_name: &'a str,
 
     #[serde(skip_serializing_if = "Vec::is_empty")]
     file_paths: Vec<PathBuf>,
 
     #[serde(skip_serializing_if = "Vec::is_empty")]
-    children: Vec<SourcemapNode>,
+    children: Vec<SourcemapNode<'a>>,
 }
 
 /// Generates a sourcemap file from the Rojo project.
@@ -75,6 +76,13 @@ impl SourcemapCommand {
             filter_non_scripts
         };
 
+        // Pre-build a rayon threadpool with a low number of threads to avoid
+        // dynamic creation overhead on systems with a high number of cpus.
+        rayon::ThreadPoolBuilder::new()
+            .num_threads(num_cpus::get().min(6))
+            .build_global()
+            .unwrap();
+
         write_sourcemap(&session, self.output.as_deref(), filter)?;
 
         if self.watch {
@@ -108,20 +116,19 @@ fn filter_non_scripts(instance: &InstanceWithMeta) -> bool {
     )
 }
 
-fn recurse_create_node(
-    tree: &RojoTree,
+fn recurse_create_node<'a>(
+    tree: &'a RojoTree,
     referent: Ref,
     project_dir: &Path,
     filter: fn(&InstanceWithMeta) -> bool,
-) -> Option<SourcemapNode> {
+) -> Option<SourcemapNode<'a>> {
     let instance = tree.get_instance(referent).expect("instance did not exist");
 
-    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) {
-            children.push(child_node);
-        }
-    }
+    let children: Vec<_> = instance
+        .children()
+        .par_iter()
+        .filter_map(|&child_id| recurse_create_node(tree, child_id, project_dir, filter))
+        .collect();
 
     // If this object has no children and doesn't pass the filter, it doesn't
     // contain any information we're looking for.
@@ -140,8 +147,8 @@ fn recurse_create_node(
         .collect();
 
     Some(SourcemapNode {
-        name: instance.name().to_string(),
-        class_name: instance.class_name().to_string(),
+        name: instance.name(),
+        class_name: instance.class_name(),
         file_paths,
         children,
     })
@@ -157,7 +164,7 @@ fn write_sourcemap(
     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)?);
+        let mut file = BufWriter::new(File::create(output_path)?);
         serde_json::to_writer(&mut file, &root_node)?;
         file.flush()?;