1use std::collections::HashSet;
9
10use ast_grep_core::{AstGrep, Doc};
11use serde::{Deserialize, Serialize};
12use strum::Display;
13
14use crate::rule::Language;
15use crate::test_filter::CfgTestRanges;
16
17#[derive(Debug, Display, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20#[strum(serialize_all = "snake_case")]
21pub enum Metric {
22 FileLines,
23 ImportCount,
24 FunctionCount,
25 CommentRatio,
26}
27
28fn import_kinds(lang: &Language) -> &'static [&'static str] {
30 match lang {
31 Language::Rust => &["use_declaration", "extern_crate_declaration"],
32 Language::TypeScript => &["import_statement"],
33 }
34}
35
36fn function_kinds(lang: &Language) -> &'static [&'static str] {
39 match lang {
40 Language::Rust => &["function_item"],
41 Language::TypeScript => &[
42 "function_declaration",
43 "generator_function_declaration",
44 "function_expression",
45 "arrow_function",
46 "method_definition",
47 ],
48 }
49}
50
51fn comment_kinds(lang: &Language) -> &'static [&'static str] {
53 match lang {
54 Language::Rust => &["line_comment", "block_comment", "doc_comment"],
55 Language::TypeScript => &["comment"],
56 }
57}
58
59pub fn compute<D: Doc>(
65 metric: Metric,
66 lang: &Language,
67 root: &AstGrep<D>,
68 source: &str,
69 exclude_tests: Option<&CfgTestRanges>,
70) -> f64 {
71 match metric {
72 Metric::FileLines => count_lines(source, exclude_tests) as f64,
73 Metric::ImportCount => count_kinds(root, import_kinds(lang), exclude_tests) as f64,
74 Metric::FunctionCount => count_kinds(root, function_kinds(lang), exclude_tests) as f64,
75 Metric::CommentRatio => {
76 let total_lines = count_lines(source, exclude_tests);
77 if total_lines == 0 {
78 return 0.0;
79 }
80 comment_lines(root, comment_kinds(lang), exclude_tests) as f64 / total_lines as f64
81 }
82 }
83}
84
85fn in_test(exclude_tests: Option<&CfgTestRanges>, zero_indexed_line: usize) -> bool {
88 exclude_tests.is_some_and(|r| r.contains_line(zero_indexed_line + 1))
89}
90
91fn count_lines(source: &str, exclude_tests: Option<&CfgTestRanges>) -> usize {
96 source
97 .lines()
98 .enumerate()
99 .filter(|(i, _)| !in_test(exclude_tests, *i))
100 .count()
101}
102
103fn count_kinds<D: Doc>(
104 root: &AstGrep<D>,
105 kinds: &[&str],
106 exclude_tests: Option<&CfgTestRanges>,
107) -> usize {
108 let mut count = 0;
109 for node in root.root().dfs() {
110 if kinds.contains(&node.kind().as_ref()) && !in_test(exclude_tests, node.start_pos().line())
111 {
112 count += 1;
113 }
114 }
115 count
116}
117
118fn comment_lines<D: Doc>(
127 root: &AstGrep<D>,
128 kinds: &[&str],
129 exclude_tests: Option<&CfgTestRanges>,
130) -> usize {
131 let mut lines = HashSet::new();
132 let mut doc_lines = HashSet::new();
133 for node in root.root().dfs() {
134 let kind = node.kind();
135 if kind.as_ref() == "doc_comment" {
136 doc_lines.extend(node.start_pos().line()..=node.end_pos().line());
137 continue;
138 }
139 if !kinds.contains(&kind.as_ref()) {
140 continue;
141 }
142 let range = node.start_pos().line()..=node.end_pos().line();
143 if node.text().trim_start().starts_with("/**") {
144 doc_lines.extend(range);
145 } else {
146 lines.extend(range);
147 }
148 }
149 lines
150 .into_iter()
151 .filter(|line| !doc_lines.contains(line) && !in_test(exclude_tests, *line))
152 .count()
153}
154
155impl Metric {
156 pub fn format_value(&self, value: f64) -> String {
159 match self {
160 Metric::CommentRatio => format!("{value:.2}"),
161 _ => (value as u64).to_string(),
162 }
163 }
164
165 pub fn describe(&self, value: f64) -> String {
168 let rendered = self.format_value(value);
169 match self {
170 Metric::FileLines => format!("{rendered} lines"),
171 Metric::ImportCount => format!("{rendered} imports"),
172 Metric::FunctionCount => format!("{rendered} functions"),
173 Metric::CommentRatio => format!("{rendered} comment ratio"),
174 }
175 }
176}
177
178pub fn exceeds(value: f64, threshold: f64) -> bool {
181 value > threshold
182}
183
184#[cfg(test)]
185mod tests {
186 use ast_grep_core::tree_sitter::LanguageExt;
187 use ast_grep_language::SupportLang;
188
189 use super::*;
190
191 fn rust(metric: Metric, source: &str) -> f64 {
192 let root = SupportLang::Rust.ast_grep(source);
193 compute(metric, &Language::Rust, &root, source, None)
194 }
195
196 fn typescript(metric: Metric, source: &str) -> f64 {
197 let root = SupportLang::TypeScript.ast_grep(source);
198 compute(metric, &Language::TypeScript, &root, source, None)
199 }
200
201 fn cfg_ranges(source: &str) -> CfgTestRanges {
202 CfgTestRanges::from_root(&SupportLang::Rust.ast_grep(source))
203 }
204
205 fn rust_no_tests(metric: Metric, source: &str) -> f64 {
206 let root = SupportLang::Rust.ast_grep(source);
207 compute(
208 metric,
209 &Language::Rust,
210 &root,
211 source,
212 Some(&cfg_ranges(source)),
213 )
214 }
215
216 #[test]
217 fn file_lines_counts_lines() {
218 let source = "fn a() {}\nfn b() {}\nfn c() {}\n";
219 assert_eq!(rust(Metric::FileLines, source), 3.0);
220 assert_eq!(
221 rust(Metric::FileLines, "fn a() {}\n"),
222 1.0,
223 "a trailing newline terminates the last line"
224 );
225 }
226
227 #[test]
228 fn file_lines_empty_file_is_zero() {
229 assert_eq!(rust(Metric::FileLines, ""), 0.0);
230 }
231
232 #[test]
233 fn import_count_rust() {
234 let source = "\
235use std::fmt;
236use std::io::Write;
237extern crate serde;
238
239fn f() {
240 use std::collections::HashMap;
241 let _: HashMap<u8, u8> = HashMap::new();
242}
243";
244 assert_eq!(
245 rust(Metric::ImportCount, source),
246 4.0,
247 "a nested `use` inside a function still counts"
248 );
249 }
250
251 #[test]
252 fn import_count_typescript() {
253 let source = "\
254import { a } from \"./a\";
255import b from \"./b\";
256const c = require(\"./c\");
257";
258 assert_eq!(
259 typescript(Metric::ImportCount, source),
260 2.0,
261 "a require() call is not an import statement"
262 );
263 }
264
265 #[test]
266 fn function_count_rust() {
267 let source = "\
268trait T {
269 fn signature_only(&self);
270}
271
272struct S;
273
274impl S {
275 fn method(&self) {}
276 fn other(&self) {}
277}
278
279fn free() {}
280";
281 assert_eq!(
282 rust(Metric::FunctionCount, source),
283 3.0,
284 "a trait signature has no body and does not count"
285 );
286 }
287
288 #[test]
289 fn function_count_typescript() {
290 let source = "\
291function declared() {}
292const arrow = () => 1;
293class C {
294 method() {}
295}
296";
297 assert_eq!(typescript(Metric::FunctionCount, source), 3.0);
298 }
299
300 #[test]
301 fn excludes_cfg_test_lines_and_functions() {
302 let source = "\
303use std::fmt;
304
305fn prod_one() {}
306fn prod_two() {}
307
308#[cfg(test)]
309mod tests {
310 use super::*;
311
312 fn helper() {}
313 fn t_one() {}
314 fn t_two() {}
315}
316";
317 assert_eq!(rust(Metric::FunctionCount, source), 5.0);
319 assert_eq!(rust(Metric::ImportCount, source), 2.0);
320 assert_eq!(rust(Metric::FileLines, source), 13.0);
321 assert_eq!(rust_no_tests(Metric::FunctionCount, source), 2.0);
324 assert_eq!(rust_no_tests(Metric::ImportCount, source), 1.0);
325 assert_eq!(rust_no_tests(Metric::FileLines, source), 5.0);
326 }
327
328 #[test]
329 fn excludes_cfg_test_comment_lines() {
330 let source = "\
331// prod comment
332fn prod() {}
333#[cfg(test)]
334mod tests {
335 // test comment
336 // another test comment
337 fn t() {}
338}
339";
340 assert_eq!(rust(Metric::CommentRatio, source), 3.0 / 8.0);
342 assert_eq!(rust_no_tests(Metric::CommentRatio, source), 1.0 / 2.0);
344 }
345
346 #[test]
347 fn no_cfg_test_leaves_metrics_unchanged() {
348 let source = "use std::fmt;\nfn a() {}\nfn b() {}\n";
349 assert_eq!(
350 rust_no_tests(Metric::FunctionCount, source),
351 rust(Metric::FunctionCount, source)
352 );
353 assert_eq!(
354 rust_no_tests(Metric::FileLines, source),
355 rust(Metric::FileLines, source)
356 );
357 }
358
359 #[test]
360 fn comment_ratio_line_comments() {
361 let source = "\
362// one
363// two
364fn a() {}
365// three
366fn b() {}
367// four
368fn c() {}
369fn d() {}
370";
371 assert_eq!(rust(Metric::CommentRatio, source), 0.5);
372 }
373
374 #[test]
375 fn comment_ratio_block_comment_spans_lines() {
376 let source = "\
377/* one
378 two
379 three
380 four
381 five */
382fn a() {}
383fn b() {}
384fn c() {}
385fn d() {}
386fn e() {}
387";
388 assert_eq!(
389 rust(Metric::CommentRatio, source),
390 0.5,
391 "a 5-line block comment contributes 5 comment lines out of 10"
392 );
393 }
394
395 #[test]
396 fn comment_ratio_ignores_doc_comments() {
397 let source = "\
399//! module doc
400/// outer doc for a
401/// still doc
402fn a() {}
403/** block doc */
404fn b() {}
405";
406 assert_eq!(
407 rust(Metric::CommentRatio, source),
408 0.0,
409 "a file with only doc comments has a zero narration ratio"
410 );
411 }
412
413 #[test]
414 fn comment_ratio_counts_narration_but_not_docs() {
415 let source = "\
418// narrate one
419// narrate two
420/// doc one
421/// doc two
422fn a() {}
423fn b() {}
424";
425 assert_eq!(rust(Metric::CommentRatio, source), 2.0 / 6.0);
426 }
427
428 #[test]
429 fn comment_ratio_typescript_ignores_jsdoc() {
430 let source = "\
431/** jsdoc line one
432 * jsdoc line two
433 */
434// real narration
435function f() {}
436";
437 assert_eq!(typescript(Metric::CommentRatio, source), 1.0 / 5.0);
439 }
440
441 #[test]
442 fn comment_ratio_empty_file_is_zero() {
443 assert_eq!(rust(Metric::CommentRatio, ""), 0.0);
444 }
445
446 #[test]
447 fn comment_ratio_no_comments_is_zero() {
448 let source = "fn a() {}\nfn b() {}\n";
449 assert_eq!(rust(Metric::CommentRatio, source), 0.0);
450 }
451
452 #[test]
453 fn format_value_and_describe() {
454 assert_eq!(Metric::FileLines.format_value(547.0), "547");
455 assert_eq!(Metric::FileLines.describe(547.0), "547 lines");
456 assert_eq!(Metric::ImportCount.describe(42.0), "42 imports");
457 assert_eq!(Metric::FunctionCount.describe(31.0), "31 functions");
458 assert_eq!(Metric::CommentRatio.format_value(0.617), "0.62");
459 let ratio = Metric::CommentRatio.describe(0.617);
460 assert_eq!(ratio, "0.62 comment ratio");
461 }
462
463 #[test]
464 fn exceeds_is_strict() {
465 assert!(!exceeds(500.0, 500.0));
466 assert!(exceeds(501.0, 500.0));
467 assert!(!exceeds(499.0, 500.0));
468 }
469
470 #[test]
471 fn metric_display_is_snake_case() {
472 assert_eq!(Metric::FileLines.to_string(), "file_lines");
473 assert_eq!(Metric::CommentRatio.to_string(), "comment_ratio");
474 }
475}