From 9bf5bd11e2ddc7dce625e76d3ca6412a6005ca0e Mon Sep 17 00:00:00 2001
From: Lucien Greathouse <me@lpghatguy.com>
Date: Wed, 2 Jan 2019 00:18:28 -0800
Subject: [PATCH] Add visualization stuff using GraphViz at /api/visualize

---
 server/src/lib.rs            |  1 +
 server/src/visualize_tree.rs | 42 ++++++++++++++++++++++++++++++++++++
 server/src/web.rs            | 26 ++++++++++++++++++++++
 3 files changed, 69 insertions(+)
 create mode 100644 server/src/visualize_tree.rs

diff --git a/server/src/lib.rs b/server/src/lib.rs
index fd3d600b..80b7235c 100644
--- a/server/src/lib.rs
+++ b/server/src/lib.rs
@@ -18,5 +18,6 @@ pub mod rbx_session;
 pub mod rbx_snapshot;
 pub mod session;
 pub mod session_id;
+pub mod visualize_tree;
 pub mod web;
 pub mod web_util;
\ No newline at end of file
diff --git a/server/src/visualize_tree.rs b/server/src/visualize_tree.rs
new file mode 100644
index 00000000..9faee18f
--- /dev/null
+++ b/server/src/visualize_tree.rs
@@ -0,0 +1,42 @@
+use std::fmt;
+use rbx_tree::{RbxTree, RbxId};
+
+static GRAPHVIZ_HEADER: &str = r#"
+digraph RojoTree {
+    rankdir = "LR";
+    graph [
+        ranksep = "0.7",
+        nodesep = "0.5",
+    ];
+    node [
+        fontname = "Hack",
+        shape = "record",
+    ];
+"#;
+
+pub struct VisualizeTree<'a>(pub &'a RbxTree);
+
+impl<'a> fmt::Display for VisualizeTree<'a> {
+    fn fmt(&self, output: &mut fmt::Formatter) -> fmt::Result {
+        writeln!(output, "{}", GRAPHVIZ_HEADER)?;
+
+        visualize_node(self.0, self.0.get_root_id(), output)?;
+
+        writeln!(output, "}}")?;
+
+        Ok(())
+    }
+}
+
+fn visualize_node(tree: &RbxTree, id: RbxId, output: &mut fmt::Formatter) -> fmt::Result {
+    let node = tree.get_instance(id).unwrap();
+
+    writeln!(output, "    \"{}\" [label=\"{}\"]", id, node.name)?;
+
+    for &child_id in node.get_children_ids() {
+        writeln!(output, "    \"{}\" -> \"{}\"", id, child_id)?;
+        visualize_node(tree, child_id, output)?;
+    }
+
+    Ok(())
+}
\ No newline at end of file
diff --git a/server/src/web.rs b/server/src/web.rs
index 7347f837..f73ecef2 100644
--- a/server/src/web.rs
+++ b/server/src/web.rs
@@ -1,6 +1,8 @@
 use std::{
     borrow::Cow,
     collections::HashMap,
+    io::Write,
+    process::{Command, Stdio},
     sync::{mpsc, Arc},
 };
 
@@ -17,6 +19,7 @@ use crate::{
     session_id::SessionId,
     project::InstanceProjectNodeMetadata,
     rbx_snapshot::InstanceChanges,
+    visualize_tree::VisualizeTree,
 };
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -160,6 +163,29 @@ impl Server {
                 })
             },
 
+            (GET) (/api/visualize) => {
+                let rbx_session = self.session.rbx_session.lock().unwrap();
+                let tree = rbx_session.get_tree();
+
+                let dot_source = format!("{}", VisualizeTree(tree));
+
+                let mut child = Command::new("dot")
+                    .arg("-Tsvg")
+                    .stdin(Stdio::piped())
+                    .stdout(Stdio::piped())
+                    .spawn()
+                    .expect("Failed to spawn GraphViz process -- make sure it's installed in order to use /api/visualize");
+
+                {
+                    let stdin = child.stdin.as_mut().expect("Failed to open stdin");
+                    stdin.write_all(dot_source.as_bytes()).expect("Failed to write to stdin");
+                }
+
+                let output = child.wait_with_output().expect("Failed to read stdout");
+
+                Response::svg(String::from_utf8_lossy(&output.stdout))
+            },
+
             _ => Response::empty_404()
         )
     }