cpd_tokenizer/
functions.rs

1//! Function extraction for similarity scoring (issue #999, stage 2).
2//!
3//! A [`FunctionExtractor`] turns a source file into its functions, each with
4//! a name, a span and the pre-order sequence of syntax-tree node types
5//! inside it. Node *types* only: identifiers and literal values are not part
6//! of the sequence, so the summary describes structure. The scoring in
7//! `cpd_core::similarity` is grammar-agnostic; it only ever compares two
8//! functions that carry the same [`FunctionExtractor::grammar`] id.
9//!
10//! # Adding a language
11//!
12//! 1. Implement [`FunctionExtractor`]: pick a stable `grammar` id (for a
13//!    tree-sitter grammar, its name), list the jscpd `formats` it serves,
14//!    and in `extract` walk the tree, opening a [`RawFunction`] at every
15//!    function-like node and appending each visited node's type id (any
16//!    dense `u16`, e.g. tree-sitter's `node.kind_id()`) to every open
17//!    function.
18//! 2. Add the extractor to [`EXTRACTORS`].
19//!
20//! Nothing else changes: the CLI, the MCP tool, the reporters and the
21//! fixtures pick the new formats up through [`supports_functions`].
22
23use crate::line_index::LineIndex;
24use cpd_core::models::Location;
25use oxc_allocator::Allocator;
26use oxc_ast::AstKind;
27use oxc_ast_visit::Visit;
28use oxc_parser::Parser;
29use oxc_span::GetSpan;
30
31/// A function found in a source, before token ranges are attached.
32#[derive(Debug, Clone, PartialEq)]
33pub struct RawFunction {
34    /// Grammar that produced `kinds`; functions of different grammars are
35    /// never compared.
36    pub grammar: &'static str,
37    pub name: String,
38    pub start: Location,
39    pub end: Location,
40    /// Pre-order syntax-tree node types of the function, itself included.
41    pub kinds: Vec<u16>,
42}
43
44/// Language plug-in for function extraction.
45pub trait FunctionExtractor: Send + Sync {
46    /// Stable identifier of the grammar behind the node-type ids.
47    fn grammar(&self) -> &'static str;
48    /// jscpd format names this extractor serves.
49    fn formats(&self) -> &'static [&'static str];
50    /// All functions of `source`; empty when the source does not parse.
51    fn extract(&self, source: &str, format: &str) -> Vec<RawFunction>;
52}
53
54/// Registered extractors, consulted in order. Add new languages here.
55pub static EXTRACTORS: &[&dyn FunctionExtractor] = &[&OxcExtractor];
56
57/// The extractor serving `format`, if any.
58pub fn extractor_for(format: &str) -> Option<&'static dyn FunctionExtractor> {
59    EXTRACTORS
60        .iter()
61        .copied()
62        .find(|e| e.formats().contains(&format))
63}
64
65/// Formats handled by [`extract_functions`].
66pub fn supports_functions(format: &str) -> bool {
67    extractor_for(format).is_some()
68}
69
70/// Every format some extractor serves, for messages and docs.
71pub fn supported_function_formats() -> Vec<&'static str> {
72    EXTRACTORS
73        .iter()
74        .flat_map(|e| e.formats().iter().copied())
75        .collect()
76}
77
78/// Extract every function of a source. Returns an empty vector for formats
79/// without an extractor and for sources that fail to parse.
80pub fn extract_functions(source: &str, format: &str) -> Vec<RawFunction> {
81    match extractor_for(format) {
82        Some(extractor) if !source.is_empty() => extractor.extract(source, format),
83        _ => Vec::new(),
84    }
85}
86
87/// JavaScript, TypeScript, JSX and TSX through the oxc parser.
88pub struct OxcExtractor;
89
90impl FunctionExtractor for OxcExtractor {
91    fn grammar(&self) -> &'static str {
92        "oxc"
93    }
94
95    fn formats(&self) -> &'static [&'static str] {
96        &["javascript", "typescript", "jsx", "tsx"]
97    }
98
99    fn extract(&self, source: &str, format: &str) -> Vec<RawFunction> {
100        extract_with_oxc(source, format)
101    }
102}
103
104fn extract_with_oxc(source: &str, format: &str) -> Vec<RawFunction> {
105    let allocator = Allocator::new();
106    let source_type = crate::javascript::source_type_for_format(format);
107    let parsed = Parser::new(&allocator, source, source_type).parse();
108    // Recoverable diagnostics leave a usable (possibly partial) AST; only a
109    // parser that gave up yields nothing (issue #1023).
110    if parsed.panicked {
111        return Vec::new();
112    }
113    let line_index = LineIndex::new(source.as_bytes());
114    let mut extractor = Extractor {
115        frames: Vec::new(),
116        out: Vec::new(),
117        pending_name: None,
118        line_index: &line_index,
119        len: source.len(),
120    };
121    extractor.visit_program(&parsed.program);
122    extractor.out
123}
124
125struct Frame {
126    name: String,
127    start: u32,
128    end: u32,
129    kinds: Vec<u16>,
130}
131
132struct Extractor<'i> {
133    frames: Vec<Frame>,
134    out: Vec<RawFunction>,
135    /// Name from the enclosing declarator, property or method, consumed by
136    /// the next function node.
137    pending_name: Option<String>,
138    line_index: &'i LineIndex,
139    len: usize,
140}
141
142impl Extractor<'_> {
143    fn open(&mut self, name: String, start: u32, end: u32) {
144        self.frames.push(Frame {
145            name,
146            start,
147            end,
148            kinds: Vec::new(),
149        });
150    }
151
152    fn close(&mut self) {
153        let Some(frame) = self.frames.pop() else {
154            return;
155        };
156        let start = (frame.start as usize).min(self.len);
157        let end = (frame.end as usize).min(self.len);
158        self.out.push(RawFunction {
159            grammar: OxcExtractor.grammar(),
160            name: frame.name,
161            start: self.line_index.location(start),
162            end: self.line_index.location(end),
163            kinds: frame.kinds,
164        });
165    }
166}
167
168impl<'a> Visit<'a> for Extractor<'_> {
169    fn enter_node(&mut self, kind: AstKind<'a>) {
170        match kind {
171            AstKind::VariableDeclarator(d) => {
172                self.pending_name = d.id.get_identifier_name().map(|n| n.to_string());
173            }
174            AstKind::MethodDefinition(m) => {
175                self.pending_name = m.key.static_name().map(|n| n.into_owned());
176            }
177            AstKind::PropertyDefinition(p) => {
178                self.pending_name = p.key.static_name().map(|n| n.into_owned());
179            }
180            AstKind::ObjectProperty(p) => {
181                self.pending_name = p.key.static_name().map(|n| n.into_owned());
182            }
183            AstKind::Function(f) => {
184                let name =
185                    f.id.as_ref()
186                        .map(|id| id.name.to_string())
187                        .or_else(|| self.pending_name.take())
188                        .unwrap_or_else(|| "<anonymous>".to_string());
189                let span = f.span;
190                self.open(name, span.start, span.end);
191            }
192            AstKind::ArrowFunctionExpression(a) => {
193                let name = self
194                    .pending_name
195                    .take()
196                    .unwrap_or_else(|| "<arrow>".to_string());
197                let span = a.span;
198                self.open(name, span.start, span.end);
199            }
200            _ => {}
201        }
202        let ty = kind.ty() as u16;
203        for frame in &mut self.frames {
204            frame.kinds.push(ty);
205        }
206    }
207
208    fn leave_node(&mut self, kind: AstKind<'a>) {
209        match kind {
210            AstKind::Function(_) | AstKind::ArrowFunctionExpression(_) => self.close(),
211            AstKind::VariableDeclarator(_)
212            | AstKind::MethodDefinition(_)
213            | AstKind::PropertyDefinition(_)
214            | AstKind::ObjectProperty(_) => self.pending_name = None,
215            _ => {}
216        }
217        let _ = kind.span();
218    }
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224
225    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";
226
227    #[test]
228    fn extracts_declarations_arrows_methods_and_properties_with_names() {
229        let fns = extract_functions(SRC, "javascript");
230        let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect();
231        assert_eq!(names, vec!["total", "double", "add", "run", "cb"]);
232        let total = &fns[0];
233        assert_eq!((total.start.line, total.end.line), (1, 5));
234        assert!(total.kinds.len() > 20, "{}", total.kinds.len());
235        assert_eq!(total.kinds[0], oxc_ast::AstType::Function as u16);
236    }
237
238    #[test]
239    fn nested_functions_are_emitted_separately_and_contribute_to_the_outer() {
240        let src = "function outer() {\n  const inner = () => 1;\n  return inner();\n}\n";
241        let fns = extract_functions(src, "typescript");
242        assert_eq!(fns.len(), 2);
243        assert_eq!(fns[0].name, "inner"); // closed first
244        assert_eq!(fns[1].name, "outer");
245        assert!(fns[1].kinds.len() > fns[0].kinds.len());
246    }
247
248    #[test]
249    fn registry_dispatches_by_format_and_tags_the_grammar() {
250        assert_eq!(extractor_for("typescript").unwrap().grammar(), "oxc");
251        assert!(extractor_for("python").is_none());
252        let formats = supported_function_formats();
253        for f in ["javascript", "typescript", "jsx", "tsx"] {
254            assert!(formats.contains(&f), "{f}");
255            assert!(supports_functions(f));
256        }
257        let fns = extract_functions("const f = () => 1;", "jsx");
258        assert_eq!(fns.len(), 1);
259        assert_eq!(fns[0].grammar, "oxc");
260    }
261
262    #[test]
263    fn unsupported_or_empty_sources_yield_nothing() {
264        assert!(extract_functions("def f():\n  pass\n", "python").is_empty());
265        assert!(extract_functions("", "javascript").is_empty());
266    }
267
268    #[test]
269    fn redeclared_functions_are_still_extracted() {
270        let src = "function f(a) { return a + 1; }\nfunction f(b) { return b + 1; }\n";
271        let fns = extract_functions(src, "javascript");
272        assert_eq!(
273            fns.len(),
274            2,
275            "a redeclaration diagnostic must not drop the file"
276        );
277        assert_eq!(fns[0].kinds, fns[1].kinds);
278    }
279
280    #[test]
281    fn renamed_copies_share_the_same_kind_sequence() {
282        let a = extract_functions("function a(x) { return x + 1; }", "javascript");
283        let b = extract_functions("function b(y) { return y + 1; }", "javascript");
284        assert_eq!(a[0].kinds, b[0].kinds);
285    }
286}
cpd_tokenizer/functions.rs

cpd_tokenizer/
functions.rs
