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    /// Where the code that names the function starts: the key of a method
41    /// or property, or the variable a function is assigned to, when that
42    /// code precedes `start`; `start` otherwise.
43    pub head: Location,
44    /// Pre-order syntax-tree node types of the function, itself included.
45    pub kinds: Vec<u16>,
46}
47
48/// Language plug-in for function extraction.
49pub trait FunctionExtractor: Send + Sync {
50    /// Stable identifier of the grammar behind the node-type ids.
51    fn grammar(&self) -> &'static str;
52    /// jscpd format names this extractor serves.
53    fn formats(&self) -> &'static [&'static str];
54    /// All functions of `source`; empty when the source does not parse.
55    fn extract(&self, source: &str, format: &str) -> Vec<RawFunction>;
56}
57
58/// Registered extractors, consulted in order. Add new languages here.
59pub static EXTRACTORS: &[&dyn FunctionExtractor] = &[&OxcExtractor];
60
61/// The extractor serving `format`, if any.
62pub fn extractor_for(format: &str) -> Option<&'static dyn FunctionExtractor> {
63    EXTRACTORS
64        .iter()
65        .copied()
66        .find(|e| e.formats().contains(&format))
67}
68
69/// Formats handled by [`extract_functions`].
70pub fn supports_functions(format: &str) -> bool {
71    extractor_for(format).is_some()
72}
73
74/// Every format some extractor serves, for messages and docs.
75pub fn supported_function_formats() -> Vec<&'static str> {
76    EXTRACTORS
77        .iter()
78        .flat_map(|e| e.formats().iter().copied())
79        .collect()
80}
81
82/// Extract every function of a source. Returns an empty vector for formats
83/// without an extractor and for sources that fail to parse.
84pub fn extract_functions(source: &str, format: &str) -> Vec<RawFunction> {
85    extract_with(extractor_for(format), source, format)
86}
87
88/// Every function of `source` as `extractor` finds them; empty without an
89/// extractor and for an empty source. For callers that pick extractors from
90/// a registry of their own, like `--semantic`'s.
91pub fn extract_with(
92    extractor: Option<&dyn FunctionExtractor>,
93    source: &str,
94    format: &str,
95) -> Vec<RawFunction> {
96    match extractor {
97        Some(extractor) if !source.is_empty() => extractor.extract(source, format),
98        _ => Vec::new(),
99    }
100}
101
102/// JavaScript, TypeScript, JSX and TSX through the oxc parser.
103pub struct OxcExtractor;
104
105impl FunctionExtractor for OxcExtractor {
106    fn grammar(&self) -> &'static str {
107        "oxc"
108    }
109
110    fn formats(&self) -> &'static [&'static str] {
111        &["javascript", "typescript", "jsx", "tsx"]
112    }
113
114    fn extract(&self, source: &str, format: &str) -> Vec<RawFunction> {
115        extract_with_oxc(source, format)
116    }
117}
118
119fn extract_with_oxc(source: &str, format: &str) -> Vec<RawFunction> {
120    let allocator = Allocator::new();
121    let source_type = crate::javascript::source_type_for_format(format);
122    let parsed = Parser::new(&allocator, source, source_type).parse();
123    // Recoverable diagnostics leave a usable (possibly partial) AST; only a
124    // parser that gave up yields nothing (issue #1023).
125    if parsed.fatal_error {
126        return Vec::new();
127    }
128    let line_index = LineIndex::new(source.as_bytes());
129    let mut extractor = Extractor {
130        frames: Vec::new(),
131        out: Vec::new(),
132        pending_name: None,
133        pending_head: None,
134        line_index: &line_index,
135        len: source.len(),
136    };
137    extractor.visit_program(&parsed.program);
138    extractor.out
139}
140
141struct Frame {
142    name: String,
143    head: u32,
144    start: u32,
145    end: u32,
146    kinds: Vec<u16>,
147}
148
149struct Extractor<'i> {
150    frames: Vec<Frame>,
151    out: Vec<RawFunction>,
152    /// Name from the enclosing declarator, property or method, consumed by
153    /// the next function node.
154    pending_name: Option<String>,
155    /// Where the code that named `pending_name` starts, and where its
156    /// value starts: the function that starts there takes that code as its
157    /// head.
158    pending_head: Option<(u32, u32)>,
159    line_index: &'i LineIndex,
160    len: usize,
161}
162
163impl Extractor<'_> {
164    fn open(&mut self, name: String, start: u32, end: u32) {
165        let head = self
166            .pending_head
167            .take()
168            .filter(|&(_, value)| value == start)
169            .map_or(start, |(head, _)| head);
170        self.frames.push(Frame {
171            name,
172            head,
173            start,
174            end,
175            kinds: Vec::new(),
176        });
177    }
178
179    fn close(&mut self) {
180        let Some(frame) = self.frames.pop() else {
181            return;
182        };
183        let start = (frame.start as usize).min(self.len);
184        let end = (frame.end as usize).min(self.len);
185        let head = (frame.head as usize).min(start);
186        self.out.push(RawFunction {
187            grammar: OxcExtractor.grammar(),
188            name: frame.name,
189            start: self.line_index.location(start),
190            end: self.line_index.location(end),
191            head: self.line_index.location(head),
192            kinds: frame.kinds,
193        });
194    }
195
196    /// Remember where the code naming the next function starts, for the
197    /// function that is the named value itself (`value` starts there), not
198    /// one nested in it.
199    fn name_head(&mut self, head: u32, value: Option<u32>) {
200        self.pending_head = match (&self.pending_name, value) {
201            (Some(_), Some(value)) => Some((head, value)),
202            _ => None,
203        };
204    }
205}
206
207impl<'a> Visit<'a> for Extractor<'_> {
208    fn enter_node(&mut self, kind: AstKind<'a>) {
209        match kind {
210            AstKind::VariableDeclarator(d) => {
211                self.pending_name = d.id.get_identifier_name().map(|n| n.to_string());
212                self.name_head(d.span.start, d.init.as_ref().map(|v| v.span().start));
213            }
214            AstKind::MethodDefinition(m) => {
215                self.pending_name = m.key.static_name().map(|n| n.into_owned());
216                self.name_head(m.key.span().start, Some(m.value.span.start));
217            }
218            AstKind::PropertyDefinition(p) => {
219                self.pending_name = p.key.static_name().map(|n| n.into_owned());
220                self.name_head(p.key.span().start, p.value.as_ref().map(|v| v.span().start));
221            }
222            AstKind::ObjectProperty(p) => {
223                self.pending_name = p.key.static_name().map(|n| n.into_owned());
224                self.name_head(p.key.span().start, Some(p.value.span().start));
225            }
226            AstKind::Function(f) => {
227                let name =
228                    f.id.as_ref()
229                        .map(|id| id.name.to_string())
230                        .or_else(|| self.pending_name.take())
231                        .unwrap_or_else(|| "<anonymous>".to_string());
232                let span = f.span;
233                self.open(name, span.start, span.end);
234            }
235            AstKind::ArrowFunctionExpression(a) => {
236                let name = self
237                    .pending_name
238                    .take()
239                    .unwrap_or_else(|| "<arrow>".to_string());
240                let span = a.span;
241                self.open(name, span.start, span.end);
242            }
243            _ => {}
244        }
245        let ty = kind.ty() as u16;
246        for frame in &mut self.frames {
247            frame.kinds.push(ty);
248        }
249    }
250
251    fn leave_node(&mut self, kind: AstKind<'a>) {
252        match kind {
253            AstKind::Function(_) | AstKind::ArrowFunctionExpression(_) => self.close(),
254            AstKind::VariableDeclarator(_)
255            | AstKind::MethodDefinition(_)
256            | AstKind::PropertyDefinition(_)
257            | AstKind::ObjectProperty(_) => {
258                self.pending_name = None;
259                self.pending_head = None;
260            }
261            _ => {}
262        }
263        let _ = kind.span();
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270
271    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";
272
273    #[test]
274    fn extracts_declarations_arrows_methods_and_properties_with_names() {
275        let fns = extract_functions(SRC, "javascript");
276        let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect();
277        assert_eq!(names, vec!["total", "double", "add", "run", "cb"]);
278        let total = &fns[0];
279        assert_eq!((total.start.line, total.end.line), (1, 5));
280        assert!(total.kinds.len() > 20, "{}", total.kinds.len());
281        assert_eq!(total.kinds[0], oxc_ast::AstType::Function as u16);
282    }
283
284    #[test]
285    fn nested_functions_are_emitted_separately_and_contribute_to_the_outer() {
286        let src = "function outer() {\n  const inner = () => 1;\n  return inner();\n}\n";
287        let fns = extract_functions(src, "typescript");
288        assert_eq!(fns.len(), 2);
289        assert_eq!(fns[0].name, "inner"); // closed first
290        assert_eq!(fns[1].name, "outer");
291        assert!(fns[1].kinds.len() > fns[0].kinds.len());
292    }
293
294    #[test]
295    fn registry_dispatches_by_format_and_tags_the_grammar() {
296        assert_eq!(extractor_for("typescript").unwrap().grammar(), "oxc");
297        assert!(extractor_for("python").is_none());
298        let formats = supported_function_formats();
299        for f in ["javascript", "typescript", "jsx", "tsx"] {
300            assert!(formats.contains(&f), "{f}");
301            assert!(supports_functions(f));
302        }
303        let fns = extract_functions("const f = () => 1;", "jsx");
304        assert_eq!(fns.len(), 1);
305        assert_eq!(fns[0].grammar, "oxc");
306    }
307
308    #[test]
309    fn unsupported_or_empty_sources_yield_nothing() {
310        assert!(extract_functions("def f():\n  pass\n", "python").is_empty());
311        assert!(extract_functions("", "javascript").is_empty());
312    }
313
314    #[test]
315    fn redeclared_functions_are_still_extracted() {
316        let src = "function f(a) { return a + 1; }\nfunction f(b) { return b + 1; }\n";
317        let fns = extract_functions(src, "javascript");
318        assert_eq!(
319            fns.len(),
320            2,
321            "a redeclaration diagnostic must not drop the file"
322        );
323        assert_eq!(fns[0].kinds, fns[1].kinds);
324    }
325
326    #[test]
327    fn renamed_copies_share_the_same_kind_sequence() {
328        let a = extract_functions("function a(x) { return x + 1; }", "javascript");
329        let b = extract_functions("function b(y) { return y + 1; }", "javascript");
330        assert_eq!(a[0].kinds, b[0].kinds);
331    }
332}
cpd_tokenizer/functions.rs

cpd_tokenizer/
functions.rs
