diff --git a/CHANGELOG.md b/CHANGELOG.md index 171326f..6056dc7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Update Rust dependencies (#150) - Update default versions of Prometheus and Pushgateway (#150) - Add ability to scrape the metrics of `am`s own web server with `am start --scrape-self` (#153) +- `am list` now properly detects functions instrumented in Typescript using the `Autometrics` decorator (#152) +- `am instrument` is a new subcommand that can automatically add annotations to instrument a project (#152) + + it works in Go, Python, Typescript, and Rust projects ## [0.5.0] diff --git a/Cargo.lock b/Cargo.lock index 6c1df33..49d46a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -65,6 +65,7 @@ dependencies = [ "hex", "http", "humantime", + "ignore", "include_dir", "indicatif", "itertools", @@ -96,6 +97,8 @@ name = "am_list" version = "0.5.0" dependencies = [ "anyhow", + "crop", + "ignore", "itertools", "log", "pretty_assertions", @@ -352,6 +355,16 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bstr" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c79ad7fb2dd38f3dabd76b09c6a5a20c038fc0213ef1e9afd30eb777f120f019" +dependencies = [ + "memchr", + "serde", +] + [[package]] name = "bumpalo" version = "3.14.0" @@ -499,6 +512,15 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crop" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57e7b88784eea5a7895d70cde2535b36030ae387adde16a1d5c962c5ec5f273f" +dependencies = [ + "str_indices", +] + [[package]] name = "crossbeam-deque" version = "0.8.3" @@ -852,6 +874,19 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b" +[[package]] +name = "globset" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "759c97c1e17c55525b57192c06a267cda0ac5210b222d6b82189a2338fa1c13d" +dependencies = [ + "aho-corasick", + "bstr", + "fnv", + "log", + "regex", +] + [[package]] name = "h2" version = "0.3.21" @@ -1051,6 +1086,23 @@ dependencies = [ "unicode-normalization", ] +[[package]] +name = "ignore" +version = "0.4.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbe7873dab538a9a44ad79ede1faf5f30d49f9a5c883ddbab48bce81b64b7492" +dependencies = [ + "globset", + "lazy_static", + "log", + "memchr", + "regex", + "same-file", + "thread_local", + "walkdir", + "winapi-util", +] + [[package]] name = "include_dir" version = "0.7.3" @@ -2375,6 +2427,12 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +[[package]] +name = "str_indices" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8eeaedde8e50d8a331578c9fa9a288df146ce5e16173ad26ce82f6e263e2be4" + [[package]] name = "strsim" version = "0.10.0" diff --git a/am/Cargo.toml b/am/Cargo.toml index 171fbe8..0d82708 100644 --- a/am/Cargo.toml +++ b/am/Cargo.toml @@ -27,6 +27,7 @@ futures-util = { version = "0.3.28", features = ["io"] } hex = "0.4.3" http = "0.2.9" humantime = { workspace = true } +ignore = "0.4.20" include_dir = "0.7.3" indicatif = "0.17.5" itertools = "0.11.0" diff --git a/am/src/commands.rs b/am/src/commands.rs index b839276..bad9998 100644 --- a/am/src/commands.rs +++ b/am/src/commands.rs @@ -7,6 +7,7 @@ use tracing::info; mod explore; mod init; +mod instrument; mod list; mod proxy; pub mod start; @@ -64,6 +65,14 @@ pub enum SubCommands { /// List the functions in a project List(list::Arguments), + /// Instrument a project entirely. + /// + /// IMPORTANT: This will add code in your files! If you want to easily + /// undo the effects of this command, stage your work in progress (using `git add` or similar) + /// So that a command like `git restore .` can undo all unstaged changes, leaving your work + /// in progress alone. + Instrument(instrument::Arguments), + #[clap(hide = true)] MarkdownHelp, } @@ -86,6 +95,7 @@ pub async fn handle_command(app: Application, config: AmConfig, mp: MultiProgres } SubCommands::Update(args) => update::handle_command(args, mp).await, SubCommands::List(args) => list::handle_command(args), + SubCommands::Instrument(args) => instrument::handle_command(args), SubCommands::MarkdownHelp => { let disable_toc = true; clap_markdown::print_help_markdown::(Some(disable_toc)); diff --git a/am/src/commands/instrument.rs b/am/src/commands/instrument.rs new file mode 100644 index 0000000..12de212 --- /dev/null +++ b/am/src/commands/instrument.rs @@ -0,0 +1,159 @@ +use crate::interactive; +use am_list::Language; +use anyhow::Context; +use clap::{Args, Subcommand}; +use std::{ + path::{Path, PathBuf}, + process, +}; +use tracing::info; + +#[derive(Args)] +pub struct Arguments { + #[command(subcommand)] + command: Command, +} + +#[derive(Subcommand)] +enum Command { + /// Instrument functions in a single project, giving the language implementation + /// + /// IMPORTANT: This will add code in your files! If you want to easily + /// undo the effects of this command, stage your work in progress (using `git add` or similar) + /// So that a command like `git restore .` can undo all unstaged changes, leaving your work + /// in progress alone. + Single(SingleProject), + /// Instrument functions in all projects under the given directory, detecting languages on a best-effort basis. + /// + /// IMPORTANT: This will add code in your files! If you want to easily + /// undo the effects of this command, stage your work in progress (using `git add` or similar) + /// So that a command like `git restore .` can undo all unstaged changes, leaving your work + /// in progress alone. + All(AllProjects), +} + +#[derive(Args)] +struct SingleProject { + /// Language to detect autometrics functions for. Valid values are: + /// - 'rust' or 'rs' for Rust, + /// - 'go' for Golang, + /// - 'typescript', 'ts', 'javascript', or 'js' for Typescript/Javascript, + /// - 'python' or 'py' for Python. + #[arg(short, long, value_name = "LANGUAGE", verbatim_doc_comment)] + language: Language, + /// Root of the project to start the search on: + /// - For Rust projects it must be where the Cargo.toml lie, + /// - For Go projects it must be the root of the repository, + /// - For Python projects it must be the root of the library, + /// - For Typescript projects it must be where the package.json lie. + #[arg(value_name = "ROOT", verbatim_doc_comment)] + root: PathBuf, + /// A list of patterns to exclude from instrumentation. The patterns follow .gitignore rules, so + /// `--exclude "/vendor/"` will exclude all the vendor subdirectory only at the root, and adding + /// a pattern that starts with `!` will unignore a file or directory + #[arg(short, long, value_name = "PATTERNS")] + exclude: Vec, +} + +#[derive(Args)] +struct AllProjects { + /// Main directory to start the subprojects search on. am currently detects + /// Rust (Cargo.toml), Typescript (package.json), and Golang (go.mod) + /// projects. + #[arg(value_name = "ROOT")] + root: PathBuf, + /// A list of patterns to exclude from instrumentation. The patterns follow .gitignore rules, so + /// `--exclude "/vendor/"` will exclude all the vendor subdirectory only at the root, and adding + /// a pattern that starts with `!` will unignore a file or directory + #[arg(short, long, value_name = "PATTERNS")] + exclude: Vec, +} + +pub fn handle_command(args: Arguments) -> anyhow::Result<()> { + match args.command { + Command::Single(args) => handle_single_project(args), + Command::All(args) => handle_all_projects(args), + } +} + +fn folder_has_unstaged_changes(root: &Path) -> Option { + if cfg!(windows) { + // TODO: figure out the Windows story + return None; + } + + if cfg!(unix) { + // TODO: Figure out the non git story + let git_diff = process::Command::new("git") + .arg("-C") + .arg(root.as_os_str()) + .arg("diff") + .output(); + return match git_diff { + Ok(output) => Some(!output.stdout.is_empty()), + Err(_) => { + // We either don't have git, or root is not within a repository + None + } + }; + } + + None +} + +fn handle_all_projects(args: AllProjects) -> Result<(), anyhow::Error> { + let root = args + .root + .canonicalize() + .context("The path must be resolvable to an absolute path")?; + + if let Some(true) = folder_has_unstaged_changes(&root) { + let cont = interactive::confirm("The targeted root folder seems to have unstaged changes. `am` will also change files in this folder.\nDo you wish to continue?")?; + if !cont { + return Ok(()); + } + } + + info!("Instrumenting functions in {}:", root.display()); + + let mut exclude_patterns_builder = ignore::gitignore::GitignoreBuilder::new(&root); + for pattern in args.exclude { + exclude_patterns_builder.add_line(None, &pattern)?; + } + let exclude_patterns = exclude_patterns_builder.build()?; + + am_list::instrument_all_project_files(&root, &exclude_patterns)?; + + println!("If your project has Golang files, you need to run `go generate` now."); + + Ok(()) +} + +fn handle_single_project(args: SingleProject) -> Result<(), anyhow::Error> { + let root = args + .root + .canonicalize() + .context("The path must be resolvable to an absolute path")?; + + if let Some(true) = folder_has_unstaged_changes(&root) { + let cont = interactive::confirm("The targeted root folder seems to have unstaged changes. `am` will also change files in this folder.\nDo you wish to continue?")?; + if !cont { + return Ok(()); + } + } + info!("Instrumenting functions in {}:", root.display()); + + let mut exclude_patterns_builder = ignore::gitignore::GitignoreBuilder::new(&root); + for pattern in args.exclude { + exclude_patterns_builder.add_line(None, &pattern)?; + } + let exclude_patterns = exclude_patterns_builder.build()?; + + am_list::instrument_single_project_files(&root, args.language, &exclude_patterns)?; + + if args.language == Language::Go { + println!("You need to run `go generate` now."); + } + + Ok(()) +} diff --git a/am/src/commands/list.rs b/am/src/commands/list.rs index 6f4ba82..e9f7307 100644 --- a/am/src/commands/list.rs +++ b/am/src/commands/list.rs @@ -3,10 +3,6 @@ use clap::{Args, Subcommand}; use std::path::PathBuf; use tracing::info; -// TODO(gagbo): add an additional subcommand that makes use of am_list::find_roots to -// list all the functions under a given folder, by detecting the languages and all the -// subprojects included. - #[derive(Args)] pub struct Arguments { #[command(subcommand)] diff --git a/am_list/Cargo.toml b/am_list/Cargo.toml index 03a6845..099124c 100644 --- a/am_list/Cargo.toml +++ b/am_list/Cargo.toml @@ -11,6 +11,8 @@ license.workspace = true [dependencies] anyhow = "1.0.71" +crop = "0.4.0" +ignore = "0.4.20" itertools = "0.11.0" log = "0.4.18" rayon = "1.7.0" diff --git a/am_list/runtime/queries/typescript/all_functions.scm b/am_list/runtime/queries/typescript/all_functions.scm index bf33e26..816ebf8 100644 --- a/am_list/runtime/queries/typescript/all_functions.scm +++ b/am_list/runtime/queries/typescript/all_functions.scm @@ -7,7 +7,12 @@ (lexical_declaration (variable_declarator name: (identifier) @func.name - value: (arrow_function))) + value: (arrow_function) @func.value)) + +(lexical_declaration + (variable_declarator + name: (identifier) @func.name + value: (function) @func.value)) (class_declaration name: (type_identifier) @type.name diff --git a/am_list/runtime/queries/typescript/autometrics.scm b/am_list/runtime/queries/typescript/autometrics.scm index 54e10a6..0fe1399 100644 --- a/am_list/runtime/queries/typescript/autometrics.scm +++ b/am_list/runtime/queries/typescript/autometrics.scm @@ -57,3 +57,14 @@ (method_definition name: (property_identifier) @method.name)])) (#eq? @decorator.name "Autometrics")) + +((class_declaration + decorator: (decorator (call_expression + function: (identifier) @decorator.name)) + name: (type_identifier) @type.name + body: (class_body + [(method_signature + name: (property_identifier) @method.name) + (method_definition + name: (property_identifier) @method.name)])) + (#eq? @decorator.name "Autometrics")) diff --git a/am_list/src/go.rs b/am_list/src/go.rs index db3429b..9d23817 100644 --- a/am_list/src/go.rs +++ b/am_list/src/go.rs @@ -1,6 +1,7 @@ mod queries; -use crate::{FunctionInfo, ListAmFunctions, Result}; +use crate::{FunctionInfo, InstrumentFile, ListAmFunctions, Result}; +use log::debug; use queries::{AllFunctionsQuery, AmQuery}; use rayon::prelude::*; use std::{ @@ -34,17 +35,30 @@ impl Impl { .map(|s| s.ends_with(".go")) .unwrap_or(false) } -} -impl ListAmFunctions for Impl { - fn list_autometrics_functions(&mut self, project_root: &Path) -> Result> { + fn list_files( + project_root: &Path, + exclude_patterns: Option<&ignore::gitignore::Gitignore>, + ) -> Vec { const PREALLOCATED_ELEMS: usize = 100; - let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); - let walker = WalkDir::new(project_root).into_iter(); - let mut source_mod_pairs = Vec::with_capacity(PREALLOCATED_ELEMS); - source_mod_pairs.extend(walker.filter_entry(Self::is_valid).filter_map(|entry| { + let mut project_files = Vec::with_capacity(PREALLOCATED_ELEMS); + project_files.extend(walker.filter_entry(Self::is_valid).filter_map(|entry| { let entry = entry.ok()?; + + if let Some(pattern) = exclude_patterns { + let ignore_match = + pattern.matched_path_or_any_parents(entry.path(), entry.file_type().is_dir()); + if matches!(ignore_match, ignore::Match::Ignore(_)) { + debug!( + "The exclusion pattern got a match on {}: {:?}", + entry.path().display(), + ignore_match + ); + return None; + } + } + Some( entry .path() @@ -54,7 +68,19 @@ impl ListAmFunctions for Impl { ) })); - list.par_extend(source_mod_pairs.par_iter().filter_map(move |path| { + project_files + } +} + +impl ListAmFunctions for Impl { + fn list_autometrics_functions(&mut self, project_root: &Path) -> Result> { + const PREALLOCATED_ELEMS: usize = 100; + let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); + + let project_files = Self::list_files(project_root, None); + let query = AmQuery::try_new()?; + + list.par_extend(project_files.par_iter().filter_map(move |path| { let source = read_to_string(path).ok()?; let file_name = PathBuf::from(path) .strip_prefix(project_root) @@ -62,7 +88,6 @@ impl ListAmFunctions for Impl { .to_str() .expect("file_name is a valid path as it is part of `path`") .to_string(); - let query = AmQuery::try_new().ok()?; let names = query .list_function_names(&file_name, &source) .unwrap_or_default(); @@ -78,20 +103,10 @@ impl ListAmFunctions for Impl { const PREALLOCATED_ELEMS: usize = 100; let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); - let walker = WalkDir::new(project_root).into_iter(); - let mut source_mod_pairs = Vec::with_capacity(PREALLOCATED_ELEMS); - source_mod_pairs.extend(walker.filter_entry(Self::is_valid).filter_map(|entry| { - let entry = entry.ok()?; - Some( - entry - .path() - .to_str() - .map(ToString::to_string) - .unwrap_or_default(), - ) - })); + let project_files = Self::list_files(project_root, None); + let query = AllFunctionsQuery::try_new()?; - list.par_extend(source_mod_pairs.par_iter().filter_map(move |path| { + list.par_extend(project_files.par_iter().filter_map(move |path| { let source = read_to_string(path).ok()?; let file_name = PathBuf::from(path) .strip_prefix(project_root) @@ -99,7 +114,6 @@ impl ListAmFunctions for Impl { .to_str() .expect("file_name is a valid path as it is part of `path`") .to_string(); - let query = AllFunctionsQuery::try_new().ok()?; let names = query .list_function_names(&file_name, &source) .unwrap_or_default(); @@ -110,6 +124,82 @@ impl ListAmFunctions for Impl { result.extend(list.into_iter().flatten()); Ok(result) } + + fn list_autometrics_functions_in_single_file( + &mut self, + source_code: &str, + ) -> Result> { + let query = AmQuery::try_new()?; + query.list_function_names("", source_code) + } + + fn list_all_function_definitions_in_single_file( + &mut self, + source_code: &str, + ) -> Result> { + let query = AllFunctionsQuery::try_new()?; + query.list_function_names("", source_code) + } +} + +impl InstrumentFile for Impl { + fn instrument_source_code(&mut self, source: &str) -> Result { + let mut locations = self.list_all_functions_in_single_file(source)?; + locations.sort_by_key(|info| { + info.definition + .as_ref() + .map(|def| def.range.start.line) + .unwrap_or_default() + }); + + let has_am_directive = source + .lines() + .any(|line| line.starts_with("//go:generate autometrics")); + + let mut new_code = crop::Rope::from(source); + // Keeping track of inserted lines to update the byte offset to insert code to, + // only works if the locations list is sorted from top to bottom + let mut inserted_lines = 0; + + if !has_am_directive { + new_code.insert(0, "//go:generate autometrics --otel\n"); + inserted_lines += 1; + } + + for function_info in locations { + if function_info.definition.is_none() || function_info.instrumentation.is_some() { + continue; + } + + let def_line = function_info.definition.as_ref().unwrap().range.start.line; + let byte_offset = new_code.byte_of_line(inserted_lines + def_line); + new_code.insert(byte_offset, "//autometrics:inst\n"); + inserted_lines += 1; + } + + Ok(new_code.to_string()) + } + + fn instrument_project( + &mut self, + project_root: &Path, + exclude_patterns: Option<&ignore::gitignore::Gitignore>, + ) -> Result<()> { + let sources_modules = Self::list_files(project_root, exclude_patterns); + debug!("Found sources {sources_modules:?}"); + + for path in sources_modules { + if std::fs::metadata(&path)?.is_dir() { + continue; + } + debug!("Instrumenting {path}"); + let old_source = read_to_string(&path)?; + let new_source = self.instrument_source_code(&old_source)?; + std::fs::write(path, new_source)?; + } + + Ok(()) + } } #[cfg(test)] diff --git a/am_list/src/go/tests.rs b/am_list/src/go/tests.rs index e79bdb0..14413b3 100644 --- a/am_list/src/go/tests.rs +++ b/am_list/src/go/tests.rs @@ -256,3 +256,66 @@ fn detect_method_pointer_receiver() { assert_eq!(all_list.len(), 1); assert_eq!(all_list[0], the_one_all_functions); } + +#[test] +fn instrument_method() { + let source = r#" + package lambda + + func (s Server) the_one() { + return nil + } + "#; + + let expected = r#"//go:generate autometrics --otel + + package lambda + +//autometrics:inst + func (s Server) the_one() { + return nil + } + "#; + + let mut implementation = Impl {}; + let actual = implementation.instrument_source_code(source).unwrap(); + assert_eq!(&actual, expected); +} + +#[test] +fn instrument_multiple() { + let source = r#"package beta + + func not_the_one() { + } + + //autometrics:doc + func sandwiched_function() { + return nil + } + + func not_that_one_either() { + } + "#; + + let expected = r#"//go:generate autometrics --otel +package beta + +//autometrics:inst + func not_the_one() { + } + + //autometrics:doc + func sandwiched_function() { + return nil + } + +//autometrics:inst + func not_that_one_either() { + } + "#; + + let mut implementation = Impl {}; + let actual = implementation.instrument_source_code(source).unwrap(); + assert_eq!(&actual, expected); +} diff --git a/am_list/src/lib.rs b/am_list/src/lib.rs index 444b1cf..6994d8a 100644 --- a/am_list/src/lib.rs +++ b/am_list/src/lib.rs @@ -5,12 +5,12 @@ pub mod rust; pub mod typescript; use log::info; +use rayon::prelude::*; pub use roots::find_project_roots; - use serde::{Deserialize, Serialize}; use std::{ collections::{BTreeMap, HashMap}, - fmt::Display, + fmt::{Display, Write}, path::{Path, PathBuf}, str::FromStr, }; @@ -167,6 +167,62 @@ pub trait ListAmFunctions { } Ok(info_set.into_values().collect()) } + + /// List all the autometricized functions in the given source code. + fn list_autometrics_functions_in_single_file( + &mut self, + source_code: &str, + ) -> Result>; + + /// List all the functions defined in the given source code. + fn list_all_function_definitions_in_single_file( + &mut self, + source_code: &str, + ) -> Result>; + + /// List all the functions in the given source code, instrumented or just defined. + /// + /// This is guaranteed to return the most complete set of information + fn list_all_functions_in_single_file( + &mut self, + source_code: &str, + ) -> Result> { + let am_functions = self.list_autometrics_functions_in_single_file(source_code)?; + let all_function_definitions = + self.list_all_function_definitions_in_single_file(source_code)?; + let mut info_set: HashMap = am_functions + .into_iter() + .map(|full_info| (full_info.id.clone(), full_info)) + .collect(); + + // Only the definition field is expected to differ + // between am_functions and all_function_definitions + for function in all_function_definitions { + info_set + .entry(function.id.clone()) + .and_modify(|info| info.definition = function.definition.clone()) + .or_insert(function); + } + Ok(info_set.into_values().collect()) + } +} + +/// Instrument a file, adding autometrics annotations as necessary. +/// +/// Each language is responsible to reuse its queries/create additonal queries to add the +/// necessary code and produce a new version of the file that has _all_ functions instrumented. +/// +/// The invariant to maintain here is that after being done with a file, all functions defined +/// in the file should be instrumented. +pub trait InstrumentFile { + /// Instrument all functions in the file + fn instrument_source_code(&mut self, source: &str) -> Result; + /// Instrument all functions under the given project. + fn instrument_project( + &mut self, + project_root: &Path, + exclude_patterns: Option<&ignore::gitignore::Gitignore>, + ) -> Result<()>; } pub type Result = std::result::Result; @@ -192,6 +248,16 @@ pub enum AmlError { /// Issue when trying to extract a path to a project #[error("Invalid path to project")] InvalidPath, + /// Issue when trying to interact with the filesystem + #[error("I/O error when interacting with the filesystem")] + IO(#[from] std::io::Error), + /// Collection of issues when working with multiple projects + #[error("Some projects had errors: {}", + .0.iter().fold(String::new(), |mut output, (path, err)| { + let _ = write!(output, "in {}: {}; ", path.display(), err); + output + }))] + Projects(Vec<(PathBuf, Self)>), } /// Languages supported by `am_list`. @@ -246,19 +312,41 @@ pub fn list_all_project_functions( let projects = find_project_roots(root)?; let mut res: BTreeMap)> = BTreeMap::new(); - // TODO: try to parallelize this loop if possible - for (path, language) in projects.iter() { - info!( - "Listing functions in {} (Language: {})", - path.display(), - language - ); - let project_fns = list_single_project_functions(path, *language, true)?; - - res.entry(path.to_path_buf()) - .or_insert_with(|| (*language, Vec::new())) - .1 - .extend(project_fns); + let per_project_results = projects + .into_par_iter() + .map(|(path, language)| { + info!( + "Listing functions in {} (Language: {})", + path.display(), + language + ); + match list_single_project_functions(&path, language, true) { + Ok(project_fns) => Ok((path.clone(), language, project_fns)), + Err(err) => Err((path.clone(), err)), + } + }) + .collect::>(); + + // This filter_map collects the errors into a single one, and has the side-effect of + // filling the Ok result in case no error happened along the way. + // This is done in a second step rather than within the par_iter because of the mutation + // that needs to happen in the resulting map. + let errors = per_project_results + .into_iter() + .filter_map(|result| match result { + Ok((path, language, project_fns)) => { + res.entry(path.to_path_buf()) + .or_insert_with(|| (language, Vec::new())) + .1 + .extend(project_fns); + None + } + Err(err) => Some(err), + }) + .collect::>(); + + if !errors.is_empty() { + return Err(AmlError::Projects(errors)); } Ok(res) @@ -283,3 +371,49 @@ pub fn list_single_project_functions( res.sort(); Ok(res) } + +pub fn instrument_all_project_files( + root: &Path, + exclude_patterns: &ignore::gitignore::Gitignore, +) -> Result<()> { + let projects = find_project_roots(root)?; + + let errors = projects + .into_par_iter() + .filter_map(|(path, language)| { + info!( + "Instrumenting functions in {} (Language: {})", + path.display(), + language + ); + + if let Some(err) = + instrument_single_project_files(&path, language, exclude_patterns).err() + { + return Some((path.clone(), err)); + } + + None + }) + .collect::>(); + + if !errors.is_empty() { + return Err(AmlError::Projects(errors)); + } + + Ok(()) +} + +pub fn instrument_single_project_files( + root: &Path, + language: Language, + exclude_patterns: &ignore::gitignore::Gitignore, +) -> Result<()> { + let mut implementor: Box = match language { + Language::Rust => Box::new(crate::rust::Impl {}), + Language::Go => Box::new(crate::go::Impl {}), + Language::Typescript => Box::new(crate::typescript::Impl {}), + Language::Python => Box::new(crate::python::Impl {}), + }; + implementor.instrument_project(root, Some(exclude_patterns)) +} diff --git a/am_list/src/python.rs b/am_list/src/python.rs index 0922c3c..99a3a3c 100644 --- a/am_list/src/python.rs +++ b/am_list/src/python.rs @@ -1,6 +1,7 @@ mod queries; -use crate::{FunctionInfo, ListAmFunctions, Result}; +use crate::{FunctionInfo, InstrumentFile, ListAmFunctions, Result}; +use log::debug; use queries::{AllFunctionsQuery, AmImportQuery, AmQuery}; use rayon::prelude::*; use std::{ @@ -33,21 +34,30 @@ impl Impl { .extension() .map_or(false, |ext| ext == "py" || ext == "py3") } -} -impl ListAmFunctions for Impl { - fn list_autometrics_functions(&mut self, project_root: &Path) -> Result> { + fn list_files( + project_root: &Path, + exclude_patterns: Option<&ignore::gitignore::Gitignore>, + ) -> Vec { const PREALLOCATED_ELEMS: usize = 100; - let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); - let root_name = project_root - .file_name() - .map(|s| s.to_str().unwrap_or_default()) - .unwrap_or(""); - let walker = WalkDir::new(project_root).into_iter(); - let mut source_mod_pairs = Vec::with_capacity(PREALLOCATED_ELEMS); - source_mod_pairs.extend(walker.filter_entry(Self::is_valid).filter_map(|entry| { + let mut project_files = Vec::with_capacity(PREALLOCATED_ELEMS); + project_files.extend(walker.filter_entry(Self::is_valid).filter_map(|entry| { let entry = entry.ok()?; + + if let Some(pattern) = exclude_patterns { + let ignore_match = + pattern.matched_path_or_any_parents(entry.path(), entry.file_type().is_dir()); + if matches!(ignore_match, ignore::Match::Ignore(_)) { + debug!( + "The exclusion pattern got a match on {}: {:?}", + entry.path().display(), + ignore_match + ); + return None; + } + } + Some( entry .path() @@ -57,7 +67,21 @@ impl ListAmFunctions for Impl { ) })); - list.par_extend(source_mod_pairs.par_iter().filter_map(move |path| { + project_files + } +} + +impl ListAmFunctions for Impl { + fn list_autometrics_functions(&mut self, project_root: &Path) -> Result> { + const PREALLOCATED_ELEMS: usize = 100; + let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); + let root_name = project_root + .file_name() + .map(|s| s.to_str().unwrap_or_default()) + .unwrap_or(""); + let project_files = Self::list_files(project_root, None); + + list.par_extend(project_files.par_iter().filter_map(move |path| { let relative_module_name = Path::new(path) .strip_prefix(project_root) .ok()? @@ -94,20 +118,9 @@ impl ListAmFunctions for Impl { .map(|s| s.to_str().unwrap_or_default()) .unwrap_or(""); - let walker = WalkDir::new(project_root).into_iter(); - let mut source_mod_pairs = Vec::with_capacity(PREALLOCATED_ELEMS); - source_mod_pairs.extend(walker.filter_entry(Self::is_valid).filter_map(|entry| { - let entry = entry.ok()?; - Some( - entry - .path() - .to_str() - .map(ToString::to_string) - .unwrap_or_default(), - ) - })); + let project_files = Self::list_files(project_root, None); - list.par_extend(source_mod_pairs.par_iter().filter_map(move |path| { + list.par_extend(project_files.par_iter().filter_map(move |path| { let relative_module_name = Path::new(path) .strip_prefix(project_root) .ok()? @@ -133,6 +146,98 @@ impl ListAmFunctions for Impl { result.extend(list.into_iter().flatten()); Ok(result) } + + fn list_autometrics_functions_in_single_file( + &mut self, + source_code: &str, + ) -> Result> { + let import_query = AmImportQuery::try_new()?; + let decorator_name = import_query.get_decorator_name(source_code).ok(); + if decorator_name.is_none() { + return Ok(Vec::new()); + } + let query = AmQuery::try_new(decorator_name.as_ref().unwrap())?; + query.list_function_names("", source_code, "") + } + + fn list_all_function_definitions_in_single_file( + &mut self, + source_code: &str, + ) -> Result> { + let query = AllFunctionsQuery::try_new()?; + query.list_function_names("", source_code, "") + } +} + +impl InstrumentFile for Impl { + fn instrument_source_code(&mut self, source: &str) -> Result { + const DEF_LEN: usize = "def ".len(); + + let mut locations = self.list_all_functions_in_single_file(source)?; + locations.sort_by_key(|info| { + info.definition + .as_ref() + .map(|def| def.range.start.line) + .unwrap_or_default() + }); + + let has_am_directive = source + .lines() + .any(|line| line.contains("from autometrics import autometrics")); + + let mut new_code = crop::Rope::from(source); + // Keeping track of inserted lines to update the byte offset to insert code to, + // only works if the locations list is sorted from top to bottom + let mut inserted_lines = 0; + + if !has_am_directive { + new_code.insert(0, "from autometrics import autometrics\n"); + inserted_lines += 1; + } + + for function_info in locations { + if function_info.definition.is_none() || function_info.instrumentation.is_some() { + continue; + } + + let def_line = function_info.definition.as_ref().unwrap().range.start.line; + let def_col = function_info + .definition + .unwrap() + .range + .start + .column + .saturating_sub(DEF_LEN); + let byte_offset = new_code.byte_of_line(inserted_lines + def_line); + new_code.insert( + byte_offset, + format!("{}@autometrics\n", " ".repeat(def_col)), + ); + inserted_lines += 1; + } + + Ok(new_code.to_string()) + } + + fn instrument_project( + &mut self, + project_root: &Path, + exclude_patterns: Option<&ignore::gitignore::Gitignore>, + ) -> Result<()> { + let sources_modules = Self::list_files(project_root, exclude_patterns); + + for path in sources_modules { + if std::fs::metadata(&path)?.is_dir() { + continue; + } + debug!("Instrumenting {path}"); + let old_source = read_to_string(&path)?; + let new_source = self.instrument_source_code(&old_source)?; + std::fs::write(path, new_source)?; + } + + Ok(()) + } } #[cfg(test)] diff --git a/am_list/src/python/tests.rs b/am_list/src/python/tests.rs index 7c7a030..3995c70 100644 --- a/am_list/src/python/tests.rs +++ b/am_list/src/python/tests.rs @@ -198,3 +198,68 @@ fn detect_nested() { assert!(all_list.contains(&the_one)); assert!(all_list.contains(&the_two)); } + +#[test] +fn instrument_nested() { + let source = r#" +import os + + def the_one(): + def the_two(): + return 'wake up, Neo' + return the_two() + "#; + + let expected = r#"from autometrics import autometrics + +import os + + @autometrics + def the_one(): + @autometrics + def the_two(): + return 'wake up, Neo' + return the_two() + "#; + + let mut implementation = Impl {}; + let actual = implementation.instrument_source_code(source).unwrap(); + assert_eq!(&actual, expected); +} + +#[test] +fn instrument_multiple() { + let source = r#" + from autometrics import autometrics + + def the_one(): + return 'wake up, Neo' + + @autometrics + def the_two(): + return 'wake up, Neo' + + def the_three(): + return 'wake up, Neo' + "#; + + let expected = r#" + from autometrics import autometrics + + @autometrics + def the_one(): + return 'wake up, Neo' + + @autometrics + def the_two(): + return 'wake up, Neo' + + @autometrics + def the_three(): + return 'wake up, Neo' + "#; + + let mut implementation = Impl {}; + let actual = implementation.instrument_source_code(source).unwrap(); + assert_eq!(&actual, expected); +} diff --git a/am_list/src/rust.rs b/am_list/src/rust.rs index 560572e..405c879 100644 --- a/am_list/src/rust.rs +++ b/am_list/src/rust.rs @@ -1,7 +1,8 @@ mod queries; use self::queries::{AllFunctionsQuery, AmQuery}; -use crate::{FunctionInfo, ListAmFunctions, Result}; +use crate::{FunctionInfo, InstrumentFile, ListAmFunctions, Result}; +use log::debug; use rayon::prelude::*; use std::{ collections::{HashSet, VecDeque}, @@ -89,18 +90,31 @@ impl Impl { itertools::intersperse(mod_name_elements, "::".to_string()).collect() } -} -impl ListAmFunctions for Impl { - fn list_autometrics_functions(&mut self, project_root: &Path) -> Result> { + fn list_files_and_modules( + project_root: &Path, + exclude_patterns: Option<&ignore::gitignore::Gitignore>, + ) -> Vec<(String, String)> { const PREALLOCATED_ELEMS: usize = 100; - let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); let walker = WalkDir::new(project_root).into_iter(); let mut source_mod_pairs = Vec::with_capacity(PREALLOCATED_ELEMS); - let query = AmQuery::try_new()?; source_mod_pairs.extend(walker.filter_entry(Self::is_valid).filter_map(|entry| { let entry = entry.ok()?; + + if let Some(pattern) = exclude_patterns { + let ignore_match = + pattern.matched_path_or_any_parents(entry.path(), entry.file_type().is_dir()); + if matches!(ignore_match, ignore::Match::Ignore(_)) { + debug!( + "The exclusion pattern got a match on {}: {:?}", + entry.path().display(), + ignore_match + ); + return None; + } + } + let module = Self::fully_qualified_module_name(&entry); Some(( entry @@ -112,6 +126,17 @@ impl ListAmFunctions for Impl { )) })); + source_mod_pairs + } +} + +impl ListAmFunctions for Impl { + fn list_autometrics_functions(&mut self, project_root: &Path) -> Result> { + const PREALLOCATED_ELEMS: usize = 100; + let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); + let query = AmQuery::try_new()?; + let source_mod_pairs = Self::list_files_and_modules(project_root, None); + list.par_extend( source_mod_pairs .par_iter() @@ -138,22 +163,8 @@ impl ListAmFunctions for Impl { fn list_all_function_definitions(&mut self, project_root: &Path) -> Result> { const PREALLOCATED_ELEMS: usize = 400; let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); - - let walker = WalkDir::new(project_root).into_iter(); - let mut source_mod_pairs = Vec::with_capacity(PREALLOCATED_ELEMS); + let source_mod_pairs = Self::list_files_and_modules(project_root, None); let query = AllFunctionsQuery::try_new()?; - source_mod_pairs.extend(walker.filter_entry(Self::is_valid).filter_map(|entry| { - let entry = entry.ok()?; - let module = Self::fully_qualified_module_name(&entry); - Some(( - entry - .path() - .to_str() - .map(ToString::to_string) - .unwrap_or_default(), - module, - )) - })); list.par_extend( source_mod_pairs @@ -177,6 +188,72 @@ impl ListAmFunctions for Impl { result.extend(list.into_iter().flatten()); Ok(result) } + + fn list_autometrics_functions_in_single_file( + &mut self, + source_code: &str, + ) -> Result> { + let query = AmQuery::try_new()?; + query.list_function_names("", String::new(), source_code) + } + + fn list_all_function_definitions_in_single_file( + &mut self, + source_code: &str, + ) -> Result> { + let query = AllFunctionsQuery::try_new()?; + query.list_function_names("", String::new(), source_code) + } +} + +impl InstrumentFile for Impl { + fn instrument_source_code(&mut self, source: &str) -> Result { + let mut locations = self.list_all_functions_in_single_file(source)?; + locations.sort_by_key(|info| { + info.definition + .as_ref() + .map(|def| def.range.start.line) + .unwrap_or_default() + }); + + let mut new_code = crop::Rope::from(source); + // Keeping track of inserted lines to update the byte offset to insert code to, + // only works if the locations list is sorted from top to bottom + let mut inserted_lines = 0; + + for function_info in locations { + if function_info.definition.is_none() || function_info.instrumentation.is_some() { + continue; + } + + let def_line = function_info.definition.as_ref().unwrap().range.start.line; + let byte_offset = new_code.byte_of_line(inserted_lines + def_line); + new_code.insert(byte_offset, "#[autometrics::autometrics]\n"); + inserted_lines += 1; + } + + Ok(new_code.to_string()) + } + + fn instrument_project( + &mut self, + project_root: &Path, + exclude_patterns: Option<&ignore::gitignore::Gitignore>, + ) -> Result<()> { + let sources_modules = Self::list_files_and_modules(project_root, exclude_patterns); + + for (path, _module) in sources_modules { + if std::fs::metadata(&path)?.is_dir() { + continue; + } + debug!("Instrumenting {path}"); + let old_source = read_to_string(&path)?; + let new_source = self.instrument_source_code(&old_source)?; + std::fs::write(path, new_source)?; + } + + Ok(()) + } } #[cfg(test)] diff --git a/am_list/src/rust/tests.rs b/am_list/src/rust/tests.rs index edad4eb..5941521 100644 --- a/am_list/src/rust/tests.rs +++ b/am_list/src/rust/tests.rs @@ -482,3 +482,82 @@ fn detect_partially_annotated_impl_block() { "Expecting the list to contain {dummy:?}\nComplete list is {all:?}" ); } + +#[test] +fn instrument_method() { + let source = r#" + impl Foo for Bar { + fn bar(&mut self) -> Self { + todo!() + } + } + "#; + + let expected = r#" + impl Foo for Bar { +#[autometrics::autometrics] + fn bar(&mut self) -> Self { + todo!() + } + } + "#; + + let mut implementation = Impl {}; + let actual = implementation.instrument_source_code(source).unwrap(); + assert_eq!(&actual, expected); +} + +#[test] +fn instrument_multiple() { + let source = r#"use autometrics::autometrics; + + fn not_the_one() { + todo!() + } + + #[autometrics] + fn sandwiched_function() { + todo!() + } + + impl Stuff { + fn not_that_one_either(&self) { + } + } + + #[autometrics] + impl AMStuff { + fn not_that_one_either(&self) { + } + } + "#; + + let expected = r#"use autometrics::autometrics; + +#[autometrics::autometrics] + fn not_the_one() { + todo!() + } + + #[autometrics] + fn sandwiched_function() { + todo!() + } + + impl Stuff { +#[autometrics::autometrics] + fn not_that_one_either(&self) { + } + } + + #[autometrics] + impl AMStuff { + fn not_that_one_either(&self) { + } + } + "#; + + let mut implementation = Impl {}; + let actual = implementation.instrument_source_code(source).unwrap(); + assert_eq!(&actual, expected); +} diff --git a/am_list/src/typescript.rs b/am_list/src/typescript.rs index dc1a525..f73b8f8 100644 --- a/am_list/src/typescript.rs +++ b/am_list/src/typescript.rs @@ -1,7 +1,8 @@ mod imports; mod queries; -use crate::{FunctionInfo, ListAmFunctions, Result}; +use crate::{FunctionInfo, InstrumentFile, ListAmFunctions, Result}; +use log::{debug, trace}; use rayon::prelude::*; use std::{ collections::{HashSet, VecDeque}, @@ -10,7 +11,7 @@ use std::{ }; use walkdir::{DirEntry, WalkDir}; -use self::queries::{AllFunctionsQuery, AmQuery}; +use self::queries::{AllFunctionsQuery, AmQuery, TypescriptFunctionInfo}; /// Implementation of the Typescript support for listing autometricized functions. #[derive(Clone, Copy, Debug, Default)] @@ -68,26 +69,57 @@ impl Impl { } itertools::intersperse(mod_name_elements, "/".to_string()).collect() } -} -impl ListAmFunctions for Impl { - fn list_autometrics_functions(&mut self, project_root: &Path) -> Result> { - const PREALLOCATED_ELEMS: usize = 100; - let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); + fn ts_function_definitions_in_single_file( + &mut self, + source_code: &str, + ) -> Result> { + let query = AllFunctionsQuery::try_new()?; + query.list_function_names("", "", source_code) + } + fn list_files_and_modules( + project_root: &Path, + exclude_patterns: Option<&ignore::gitignore::Gitignore>, + ) -> Vec<(PathBuf, String)> { + const PREALLOCATED_ELEMS: usize = 100; let walker = WalkDir::new(project_root).into_iter(); let mut source_mod_pairs = Vec::with_capacity(PREALLOCATED_ELEMS); source_mod_pairs.extend(walker.filter_entry(Self::is_valid).filter_map(|entry| { let entry = entry.ok()?; + + if let Some(pattern) = exclude_patterns { + let ignore_match = + pattern.matched_path_or_any_parents(entry.path(), entry.file_type().is_dir()); + if matches!(ignore_match, ignore::Match::Ignore(_)) { + debug!( + "The exclusion pattern got a match on {}: {:?}", + entry.path().display(), + ignore_match + ); + return None; + } + } + let module = Self::qualified_module_name(&entry); Some((entry.path().to_path_buf(), module)) })); + source_mod_pairs + } +} + +impl ListAmFunctions for Impl { + fn list_autometrics_functions(&mut self, project_root: &Path) -> Result> { + const PREALLOCATED_ELEMS: usize = 100; + let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); + let source_mod_pairs = Self::list_files_and_modules(project_root, None); + let query = AmQuery::try_new()?; + list.par_extend( source_mod_pairs .par_iter() .filter_map(move |(path, module)| { - let query = AmQuery::try_new().ok()?; let source = read_to_string(path).ok()?; let file_name = PathBuf::from(path) .strip_prefix(project_root) @@ -110,28 +142,14 @@ impl ListAmFunctions for Impl { fn list_all_function_definitions(&mut self, project_root: &Path) -> Result> { const PREALLOCATED_ELEMS: usize = 100; let mut list = HashSet::with_capacity(PREALLOCATED_ELEMS); - - let walker = WalkDir::new(project_root).into_iter(); - let mut source_mod_pairs = Vec::with_capacity(PREALLOCATED_ELEMS); - source_mod_pairs.extend(walker.filter_entry(Self::is_valid).filter_map(|entry| { - let entry = entry.ok()?; - let module = Self::qualified_module_name(&entry); - Some(( - entry - .path() - .to_str() - .map(ToString::to_string) - .unwrap_or_default(), - module, - )) - })); + let source_mod_pairs = Self::list_files_and_modules(project_root, None); + let query = AllFunctionsQuery::try_new()?; list.par_extend( source_mod_pairs .par_iter() .filter_map(move |(path, module)| { let source = read_to_string(path).ok()?; - let query = AllFunctionsQuery::try_new().ok()?; let file_name = PathBuf::from(path) .strip_prefix(project_root) .expect("path comes from a project_root WalkDir") @@ -141,14 +159,153 @@ impl ListAmFunctions for Impl { let names = query .list_function_names(&file_name, module, &source) .ok()?; - Some(names.into_iter().collect::>()) + Some( + names + .into_iter() + .map(|info| info.inner_info) + .collect::>(), + ) }), ); let mut result = Vec::with_capacity(PREALLOCATED_ELEMS); result.extend(list.into_iter().flatten()); + trace!("Item list: {result:?}"); Ok(result) } + + fn list_autometrics_functions_in_single_file( + &mut self, + source_code: &str, + ) -> Result> { + let query = AmQuery::try_new()?; + query.list_function_names("", "", source_code, None) + } + + fn list_all_function_definitions_in_single_file( + &mut self, + source_code: &str, + ) -> Result> { + Ok(self + .ts_function_definitions_in_single_file(source_code)? + .into_iter() + .map(Into::into) + .collect()) + } +} + +impl InstrumentFile for Impl { + fn instrument_source_code(&mut self, source: &str) -> Result { + let mut locations = self.list_all_functions_in_single_file(source)?; + locations.sort_by_key(|info| { + info.definition + .as_ref() + .map(|def| def.range.start.line) + .unwrap_or_default() + }); + + let mut ts_specific_locations = self.ts_function_definitions_in_single_file(source)?; + ts_specific_locations.sort_by_key(|info| { + info.inner_info + .definition + .as_ref() + .map(|def| def.range.start.line) + .unwrap_or_default() + }); + + let has_am_directive = source.lines().any(|line| { + line.contains("import { autometrics } from") + || line.contains("import { Autometrics } from") + || line.contains("import { autometrics, Autometrics } from") + }); + let mut placeholder_offset_range = None; + let mut needs_decorator_import = false; + let mut needs_wrapper_import = false; + + let mut new_code = crop::Rope::from(source); + // Keeping track of inserted lines to update the byte offset to insert code to, + // only works if the locations list is sorted from top to bottom + let mut inserted_lines = 0; + + if !has_am_directive { + new_code.insert( + 0, + "import { placeholder } from '@autometrics/autometrics';\n", + ); + inserted_lines += 1; + placeholder_offset_range = Some("import { ".len().."import { placeholder".len()); + } + + for function_info in locations { + if function_info.definition.is_none() || function_info.instrumentation.is_some() { + continue; + } + + let ts_loc = ts_specific_locations + .iter() + .find_map(|info| { + if info.inner_info.id == function_info.id { + Some(info.function_rvalue_range.clone()) + } else { + None + } + }) + .flatten(); + + match ts_loc { + Some(rvalue_range) => { + let start_byte_offset = new_code + .byte_of_line(inserted_lines + rvalue_range.start.line) + + rvalue_range.start.column; + new_code.insert(start_byte_offset, "autometrics("); + let end_byte_offset = new_code + .byte_of_line(inserted_lines + rvalue_range.end.line) + + rvalue_range.end.column; + new_code.insert(end_byte_offset, ")"); + needs_wrapper_import = true; + } + None => { + let def_line = function_info.definition.as_ref().unwrap().range.start.line; + let byte_offset = new_code.byte_of_line(inserted_lines + def_line); + new_code.insert(byte_offset, "@Autometrics()\n"); + inserted_lines += 1; + needs_decorator_import = true; + } + } + } + + if let Some(range) = placeholder_offset_range { + let imports = match (needs_wrapper_import, needs_decorator_import) { + (true, true) => "autometrics, Autometrics", + (true, false) => "autometrics", + (false, true) => "Autometrics", + (false, false) => "", + }; + new_code.replace(range, imports); + } + + Ok(new_code.to_string()) + } + + fn instrument_project( + &mut self, + project_root: &Path, + exclude_patterns: Option<&ignore::gitignore::Gitignore>, + ) -> Result<()> { + let sources_modules = Self::list_files_and_modules(project_root, exclude_patterns); + + for (path, _module) in sources_modules { + if std::fs::metadata(&path)?.is_dir() { + continue; + } + debug!("Instrumenting {}", path.display()); + let old_source = read_to_string(&path)?; + let new_source = self.instrument_source_code(&old_source)?; + std::fs::write(path, new_source)?; + } + + Ok(()) + } } #[cfg(test)] diff --git a/am_list/src/typescript/queries.rs b/am_list/src/typescript/queries.rs index 31cc927..b899005 100644 --- a/am_list/src/typescript/queries.rs +++ b/am_list/src/typescript/queries.rs @@ -1,13 +1,16 @@ use std::path::Path; use log::warn; +use serde::{Deserialize, Serialize}; use tree_sitter::{Parser, Query}; use tree_sitter_typescript::language_typescript as language; -use crate::{AmlError, FunctionInfo, Location, Result, FUNC_NAME_CAPTURE}; +use crate::{AmlError, FunctionInfo, Location, Range, Result, FUNC_NAME_CAPTURE}; use super::imports::{Identifier, ImportsMap, Source}; +const FUNCTION_RVALUE_CAPTURE: &str = "func.value"; + const TYPE_NAME_CAPTURE: &str = "type.name"; const METHOD_NAME_CAPTURE: &str = "method.name"; const WRAPPER_DIRECT_NAME_CAPTURE: &str = "wrapperdirect.name"; @@ -31,12 +34,27 @@ pub(super) struct AllFunctionsQuery { query: Query, /// Index of the capture for a function name. func_name_idx: u32, + /// Index of the capture for a function value, when the function is + /// used in an assignment. + func_value_idx: u32, /// Index of the capture for the name of a class that is defined in file. type_name_idx: u32, /// Index of the capture for the contents of a method that is defined in file. method_name_idx: u32, } +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct TypescriptFunctionInfo { + pub inner_info: FunctionInfo, + pub function_rvalue_range: Option, +} + +impl From for FunctionInfo { + fn from(value: TypescriptFunctionInfo) -> Self { + value.inner_info + } +} + impl AllFunctionsQuery { pub fn try_new() -> Result { let query = Query::new( @@ -46,6 +64,9 @@ impl AllFunctionsQuery { let func_name_idx = query .capture_index_for_name(FUNC_NAME_CAPTURE) .ok_or_else(|| AmlError::MissingNamedCapture(FUNC_NAME_CAPTURE.to_string()))?; + let func_value_idx = query + .capture_index_for_name(FUNCTION_RVALUE_CAPTURE) + .ok_or_else(|| AmlError::MissingNamedCapture(FUNCTION_RVALUE_CAPTURE.to_string()))?; let type_name_idx = query .capture_index_for_name(TYPE_NAME_CAPTURE) .ok_or_else(|| AmlError::MissingNamedCapture(TYPE_NAME_CAPTURE.to_string()))?; @@ -56,6 +77,7 @@ impl AllFunctionsQuery { Ok(Self { query, func_name_idx, + func_value_idx, type_name_idx, method_name_idx, }) @@ -66,20 +88,25 @@ impl AllFunctionsQuery { file_name: &str, module_name: &str, source: &str, - ) -> Result> { + ) -> Result> { let mut parser = new_parser()?; let parsed_source = parser.parse(source, None).ok_or(AmlError::Parsing)?; let mut cursor = tree_sitter::QueryCursor::new(); let functions = cursor .matches(&self.query, parsed_source.root_node(), source.as_bytes()) - .filter_map(|capture| -> Option { + .filter_map(|capture| -> Option { let func_name_node = capture.nodes_for_capture_index(self.func_name_idx).next(); + let func_value_node = capture.nodes_for_capture_index(self.func_value_idx).next(); let method_name_node = capture.nodes_for_capture_index(self.method_name_idx).next(); let type_name_node = capture.nodes_for_capture_index(self.type_name_idx).next(); match ( // Test for bare function capture func_name_node .map(|node| node.utf8_text(source.as_bytes()).map(ToString::to_string)), + // Test for function assignment capture + func_value_node.map(|node| { + Range::from((node.range().start_point, node.range().end_point)) + }), // Test for method name capture method_name_node .map(|node| node.utf8_text(source.as_bytes()).map(ToString::to_string)), @@ -87,7 +114,7 @@ impl AllFunctionsQuery { type_name_node .map(|node| node.utf8_text(source.as_bytes()).map(ToString::to_string)), ) { - (Some(Ok(bare_function_name)), _, _) => { + (Some(Ok(bare_function_name)), function_rvalue_range, _, _) => { let start = func_name_node .expect("just extracted a name from the node") .start_position(); @@ -96,13 +123,16 @@ impl AllFunctionsQuery { .end_position(); let instrumentation = None; let definition = Some(Location::from((file_name, start, end))); - Some(FunctionInfo { - id: (module_name, bare_function_name).into(), - instrumentation, - definition, + Some(TypescriptFunctionInfo { + inner_info: FunctionInfo { + id: (module_name, bare_function_name).into(), + instrumentation, + definition, + }, + function_rvalue_range, }) } - (_, Some(Ok(method_name)), Some(Ok(class_name))) => { + (_, _, Some(Ok(method_name)), Some(Ok(class_name))) => { let start = method_name_node .expect("just extracted a name from the node") .start_position(); @@ -112,29 +142,32 @@ impl AllFunctionsQuery { let instrumentation = None; let definition = Some(Location::from((file_name, start, end))); let qual_fn_name = format!("{class_name}.{method_name}"); - Some(FunctionInfo { - id: (module_name, qual_fn_name).into(), - instrumentation, - definition, + Some(TypescriptFunctionInfo { + inner_info: FunctionInfo { + id: (module_name, qual_fn_name).into(), + instrumentation, + definition, + }, + function_rvalue_range: None, }) } - (_, None, Some(_)) => { + (_, _, None, Some(_)) => { warn!("Found a class without a method in the capture"); None } - (_, Some(_), None) => { + (_, _, Some(_), None) => { warn!("Found a method without a class in the capture"); None } - (Some(Err(e)), _, _) => { + (Some(Err(e)), _, _, _) => { warn!("Could not extract a function name: {e}"); None } - (_, Some(Err(e)), _) => { + (_, _, Some(Err(e)), _) => { warn!("Could not extract a method name: {e}"); None } - (_, _, Some(Err(e))) => { + (_, _, _, Some(Err(e))) => { warn!("Could not extract a class name: {e}"); None } diff --git a/am_list/src/typescript/tests.rs b/am_list/src/typescript/tests.rs index 718d5cf..bda70f1 100644 --- a/am_list/src/typescript/tests.rs +++ b/am_list/src/typescript/tests.rs @@ -108,11 +108,11 @@ const asyncCallMetricized = autometrics(async function asyncCall() { "list of all functions should have 2 items, got this instead: {all:?}" ); assert!( - all.contains(&async_call), + all.iter().any(|info| info.inner_info == async_call), "List of all functions should contain {async_call:?}; complete list is {all:?}" ); assert!( - all.contains(&resolve_after_half), + all.iter().any(|info| info.inner_info == resolve_after_half), "List of all functions should contain {resolve_after_half:?}; complete list is {all:?}" ); } @@ -205,7 +205,7 @@ fn detect_class() { let source = r#" import express from "express"; -@Autometrics +@Autometrics() class Foo { x: number constructor(x = 0) { @@ -335,19 +335,19 @@ class NotGood { ); assert!( - all.contains(&foo_constructor), + all.iter().any(|info| info.inner_info == foo_constructor), "The list of all functions should contain {foo_constructor:?}; complete list is {all:?}" ); assert!( - all.contains(&method_b), + all.iter().any(|info| info.inner_info == method_b), "The list of all functions should contain {method_b:?}; complete list is {all:?}" ); assert!( - all.contains(¬_good_constructor), + all.iter().any(|info| info.inner_info == not_good_constructor), "The list of all functions should contain {not_good_constructor:?}; complete list is {all:?}" ); assert!( - all.contains(&gotgot_method), + all.iter().any(|info| info.inner_info == gotgot_method), "The list of all functions should contain {gotgot_method:?}; complete list is {all:?}" ); } @@ -578,3 +578,137 @@ fn detect_two_args_wrapper() { "list of all functions should have 0 items, got this instead: {all:?}" ); } + +#[test] +fn instrument_multiple() { + let source = r#" +import { Autometrics } from '@autometrics/autometrics'; +import express from "express"; + +@Autometrics() +class Foo { + x: number + constructor(x = 0) { + this.x = x; + } + method_b(): string { + return "you win"; + } +} + +class NotGood { + x: string + constructor(x = "got you") { + this.x = x; + } + gotgot(): string { + return "!"; + } +} + "#; + + let expected = r#" +import { Autometrics } from '@autometrics/autometrics'; +import express from "express"; + +@Autometrics() +class Foo { + x: number + constructor(x = 0) { + this.x = x; + } + method_b(): string { + return "you win"; + } +} + +class NotGood { + x: string +@Autometrics() + constructor(x = "got you") { + this.x = x; + } +@Autometrics() + gotgot(): string { + return "!"; + } +} + "#; + + let mut implementation = Impl {}; + let actual = implementation.instrument_source_code(source).unwrap(); + assert_eq!(&actual, expected); +} + +#[test] +fn instrument_methods_and_top_level_function() { + let source = r#" +class Foo { + method_b(): string { + return "you win"; + } +} + +function gotgot(): string { + return "!"; +} + +async function gotgotButAsync(): string { + return "!"; +} + "#; + + let expected = r#"import { Autometrics } from '@autometrics/autometrics'; + +class Foo { +@Autometrics() + method_b(): string { + return "you win"; + } +} + +@Autometrics() +function gotgot(): string { + return "!"; +} + +@Autometrics() +async function gotgotButAsync(): string { + return "!"; +} + "#; + + let mut implementation = Impl {}; + let actual = implementation.instrument_source_code(source).unwrap(); + assert_eq!(&actual, expected); +} + +#[test] +fn instrument_function_assignment() { + let source = r#" + +const bang = async (argument) => string { + return "!"; +}; + +const bang_too = async function (): string { + return "!"; +}; + "#; + + let expected = r#"import { autometrics } from '@autometrics/autometrics'; + + +const bang = autometrics(async (argument) => string { + return "!"; +}); + +const bang_too = autometrics(async function (): string { + return "!"; +}); + "#; + + let mut implementation = Impl {}; + let actual = implementation.instrument_source_code(source).unwrap(); + assert_eq!(&actual, expected); +}