feat: profile CLI command (ribru17#136)

Author: ribru17

README.md

@@ -245,6 +245,24 @@ ts_query_ls lint ./queries
 ts_query_ls lint --help
 ```
 
+### Profiler
+
+The server can be used to profile individual query patterns to check for
+patterns which are very slow to compile (often because they are too complex).
+This can be done via the `profile` subcommand, which prints each pattern's file
+path, start line, and the time (in milliseconds) that it took to compile.
+Alternatively, it can also time the entire query file itself (rather than each
+pattern inside of it).
+
+```sh
+ts_query_ls profile ./queries
+# Use this command for the full documentation
+ts_query_ls profile --help
+```
+
+> **NOTE:** This command will not warm up the cache for you, so it may be best
+> to run more than once.
+
 ## Checklist
 
 - [x] References for captures
@@ -266,7 +284,7 @@ ts_query_ls lint --help
     query file. This should either be implemented as a user command, or core
     methods should be exposed to gather pattern information more efficiently~~
   - For now, this has been made possible due to caching and spawning query scans
-    on a separate, blocking thread. Ideally in the future the kinks of query
+    on a separate, blocking thread. Ideally, in the future, the kinks of query
     creation will be ironed out so query creation will be quicker, and this
     logic can be simplified
 - [x] Recognize parsers built for `WASM`

src/cli/mod.rs

@@ -5,6 +5,7 @@ use walkdir::WalkDir;
 pub(super) mod check;
 pub(super) mod format;
 pub(super) mod lint;
+pub(super) mod profile;
 
 pub(in crate::cli) fn get_scm_files(directories: &[PathBuf]) -> Vec<PathBuf> {
     directories

src/cli/profile.rs

@@ -0,0 +1,107 @@
+use std::{
+    env, fs,
+    path::PathBuf,
+    sync::{Arc, LazyLock},
+    time::Instant,
+};
+
+use dashmap::DashMap;
+use futures::future::join_all;
+use tower_lsp::lsp_types::Url;
+use tree_sitter::{Parser, Query, QueryCursor, StreamingIterator as _};
+use ts_query_ls::Options;
+
+use crate::{LanguageData, QUERY_LANGUAGE, handlers::did_open::init_language_data, util};
+
+use super::get_scm_files;
+
+static LANGUAGE_CACHE: LazyLock<DashMap<String, Arc<LanguageData>>> = LazyLock::new(DashMap::new);
+
+static PATTERN_DEFINITION_QUERY: LazyLock<Query> =
+    LazyLock::new(|| Query::new(&QUERY_LANGUAGE, "(program (definition) @def)").unwrap());
+
+pub async fn profile_directories(directories: &[PathBuf], config: String, broad: bool) {
+    let Ok(options) = serde_json::from_str::<Options>(&config) else {
+        eprintln!("Could not parse the provided configuration");
+        return;
+    };
+    let scm_files = if directories.is_empty() {
+        get_scm_files(&[env::current_dir().expect("Failed to get current directory")])
+    } else {
+        get_scm_files(directories)
+    };
+    let tasks = scm_files.into_iter().filter_map(|path| {
+        let uri = Url::from_file_path(path.canonicalize().unwrap()).unwrap();
+        let path_str = path.to_string_lossy().to_string();
+        let language_name = util::get_language_name(&uri, &options);
+        let language_data = language_name.and_then(|name| {
+            LANGUAGE_CACHE.get(&name).as_deref().cloned().or_else(|| {
+                util::get_language(&name, &options).map(|lang| Arc::new(init_language_data(lang)))
+            })
+        });
+        if let Some(lang_data) = language_data {
+            let lang = lang_data.language.clone().unwrap();
+            if let Ok(source) = fs::read_to_string(&path) {
+                Some(tokio::spawn(async move {
+                    let mut results = Vec::new();
+                    if broad {
+                        let now = Instant::now();
+                        let _ = Query::new(&lang, &source);
+                        results.push((path_str.clone(), 1, now.elapsed().as_millis()));
+                    } else {
+                        let mut parser = Parser::new();
+                        parser.set_language(&QUERY_LANGUAGE).unwrap();
+                        let tree = parser.parse(&source, None).expect("Tree should exist");
+                        let mut cursor = QueryCursor::new();
+                        let source_bytes = source.as_bytes();
+                        let mut matches = cursor.matches(
+                            &PATTERN_DEFINITION_QUERY,
+                            tree.root_node(),
+                            source_bytes,
+                        );
+                        while let Some(match_) = matches.next() {
+                            for capture in match_.captures {
+                                let now = Instant::now();
+                                let _ = Query::new(
+                                    &lang,
+                                    capture
+                                        .node
+                                        .utf8_text(source_bytes)
+                                        .expect("Source should be UTF-8"),
+                                );
+                                results.push((
+                                    path_str.clone(),
+                                    capture.node.start_position().row + 1,
+                                    now.elapsed().as_millis(),
+                                ));
+                            }
+                        }
+                    }
+                    results
+                }))
+            } else {
+                eprintln!("Failed to read {:?}", path.canonicalize().unwrap());
+                None
+            }
+        } else {
+            eprintln!(
+                "Could not retrieve language for {:?}",
+                path.canonicalize().unwrap()
+            );
+            None
+        }
+    });
+    let results = join_all(tasks).await;
+    let mut results = results
+        .into_iter()
+        .flat_map(|r| r.unwrap_or_default())
+        .collect::<Vec<_>>();
+    results.sort_unstable_by(|a, b| a.2.cmp(&b.2));
+    for (path, row, time) in results {
+        if broad {
+            println!("Query at {path} took {time}ms to compile");
+        } else {
+            println!("Pattern in {path} at line {row} took {time}ms to compile");
+        }
+    }
+}

src/main.rs

@@ -1,5 +1,8 @@
 use clap::{Parser, Subcommand};
-use cli::{check::check_directories, format::format_directories, lint::lint_directories};
+use cli::{
+    check::check_directories, format::format_directories, lint::lint_directories,
+    profile::profile_directories,
+};
 use core::fmt;
 use std::{
     collections::{BTreeSet, HashMap, HashSet},
@@ -267,6 +270,19 @@ enum Commands {
         #[arg(long, short)]
         fix: bool,
     },
+    /// Profile each pattern in the given queries, outputting the time it takes them to compile.
+    Profile {
+        /// List of directories to profile
+        directories: Vec<PathBuf>,
+
+        /// String representing server's JSON configuration
+        #[arg(long, short)]
+        config: Option<String>,
+
+        /// Whether to broadly profile the entire query, rather than each pattern within the query.
+        #[arg(long, short)]
+        broad: bool,
+    },
 }
 
 /// Return the given config string, or read it from a config file if not given. This function can
@@ -308,6 +324,15 @@ async fn main() {
             let config_str = get_config_str(config);
             std::process::exit(lint_directories(&directories, config_str, fix).await)
         }
+        Some(Commands::Profile {
+            directories,
+            broad,
+            config,
+        }) => {
+            let config_str = get_config_str(config);
+            profile_directories(&directories, config_str, broad).await;
+            std::process::exit(0);
+        }
         None => {}
     }