Add visualize command

stack-graphs/src/storage.rs

@@ -473,12 +473,18 @@ impl SQLiteReader {
     /// Returns a [`Files`][] value that can be used to iterate over all descendants of a
     /// file or directory in the database.
     pub fn list_file_or_directory<'a>(
-        &'a mut self,
+        &'a self,
+        file_or_directory: &Path,
+    ) -> Result<Files<'a, [String; 1]>> {
+        Self::list_file_or_directory_inner(&self.conn, file_or_directory)
+    }
+
+    fn list_file_or_directory_inner<'a>(
+        conn: &'a Connection,
         file_or_directory: &Path,
     ) -> Result<Files<'a, [String; 1]>> {
         let file_or_directory = file_or_directory.to_string_lossy().to_string();
-        self.conn
-            .prepare("SELECT file, tag, error FROM graphs WHERE path_descendant_of(file, ?)")
+        conn.prepare("SELECT file, tag, error FROM graphs WHERE path_descendant_of(file, ?)")
             .map(|stmt| Files(stmt, [file_or_directory]))
             .map_err(|e| e.into())
     }
@@ -507,6 +513,24 @@ impl SQLiteReader {
         Ok(())
     }
 
+    pub fn load_graph_for_file_or_directory(
+        &mut self,
+        file_or_directory: &Path,
+        cancellation_flag: &dyn CancellationFlag,
+    ) -> Result<()> {
+        for file in Self::list_file_or_directory_inner(&self.conn, file_or_directory)?.try_iter()? {
+            cancellation_flag.check("loading graphs")?;
+            let file = file?;
+            Self::load_graph_for_file_inner(
+                &file.path.to_string_lossy(),
+                &mut self.graph,
+                &mut self.loaded_graphs,
+                &self.conn,
+            )?;
+        }
+        Ok(())
+    }
+
     /// Ensure the paths starting a the given node are loaded.
     fn load_paths_for_node(
         &mut self,

stack-graphs/src/visualization.rs

@@ -17,8 +17,8 @@ use crate::serde::Filter;
 use crate::stitching::Database;
 
 static CSS: &'static str = include_str!("visualization/visualization.css");
-static D3: &'static str = include_str!("visualization/d3.v7.min.js");
-static D3_DAG: &'static str = include_str!("visualization/d3-dag.v0.10.0.min.js");
+static D3: &'static str = include_str!("visualization/d3.min.js");
+static D3_DAG: &'static str = include_str!("visualization/d3-dag.min.js");
 static JS: &'static str = include_str!("visualization/visualization.js");
 
 static PKG: &'static str = env!("CARGO_PKG_NAME");

tree-sitter-stack-graphs/src/cli.rs

@@ -70,6 +70,7 @@ pub mod query;
 pub mod status;
 pub mod test;
 mod util;
+pub mod visualize;
 
 pub mod path_loading {
     use std::path::PathBuf;
@@ -87,6 +88,7 @@ pub mod path_loading {
     use crate::cli::r#match::MatchArgs;
     use crate::cli::status::StatusArgs;
     use crate::cli::test::TestArgs;
+    use crate::cli::visualize::VisualizeArgs;
 
     use super::database::DatabaseArgs;
 
@@ -102,6 +104,7 @@ pub mod path_loading {
         Query(Query),
         Status(Status),
         Test(Test),
+        Visualize(Visualize),
     }
 
     impl Subcommands {
@@ -117,6 +120,7 @@ pub mod path_loading {
                 Self::Query(cmd) => cmd.run(default_db_path),
                 Self::Status(cmd) => cmd.run(default_db_path),
                 Self::Test(cmd) => cmd.run(),
+                Self::Visualize(cmd) => cmd.run(default_db_path),
             }
         }
     }
@@ -269,6 +273,22 @@ pub mod path_loading {
             self.test_args.run(loader)
         }
     }
+
+    /// Visualize command
+    #[derive(clap::Parser)]
+    pub struct Visualize {
+        #[clap(flatten)]
+        db_args: DatabaseArgs,
+        #[clap(flatten)]
+        visualize_args: VisualizeArgs,
+    }
+
+    impl Visualize {
+        pub fn run(self, default_db_path: PathBuf) -> anyhow::Result<()> {
+            let db_path = self.db_args.get_or(default_db_path);
+            self.visualize_args.run(&db_path)
+        }
+    }
 }
 
 pub mod provided_languages {
@@ -287,6 +307,7 @@ pub mod provided_languages {
     use crate::cli::r#match::MatchArgs;
     use crate::cli::status::StatusArgs;
     use crate::cli::test::TestArgs;
+    use crate::cli::visualize::VisualizeArgs;
     use crate::loader::LanguageConfiguration;
 
     use super::database::DatabaseArgs;
@@ -303,6 +324,7 @@ pub mod provided_languages {
         Query(Query),
         Status(Status),
         Test(Test),
+        Visualize(Visualize),
     }
 
     impl Subcommands {
@@ -322,6 +344,7 @@ pub mod provided_languages {
                 Self::Query(cmd) => cmd.run(default_db_path),
                 Self::Status(cmd) => cmd.run(default_db_path),
                 Self::Test(cmd) => cmd.run(configurations),
+                Self::Visualize(cmd) => cmd.run(default_db_path),
             }
         }
     }
@@ -482,4 +505,20 @@ pub mod provided_languages {
             self.test_args.run(loader)
         }
     }
+
+    /// Visualize command
+    #[derive(clap::Parser)]
+    pub struct Visualize {
+        #[clap(flatten)]
+        db_args: DatabaseArgs,
+        #[clap(flatten)]
+        visualize_args: VisualizeArgs,
+    }
+
+    impl Visualize {
+        pub fn run(self, default_db_path: PathBuf) -> anyhow::Result<()> {
+            let db_path = self.db_args.get_or(default_db_path);
+            self.visualize_args.run(&db_path)
+        }
+    }
 }

tree-sitter-stack-graphs/src/cli/visualize.rs

@@ -0,0 +1,69 @@
+// -*- coding: utf-8 -*-
+// ------------------------------------------------------------------------------------------------
+// Copyright © 2023, stack-graphs authors.
+// Licensed under either of Apache License, Version 2.0, or MIT license, at your option.
+// Please see the LICENSE-APACHE or LICENSE-MIT files in this distribution for license details.
+// ------------------------------------------------------------------------------------------------
+
+use clap::Args;
+use clap::ValueHint;
+use stack_graphs::serde::NoFilter;
+use stack_graphs::stitching::Database;
+use stack_graphs::storage::SQLiteReader;
+use stack_graphs::NoCancellation;
+use std::path::Path;
+use std::path::PathBuf;
+
+/// Visualize database
+#[derive(Args)]
+#[clap(after_help = r#"LIMITATIONS:
+    Visualizations will only work for very small stack graphs. This command is
+    useful for debugging minimal examples, but running it on any real-world code
+    will most likely result in HTML files that will not load in any browser.
+"#)]
+pub struct VisualizeArgs {
+    /// Source file or directory paths.
+    #[clap(
+        value_name = "SOURCE_PATH",
+        value_hint = ValueHint::AnyPath,
+    )]
+    pub source_paths: Vec<PathBuf>,
+
+    #[clap(
+        long,
+        short = 'o',
+        value_name = "OUTPUT_PATH",
+        value_hint = ValueHint::AnyPath,
+        default_value = "stack-graph.html",
+    )]
+    pub output: PathBuf,
+}
+
+impl VisualizeArgs {
+    pub fn run(self, db_path: &Path) -> anyhow::Result<()> {
+        let cancellation_flag = &NoCancellation;
+        let mut db = SQLiteReader::open(&db_path)?;
+        for source_path in &self.source_paths {
+            let source_path = source_path.canonicalize()?;
+            db.load_graph_for_file_or_directory(&source_path, cancellation_flag)?;
+        }
+        let (graph, _, _) = db.get();
+        let starting_nodes = graph
+            .iter_nodes()
+            .filter(|n| graph[*n].is_reference())
+            .collect::<Vec<_>>();
+        let mut complete_paths_db = Database::new();
+        db.find_all_complete_partial_paths(starting_nodes, cancellation_flag, |g, ps, p| {
+            complete_paths_db.add_partial_path(g, ps, p.clone());
+        })?;
+        let (graph, partials, _) = db.get();
+        let html =
+            graph.to_html_string("stack-graph", partials, &mut complete_paths_db, &NoFilter)?;
+        if let Some(dir) = self.output.parent() {
+            std::fs::create_dir_all(dir)?;
+        }
+        std::fs::write(&self.output, html)?;
+        println!("Visualization at {}", self.output.display());
+        Ok(())
+    }
+}