//! Function extraction for similarity scoring (issue #999, stage 2).
//!
//! A [`FunctionExtractor`] turns a source file into its functions, each with
//! a name, a span and the pre-order sequence of syntax-tree node types
//! inside it. Node *types* only: identifiers and literal values are not part
//! of the sequence, so the summary describes structure. The scoring in
//! `cpd_core::similarity` is grammar-agnostic; it only ever compares two
//! functions that carry the same [`FunctionExtractor::grammar`] id.
//!
//! # Adding a language
//!
//! 1. Implement [`FunctionExtractor`]: pick a stable `grammar` id (for a
//! tree-sitter grammar, its name), list the jscpd `formats` it serves,
//! and in `extract` walk the tree, opening a [`RawFunction`] at every
//! function-like node and appending each visited node's type id (any
//! dense `u16`, e.g. tree-sitter's `node.kind_id()`) to every open
//! function.
//! 2. Add the extractor to [`EXTRACTORS`].
//!
//! Nothing else changes: the CLI, the MCP tool, the reporters and the
//! fixtures pick the new formats up through [`supports_functions`].
use crate::line_index::LineIndex;
use cpd_core::models::Location;
use oxc_allocator::Allocator;
use oxc_ast::AstKind;
use oxc_ast_visit::Visit;
use oxc_parser::Parser;
use oxc_span::GetSpan;
/// A function found in a source, before token ranges are attached.
#[derive(Debug, Clone, PartialEq)]
pub struct RawFunction {
/// Grammar that produced `kinds`; functions of different grammars are
/// never compared.
pub grammar: &'static str,
pub name: String,
pub start: Location,
pub end: Location,
/// Where the code that names the function starts: the key of a method
/// or property, or the variable a function is assigned to, when that
/// code precedes `start`; `start` otherwise.
pub head: Location,
/// Pre-order syntax-tree node types of the function, itself included.
pub kinds: Vec<u16>,
}
/// Language plug-in for function extraction.
pub trait FunctionExtractor: Send + Sync {
/// Stable identifier of the grammar behind the node-type ids.
fn grammar(&self) -> &'static str;
/// jscpd format names this extractor serves.
fn formats(&self) -> &'static [&'static str];
/// All functions of `source`; empty when the source does not parse.
fn extract(&self, source: &str, format: &str) -> Vec<RawFunction>;
}
/// Registered extractors, consulted in order. Add new languages here.
pub static EXTRACTORS: &[&dyn FunctionExtractor] = &[&OxcExtractor];
/// The extractor serving `format`, if any.
pub fn extractor_for(format: &str) -> Option<&'static dyn FunctionExtractor> {
EXTRACTORS
.iter()
.copied()
.find(|e| e.formats().contains(&format))
}
/// Formats handled by [`extract_functions`].
pub fn supports_functions(format: &str) -> bool {
extractor_for(format).is_some()
}
/// Every format some extractor serves, for messages and docs.
pub fn supported_function_formats() -> Vec<&'static str> {
EXTRACTORS
.iter()
.flat_map(|e| e.formats().iter().copied())
.collect()
}
/// Extract every function of a source. Returns an empty vector for formats
/// without an extractor and for sources that fail to parse.
pub fn extract_functions(source: &str, format: &str) -> Vec<RawFunction> {
extract_with(extractor_for(format), source, format)
}
/// Every function of `source` as `extractor` finds them; empty without an
/// extractor and for an empty source. For callers that pick extractors from
/// a registry of their own, like `--semantic`'s.
pub fn extract_with(
extractor: Option<&dyn FunctionExtractor>,
source: &str,
format: &str,
) -> Vec<RawFunction> {
match extractor {
Some(extractor) if !source.is_empty() => extractor.extract(source, format),
_ => Vec::new(),
}
}
/// JavaScript, TypeScript, JSX and TSX through the oxc parser.
pub struct OxcExtractor;
impl FunctionExtractor for OxcExtractor {
fn grammar(&self) -> &'static str {
"oxc"
}
fn formats(&self) -> &'static [&'static str] {
&["javascript", "typescript", "jsx", "tsx"]
}
fn extract(&self, source: &str, format: &str) -> Vec<RawFunction> {
extract_with_oxc(source, format)
}
}
fn extract_with_oxc(source: &str, format: &str) -> Vec<RawFunction> {
let allocator = Allocator::new();
let source_type = crate::javascript::source_type_for_format(format);
let parsed = Parser::new(&allocator, source, source_type).parse();
// Recoverable diagnostics leave a usable (possibly partial) AST; only a
// parser that gave up yields nothing (issue #1023).
if parsed.fatal_error {
return Vec::new();
}
let line_index = LineIndex::new(source.as_bytes());
let mut extractor = Extractor {
frames: Vec::new(),
out: Vec::new(),
pending_name: None,
pending_head: None,
pending_call: None,
line_index: &line_index,
len: source.len(),
};
extractor.visit_program(&parsed.program);
extractor.out
}
struct Frame {
name: String,
head: u32,
start: u32,
end: u32,
kinds: Vec<u16>,
}
struct Extractor<'i> {
frames: Vec<Frame>,
out: Vec<RawFunction>,
/// Name from the enclosing declarator, property or method, consumed by
/// the next function node.
pending_name: Option<String>,
/// Where the code that named `pending_name` starts, and where its
/// value starts: the function that starts there takes that code as its
/// head.
pending_head: Option<(u32, u32)>,
/// The call (by its span: `test.each(t)('title', fn)` and its inner
/// `test.each(t)` start at the same byte) that set `pending_name` for
/// its test-case callback; leaving that call drops a name no function
/// took.
pending_call: Option<(u32, u32)>,
line_index: &'i LineIndex,
len: usize,
}
impl Extractor<'_> {
fn open(&mut self, name: String, start: u32, end: u32) {
// Only the function the naming code is about takes its head; one
// that opens before it (an arrow in `test.each(table)`) leaves it.
let head = match self.pending_head {
Some((head, value)) if value == start => {
self.pending_head = None;
head
}
_ => start,
};
self.frames.push(Frame {
name,
head,
start,
end,
kinds: Vec::new(),
});
}
fn close(&mut self) {
let Some(frame) = self.frames.pop() else {
return;
};
let start = (frame.start as usize).min(self.len);
let end = (frame.end as usize).min(self.len);
let head = (frame.head as usize).min(start);
self.out.push(RawFunction {
grammar: OxcExtractor.grammar(),
name: frame.name,
start: self.line_index.location(start),
end: self.line_index.location(end),
head: self.line_index.location(head),
kinds: frame.kinds,
});
}
/// The pending name, for the function that starts at `start`. A test
/// title belongs to its callback alone: a function in `.each(table)`
/// before it, or one nested in a named callback, does not take it.
fn take_name(&mut self, start: u32) -> Option<String> {
if self.pending_call.is_some() {
if self.pending_head.map(|(_, callback)| callback) != Some(start) {
return None;
}
self.pending_call = None;
}
self.pending_name.take()
}
/// Remember where the code naming the next function starts, for the
/// function that is the named value itself (`value` starts there), not
/// one nested in it.
fn name_head(&mut self, head: u32, value: Option<u32>) {
self.pending_head = match (&self.pending_name, value) {
(Some(_), Some(value)) => Some((head, value)),
_ => None,
};
}
}
impl<'a> Visit<'a> for Extractor<'_> {
fn enter_node(&mut self, kind: AstKind<'a>) {
match kind {
AstKind::VariableDeclarator(d) => {
self.pending_name = d.id.get_identifier_name().map(|n| n.to_string());
self.name_head(d.span.start, d.init.as_ref().map(|v| v.span().start));
}
AstKind::MethodDefinition(m) => {
self.pending_name = m.key.static_name().map(|n| n.into_owned());
self.name_head(m.key.span().start, Some(m.value.span.start));
}
AstKind::PropertyDefinition(p) => {
self.pending_name = p.key.static_name().map(|n| n.into_owned());
self.name_head(p.key.span().start, p.value.as_ref().map(|v| v.span().start));
}
AstKind::ObjectProperty(p) => {
self.pending_name = p.key.static_name().map(|n| n.into_owned());
self.name_head(p.key.span().start, Some(p.value.span().start));
}
AstKind::CallExpression(call) => {
if let Some((title, callback)) = test_case(call) {
self.pending_name = Some(title);
self.pending_head = Some((call.span.start, callback));
self.pending_call = Some((call.span.start, call.span.end));
}
}
AstKind::Function(f) => {
let own = f.id.as_ref().map(|id| id.name.to_string());
let name = match own {
Some(name) => name,
None => self
.take_name(f.span.start)
.unwrap_or_else(|| "<anonymous>".to_string()),
};
let span = f.span;
self.open(name, span.start, span.end);
}
AstKind::ArrowFunctionExpression(a) => {
let name = self
.take_name(a.span.start)
.unwrap_or_else(|| "<arrow>".to_string());
let span = a.span;
self.open(name, span.start, span.end);
}
_ => {}
}
let ty = kind.ty() as u16;
for frame in &mut self.frames {
frame.kinds.push(ty);
}
}
fn leave_node(&mut self, kind: AstKind<'a>) {
match kind {
AstKind::Function(_) | AstKind::ArrowFunctionExpression(_) => self.close(),
AstKind::CallExpression(call)
if self.pending_call == Some((call.span.start, call.span.end)) =>
{
self.pending_name = None;
self.pending_head = None;
self.pending_call = None;
}
AstKind::VariableDeclarator(_)
| AstKind::MethodDefinition(_)
| AstKind::PropertyDefinition(_)
| AstKind::ObjectProperty(_) => {
self.pending_name = None;
self.pending_head = None;
}
_ => {}
}
let _ = kind.span();
}
}
/// Functions that declare one test case in the JavaScript test frameworks
/// (Jest, Vitest, Mocha, Jasmine, node:test, Bun): `it('title', fn)`, with
/// `.only`, `.skip`, `.each(table)` and the like after it. Suites
/// (`describe`) and hooks (`beforeEach`) are left out: they group or set up
/// tests, while a test case is what a port carries over one by one.
pub const TEST_CASE_CALLS: &[&str] = &["it", "test", "specify", "fit", "xit", "xtest", "bench"];
/// Members of a test function that declare something other than a test
/// case: a suite, a step inside a test, a hook, or configuration
/// (`test.describe`, `test.step`, `test.beforeEach`, `test.use`).
pub const NOT_TEST_CASES: &[&str] = &[
"describe",
"step",
"beforeEach",
"afterEach",
"beforeAll",
"afterAll",
"use",
"extend",
];
/// The title of the test case `call` declares and where its callback
/// starts, when `call` is `it('rounds cents', () => …)` or one of its
/// variants and the title is a plain string. The callback then goes by the
/// title, which is what names a test in these frameworks, where it would
/// otherwise be an anonymous arrow; its text starts at the call, so the
/// title is part of what a model sees.
fn test_case(call: &oxc_ast::ast::CallExpression<'_>) -> Option<(String, u32)> {
use oxc_ast::ast::Expression;
// `it`, `it.only`, `test.each(table)`, `it.concurrent.each(table)`,
// but not Playwright's `test.describe(…)` or `test.step(…)`.
let mut callee = &call.callee;
let root = loop {
match callee {
Expression::Identifier(id) => break id.name.as_str(),
Expression::StaticMemberExpression(member) => {
if NOT_TEST_CASES.contains(&member.property.name.as_str()) {
return None;
}
callee = &member.object;
}
Expression::CallExpression(inner) => callee = &inner.callee,
_ => return None,
}
};
if !TEST_CASE_CALLS.contains(&root) {
return None;
}
let mut args = call.arguments.iter().filter_map(|a| a.as_expression());
let title = match args.next()? {
Expression::StringLiteral(literal) => literal.value.to_string(),
Expression::TemplateLiteral(template) if template.expressions.is_empty() => {
let text = template.quasis.first()?;
text.value
.cooked
.as_ref()
.unwrap_or(&text.value.raw)
.to_string()
}
_ => return None,
};
let callback = args.find_map(|a| match a {
Expression::ArrowFunctionExpression(f) => Some(f.span.start),
Expression::FunctionExpression(f) => Some(f.span.start),
_ => None,
})?;
let title = title.split_whitespace().collect::<Vec<_>>().join(" ");
(!title.is_empty()).then_some((title, callback))
}
#[cfg(test)]
mod tests {
use super::*;
const SRC: &str = "export function total(items) {\n let sum = 0;\n for (const it of items) { sum += it.price; }\n return sum;\n}\nconst double = (x) => x * 2;\nclass Cart {\n add(item) { this.items.push(item); }\n}\nconst obj = { run() { return 1; }, cb: function () { return 2; } };\n";
#[test]
fn extracts_declarations_arrows_methods_and_properties_with_names() {
let fns = extract_functions(SRC, "javascript");
let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect();
assert_eq!(names, vec!["total", "double", "add", "run", "cb"]);
let total = &fns[0];
assert_eq!((total.start.line, total.end.line), (1, 5));
assert!(total.kinds.len() > 20, "{}", total.kinds.len());
assert_eq!(total.kinds[0], oxc_ast::AstType::Function as u16);
}
#[test]
fn nested_functions_are_emitted_separately_and_contribute_to_the_outer() {
let src = "function outer() {\n const inner = () => 1;\n return inner();\n}\n";
let fns = extract_functions(src, "typescript");
assert_eq!(fns.len(), 2);
assert_eq!(fns[0].name, "inner"); // closed first
assert_eq!(fns[1].name, "outer");
assert!(fns[1].kinds.len() > fns[0].kinds.len());
}
#[test]
fn test_case_callbacks_go_by_their_titles() {
let src = "describe('money', () => {\n beforeEach(() => reset());\n it('rounds cents', () => {\n expect(round(149)).toBe(100);\n });\n test.each([[1, 2]])('adds %i', (a, b) => {\n expect(a + b).toBe(3);\n });\n it.only(`keeps ${'x'} dynamic`, () => {});\n it('has no callback');\n const later = () => 1;\n test(\"async one\", async function () { await later(); });\n it('named', function named() { [1].map((x) => x * 2); });\n test.each([[() => 1]])('table', (f) => f());\n test.describe('suite', () => {});\n test.step('step', async () => {});\n});\n";
let fns = extract_functions(src, "typescript");
let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect();
assert_eq!(
names,
vec![
"<arrow>",
"rounds cents",
"adds %i",
"<arrow>",
"later",
"async one",
"<arrow>",
"named",
"<arrow>",
"table",
"<arrow>",
"<arrow>",
"<arrow>"
],
"hooks, suites and dynamic titles stay anonymous; a test without a callback names nothing"
);
// The test's code starts at the call, so its title is part of it.
let rounds = &fns[1];
assert!(src[rounds.head.offset as usize..].starts_with("it('rounds cents', () =>"));
assert!(src[rounds.start.offset as usize..].starts_with("() =>"));
}
#[test]
fn registry_dispatches_by_format_and_tags_the_grammar() {
assert_eq!(extractor_for("typescript").unwrap().grammar(), "oxc");
assert!(extractor_for("python").is_none());
let formats = supported_function_formats();
for f in ["javascript", "typescript", "jsx", "tsx"] {
assert!(formats.contains(&f), "{f}");
assert!(supports_functions(f));
}
let fns = extract_functions("const f = () => 1;", "jsx");
assert_eq!(fns.len(), 1);
assert_eq!(fns[0].grammar, "oxc");
}
#[test]
fn unsupported_or_empty_sources_yield_nothing() {
assert!(extract_functions("def f():\n pass\n", "python").is_empty());
assert!(extract_functions("", "javascript").is_empty());
}
#[test]
fn redeclared_functions_are_still_extracted() {
let src = "function f(a) { return a + 1; }\nfunction f(b) { return b + 1; }\n";
let fns = extract_functions(src, "javascript");
assert_eq!(
fns.len(),
2,
"a redeclaration diagnostic must not drop the file"
);
assert_eq!(fns[0].kinds, fns[1].kinds);
}
#[test]
fn renamed_copies_share_the_same_kind_sequence() {
let a = extract_functions("function a(x) { return x + 1; }", "javascript");
let b = extract_functions("function b(y) { return y + 1; }", "javascript");
assert_eq!(a[0].kinds, b[0].kinds);
}
}