Add custom LSP method for visualizing build plan

forc-pkg/Cargo.toml

@@ -41,5 +41,8 @@ url = { version = "2.2", features = ["serde"] }
 vec1 = "1.8.0"
 walkdir = "2"
 
+[dev-dependencies]
+regex = "^1.10.2"
+
 [target.'cfg(not(target_os = "macos"))'.dependencies]
 sysinfo = "0.29.0"

forc-pkg/src/pkg.rs

@@ -11,7 +11,7 @@ use forc_util::{
 };
 use fuel_abi_types::program_abi;
 use petgraph::{
-    self,
+    self, dot,
     visit::{Bfs, Dfs, EdgeRef, Walker},
     Directed, Direction,
 };
@@ -826,6 +826,28 @@ impl BuildPlan {
                 .flatten()
         })
     }
+
+    /// Returns a [String] representing the build dependency graph in GraphViz DOT format.
+    pub fn visualize(&self, url_file_prefix: Option<String>) -> String {
+        format!(
+            "{:?}",
+            dot::Dot::with_attr_getters(
+                &self.graph,
+                &[dot::Config::NodeNoLabel, dot::Config::EdgeNoLabel],
+                &|_, _| "".to_string(),
+                &|_, nr| {
+                    let url = url_file_prefix.clone().map_or("".to_string(), |prefix| {
+                        self.manifest_map
+                            .get(&nr.1.id())
+                            .map_or("".to_string(), |manifest| {
+                                format!("URL = \"{}{}\"", prefix, manifest.path().to_string_lossy())
+                            })
+                    });
+                    format!("label = \"{}\" shape = box {url}", nr.1.name)
+                },
+            )
+        )
+    }
 }
 
 /// Given a graph and the known project name retrieved from the manifest, produce an iterator
@@ -2698,28 +2720,70 @@ pub fn fuel_core_not_running(node_url: &str) -> anyhow::Error {
     Error::msg(message)
 }
 
-#[test]
-fn test_root_pkg_order() {
-    let current_dir = env!("CARGO_MANIFEST_DIR");
-    let manifest_dir = PathBuf::from(current_dir)
-        .parent()
+#[cfg(test)]
+mod test {
+    use super::*;
+    use regex::Regex;
+
+    fn setup_build_plan() -> BuildPlan {
+        let current_dir = env!("CARGO_MANIFEST_DIR");
+        let manifest_dir = PathBuf::from(current_dir)
+            .parent()
+            .unwrap()
+            .join("test/src/e2e_vm_tests/test_programs/should_pass/forc/workspace_building/");
+        let manifest_file = ManifestFile::from_dir(&manifest_dir).unwrap();
+        let member_manifests = manifest_file.member_manifests().unwrap();
+        let lock_path = manifest_file.lock_path().unwrap();
+        BuildPlan::from_lock_and_manifests(
+            &lock_path,
+            &member_manifests,
+            false,
+            false,
+            Default::default(),
+        )
         .unwrap()
-        .join("test/src/e2e_vm_tests/test_programs/should_pass/forc/workspace_building/");
-    let manifest_file = ManifestFile::from_dir(&manifest_dir).unwrap();
-    let member_manifests = manifest_file.member_manifests().unwrap();
-    let lock_path = manifest_file.lock_path().unwrap();
-    let build_plan = BuildPlan::from_lock_and_manifests(
-        &lock_path,
-        &member_manifests,
-        false,
-        false,
-        Default::default(),
-    )
-    .unwrap();
-    let graph = build_plan.graph();
-    let order: Vec<String> = build_plan
-        .member_nodes()
-        .map(|order| graph[order].name.clone())
-        .collect();
-    assert_eq!(order, vec!["test_lib", "test_contract", "test_script"])
+    }
+
+    #[test]
+    fn test_root_pkg_order() {
+        let build_plan = setup_build_plan();
+        let graph = build_plan.graph();
+        let order: Vec<String> = build_plan
+            .member_nodes()
+            .map(|order| graph[order].name.clone())
+            .collect();
+        assert_eq!(order, vec!["test_lib", "test_contract", "test_script"])
+    }
+
+    #[test]
+    fn test_visualize_with_url_prefix() {
+        let build_plan = setup_build_plan();
+        let result = build_plan.visualize(Some("some-prefix::".to_string()));
+        let re = Regex::new(r#"digraph \{
+    0 \[ label = "test_contract" shape = box URL = "some-prefix::/[[:ascii:]]+/test_contract/Forc.toml"\]
+    1 \[ label = "test_lib" shape = box URL = "some-prefix::/[[:ascii:]]+/test_lib/Forc.toml"\]
+    2 \[ label = "test_script" shape = box URL = "some-prefix::/[[:ascii:]]+/test_script/Forc.toml"\]
+    2 -> 1 \[ \]
+    2 -> 0 \[ \]
+    0 -> 1 \[ \]
+\}
+"#).unwrap();
+        assert!(!re.find(result.as_str()).unwrap().is_empty());
+    }
+
+    #[test]
+    fn test_visualize_without_prefix() {
+        let build_plan = setup_build_plan();
+        let result = build_plan.visualize(None);
+        let expected = r#"digraph {
+    0 [ label = "test_contract" shape = box ]
+    1 [ label = "test_lib" shape = box ]
+    2 [ label = "test_script" shape = box ]
+    2 -> 1 [ ]
+    2 -> 0 [ ]
+    0 -> 1 [ ]
+}
+"#;
+        assert_eq!(expected, result);
+    }
 }

sway-lsp/Cargo.toml

@@ -58,6 +58,7 @@ futures = { version = "0.3", default-features = false, features = [
     "async-await",
 ] }
 pretty_assertions = "1.4.0"
+regex = "^1.10.2"
 sway-lsp-test-utils = { path = "tests/utils" }
 tower = { version = "0.4.12", default-features = false, features = ["util"] }
 

sway-lsp/src/core/session.rs

@@ -399,7 +399,7 @@ impl Session {
 }
 
 /// Create a [BuildPlan] from the given [Url] appropriate for the language server.
-fn build_plan(uri: &Url) -> Result<BuildPlan, LanguageServerError> {
+pub(crate) fn build_plan(uri: &Url) -> Result<BuildPlan, LanguageServerError> {
     let manifest_dir = PathBuf::from(uri.path());
     let manifest =
         ManifestFile::from_dir(&manifest_dir).map_err(|_| DocumentError::ManifestFileNotFound {

sway-lsp/src/handlers/request.rs

@@ -1,7 +1,9 @@
 //! This module is responsible for implementing handlers for Language Server
 //! Protocol. This module specifically handles requests.
 
-use crate::{capabilities, lsp_ext, server_state::ServerState, utils::debug};
+use crate::{
+    capabilities, core::session::build_plan, lsp_ext, server_state::ServerState, utils::debug,
+};
 use forc_tracing::{init_tracing_subscriber, TracingSubscriberOptions, TracingWriterMode};
 use lsp_types::{
     CodeLens, CompletionResponse, DocumentFormattingParams, DocumentSymbolResponse,
@@ -407,7 +409,7 @@ pub fn handle_show_ast(
 }
 
 /// This method is triggered when the use hits enter or pastes a newline in the editor.
-pub(crate) fn on_enter(
+pub(crate) fn handle_on_enter(
     state: &ServerState,
     params: lsp_ext::OnEnterParams,
 ) -> Result<Option<WorkspaceEdit>> {
@@ -426,3 +428,22 @@ pub(crate) fn on_enter(
         }
     }
 }
+
+/// Returns a [String] of the GraphViz DOT representation of a graph.
+pub fn handle_visualize(
+    _state: &ServerState,
+    params: lsp_ext::VisualizeParams,
+) -> Result<Option<String>> {
+    match params.graph_kind.as_str() {
+        "build_plan" => match build_plan(&params.text_document.uri) {
+            Ok(build_plan) => Ok(Some(
+                build_plan.visualize(Some("vscode://file".to_string())),
+            )),
+            Err(err) => {
+                tracing::error!("{}", err.to_string());
+                Ok(None)
+            }
+        },
+        _ => Ok(None),
+    }
+}

sway-lsp/src/lib.rs

@@ -26,6 +26,7 @@ use tower_lsp::{LspService, Server};
 pub async fn start() {
     let (service, socket) = LspService::build(ServerState::new)
         .custom_method("sway/show_ast", ServerState::show_ast)
+        .custom_method("sway/visualize", ServerState::visualize)
         .custom_method("sway/on_enter", ServerState::on_enter)
         .finish();
     Server::new(tokio::io::stdin(), tokio::io::stdout(), socket)

sway-lsp/src/lsp_ext.rs

@@ -18,3 +18,10 @@ pub struct OnEnterParams {
     /// The actual content changes, including the newline.
     pub content_changes: Vec<TextDocumentContentChangeEvent>,
 }
+
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct VisualizeParams {
+    pub text_document: TextDocumentIdentifier,
+    pub graph_kind: String,
+}

sway-lsp/src/server.rs

@@ -4,7 +4,7 @@
 use crate::{
     core::document,
     handlers::{notification, request},
-    lsp_ext::{OnEnterParams, ShowAstParams},
+    lsp_ext::{OnEnterParams, ShowAstParams, VisualizeParams},
     server_state::ServerState,
 };
 use lsp_types::{
@@ -134,6 +134,10 @@ impl ServerState {
     }
 
     pub async fn on_enter(&self, params: OnEnterParams) -> Result<Option<WorkspaceEdit>> {
-        request::on_enter(self, params)
+        request::handle_on_enter(self, params)
+    }
+
+    pub async fn visualize(&self, params: VisualizeParams) -> Result<Option<String>> {
+        request::handle_visualize(self, params)
     }
 }

sway-lsp/tests/integration/lsp.rs

@@ -4,9 +4,14 @@
 
 use crate::{GotoDefinition, HoverDocumentation, Rename};
 use assert_json_diff::assert_json_eq;
+use regex::Regex;
 use serde_json::json;
 use std::{borrow::Cow, path::Path};
-use sway_lsp::{handlers::request, lsp_ext::ShowAstParams, server_state::ServerState};
+use sway_lsp::{
+    handlers::request,
+    lsp_ext::{ShowAstParams, VisualizeParams},
+    server_state::ServerState,
+};
 use tower::{Service, ServiceExt};
 use tower_lsp::{
     jsonrpc::{Id, Request, Response},
@@ -141,6 +146,22 @@ pub(crate) async fn show_ast_request(
     assert_eq!(expected, response.unwrap().unwrap());
 }
 
+pub(crate) async fn visualize_request(server: &ServerState, uri: &Url, graph_kind: &str) {
+    let params = VisualizeParams {
+        text_document: TextDocumentIdentifier { uri: uri.clone() },
+        graph_kind: graph_kind.to_string(),
+    };
+
+    let response = request::handle_visualize(server, params).unwrap().unwrap();
+    let re = Regex::new(r#"digraph \{
+    0 \[ label = "core" shape = box URL = "vscode://file/[[:ascii:]]+/sway-lib-core/Forc.toml"\]
+    1 \[ label = "struct_field_access" shape = box URL = "vscode://file/[[:ascii:]]+/struct_field_access/Forc.toml"\]
+    1 -> 0 \[ \]
+\}
+"#).unwrap();
+    assert!(!re.find(response.as_str()).unwrap().is_empty());
+}
+
 pub(crate) fn semantic_tokens_request(server: &ServerState, uri: &Url) {
     let params = SemanticTokensParams {
         text_document: TextDocumentIdentifier { uri: uri.clone() },

sway-lsp/tests/lib.rs

@@ -123,6 +123,14 @@ async fn show_ast() {
     let _ = server.shutdown_server();
 }
 
+#[tokio::test]
+async fn visualize() {
+    let server = ServerState::default();
+    let uri = open(&server, e2e_test_dir().join("src/main.sw")).await;
+    lsp::visualize_request(&server, &uri, "build_plan").await;
+    let _ = server.shutdown_server();
+}
+
 //------------------- GO TO DEFINITION -------------------//
 
 #[tokio::test]