Skip to main content

codesynapse_core/ts_extract/
java.rs

1use super::{add_node_if_missing, make_file_node, strip_docstring};
2use crate::error::{CodeSynapseError, Result};
3use crate::extract::{make_id, ImportNode, LanguageExtractor};
4use crate::types::{Edge, ExtractionFragment, Node};
5use std::collections::{HashMap, HashSet};
6use std::path::Path;
7use tree_sitter::{Node as TsNode, Parser};
8
9pub struct TsJavaExtractor;
10
11fn collect_field_call_targets(body: &TsNode<'_>, source: &[u8]) -> HashSet<String> {
12    let mut types: HashSet<String> = HashSet::new();
13    for i in 0..body.child_count() {
14        if let Some(child) = body.child(i) {
15            if child.kind() == "field_declaration" {
16                if let Some(type_node) = child.child_by_field_name("type") {
17                    let raw = jv_text(source, &type_node);
18                    let base: String = raw
19                        .chars()
20                        .take_while(|c| c.is_alphanumeric() || *c == '_')
21                        .collect();
22                    if base
23                        .chars()
24                        .next()
25                        .map(|c| c.is_uppercase())
26                        .unwrap_or(false)
27                    {
28                        types.insert(base);
29                    }
30                }
31            }
32        }
33    }
34    types
35}
36
37fn jv_text(source: &[u8], node: &TsNode<'_>) -> String {
38    std::str::from_utf8(&source[node.start_byte()..node.end_byte()])
39        .unwrap_or("")
40        .trim()
41        .to_string()
42}
43
44fn jv_javadoc(node: &TsNode<'_>, source: &[u8]) -> Option<String> {
45    node.prev_named_sibling()
46        .filter(|s| s.kind() == "block_comment")
47        .and_then(|s| {
48            let text = std::str::from_utf8(&source[s.start_byte()..s.end_byte()]).ok()?;
49            if text.starts_with("/**") {
50                strip_docstring(text)
51            } else {
52                None
53            }
54        })
55}
56
57fn jv_node(id: String, label: String, file_type: &str, path: &Path) -> Node {
58    Node {
59        id,
60        label,
61        file_type: file_type.to_string(),
62        source_file: path.to_string_lossy().to_string(),
63        source_location: None,
64        community: None,
65        rationale: None,
66        docstring: None,
67        metadata: HashMap::new(),
68    }
69}
70
71impl TsJavaExtractor {
72    pub fn extract(source: &[u8], path: &Path) -> Result<ExtractionFragment> {
73        let (file_id, _, file_node) = make_file_node(path);
74        let mut fragment = ExtractionFragment {
75            nodes: vec![file_node],
76            edges: vec![],
77        };
78
79        let lang = tree_sitter_java::LANGUAGE.into();
80        let mut parser = Parser::new();
81        parser
82            .set_language(&lang)
83            .map_err(|e| CodeSynapseError::Parse(format!("java lang: {}", e)))?;
84
85        let tree = parser
86            .parse(source, None)
87            .ok_or_else(|| CodeSynapseError::Parse("java parse failed".to_string()))?;
88
89        let stem = file_id.clone();
90
91        Self::walk(
92            tree.root_node(),
93            source,
94            path,
95            &file_id,
96            &stem,
97            &mut fragment,
98        );
99
100        Ok(fragment)
101    }
102
103    fn walk(
104        node: TsNode<'_>,
105        source: &[u8],
106        path: &Path,
107        file_id: &str,
108        stem: &str,
109        fragment: &mut ExtractionFragment,
110    ) {
111        match node.kind() {
112            "method_declaration" => {
113                // Top-level method (outside any class) — use file_id as parent
114                Self::extract_method(node, source, path, file_id, stem, file_id, fragment);
115            }
116            "import_declaration" => {
117                // scoped_identifier or identifier child
118                for i in 0..node.child_count() {
119                    if let Some(child) = node.child(i) {
120                        if child.kind() == "scoped_identifier" || child.kind() == "identifier" {
121                            let full = jv_text(source, &child);
122                            let parts: Vec<&str> = full.split('.').collect();
123                            // Emit import edge to leaf name
124                            let leaf = parts.last().copied().unwrap_or(&full);
125                            let leaf_id = make_id(&[leaf]);
126                            add_node_if_missing(
127                                fragment,
128                                jv_node(leaf_id.clone(), leaf.to_string(), "code", path),
129                            );
130                            fragment.edges.push(Edge {
131                                source: file_id.to_string(),
132                                target: leaf_id,
133                                relation: "imports".to_string(),
134                                confidence: "EXTRACTED".to_string(),
135                                source_file: Some(path.to_string_lossy().to_string()),
136                                weight: 1.0,
137                                context: Some("import".to_string()),
138                            });
139                        }
140                    }
141                }
142            }
143            "class_declaration" => {
144                let name_node = node.child_by_field_name("name");
145                let class_name = match name_node {
146                    Some(n) => jv_text(source, &n),
147                    None => return,
148                };
149                let class_id = make_id(&[stem, &class_name]);
150                let mut class_node = jv_node(class_id.clone(), class_name.clone(), "class", path);
151                class_node.docstring = jv_javadoc(&node, source);
152                class_node.source_location =
153                    Some(format!("{}:{}", node.start_byte(), node.end_byte()));
154                add_node_if_missing(fragment, class_node);
155                fragment.edges.push(Edge {
156                    source: file_id.to_string(),
157                    target: class_id.clone(),
158                    relation: "contains".to_string(),
159                    confidence: "EXTRACTED".to_string(),
160                    source_file: Some(path.to_string_lossy().to_string()),
161                    weight: 1.0,
162                    context: None,
163                });
164
165                // superclass → inherits
166                if let Some(sc) = node.child_by_field_name("superclass") {
167                    for i in 0..sc.child_count() {
168                        if let Some(t) = sc.child(i) {
169                            if t.kind() == "type_identifier" {
170                                let super_name = jv_text(source, &t);
171                                let super_id = make_id(&[stem, &super_name]);
172                                add_node_if_missing(
173                                    fragment,
174                                    jv_node(super_id.clone(), super_name, "code", path),
175                                );
176                                fragment.edges.push(Edge {
177                                    source: class_id.clone(),
178                                    target: super_id,
179                                    relation: "inherits".to_string(),
180                                    confidence: "EXTRACTED".to_string(),
181                                    source_file: Some(path.to_string_lossy().to_string()),
182                                    weight: 1.0,
183                                    context: None,
184                                });
185                            }
186                        }
187                    }
188                }
189
190                // super_interfaces → implements
191                if let Some(si) = node.child_by_field_name("interfaces") {
192                    Self::collect_type_identifiers_as(
193                        &si,
194                        source,
195                        &class_id,
196                        path,
197                        fragment,
198                        "implements",
199                        stem,
200                    );
201                }
202
203                // Walk body
204                if let Some(body) = node.child_by_field_name("body") {
205                    for type_name in collect_field_call_targets(&body, source) {
206                        let target_id = make_id(&[&type_name, &type_name]);
207                        fragment.edges.push(Edge {
208                            source: class_id.clone(),
209                            target: target_id,
210                            relation: "calls".to_string(),
211                            confidence: "EXTRACTED".to_string(),
212                            source_file: Some(path.to_string_lossy().to_string()),
213                            weight: 1.0,
214                            context: Some("field_injection".to_string()),
215                        });
216                    }
217                    for i in 0..body.child_count() {
218                        if let Some(child) = body.child(i) {
219                            if child.kind() == "method_declaration" {
220                                Self::extract_method(
221                                    child, source, path, file_id, stem, &class_id, fragment,
222                                );
223                            } else {
224                                Self::walk(child, source, path, file_id, stem, fragment);
225                            }
226                        }
227                    }
228                }
229            }
230            "interface_declaration" => {
231                let name_node = node.child_by_field_name("name");
232                if let Some(n) = name_node {
233                    let name = jv_text(source, &n);
234                    let id = make_id(&[stem, &name]);
235                    add_node_if_missing(fragment, jv_node(id.clone(), name, "interface", path));
236                    fragment.edges.push(Edge {
237                        source: file_id.to_string(),
238                        target: id,
239                        relation: "contains".to_string(),
240                        confidence: "EXTRACTED".to_string(),
241                        source_file: Some(path.to_string_lossy().to_string()),
242                        weight: 1.0,
243                        context: None,
244                    });
245                }
246            }
247            _ => {
248                for i in 0..node.child_count() {
249                    if let Some(child) = node.child(i) {
250                        Self::walk(child, source, path, file_id, stem, fragment);
251                    }
252                }
253            }
254        }
255    }
256
257    fn collect_type_identifiers_as(
258        node: &TsNode<'_>,
259        source: &[u8],
260        source_id: &str,
261        path: &Path,
262        fragment: &mut ExtractionFragment,
263        relation: &str,
264        stem: &str,
265    ) {
266        for i in 0..node.child_count() {
267            if let Some(child) = node.child(i) {
268                if child.kind() == "type_identifier" {
269                    let name = jv_text(source, &child);
270                    let id = make_id(&[stem, &name]);
271                    add_node_if_missing(fragment, jv_node(id.clone(), name, "code", path));
272                    fragment.edges.push(Edge {
273                        source: source_id.to_string(),
274                        target: id,
275                        relation: relation.to_string(),
276                        confidence: "EXTRACTED".to_string(),
277                        source_file: Some(path.to_string_lossy().to_string()),
278                        weight: 1.0,
279                        context: None,
280                    });
281                } else {
282                    Self::collect_type_identifiers_as(
283                        &child, source, source_id, path, fragment, relation, stem,
284                    );
285                }
286            }
287        }
288    }
289
290    fn extract_method(
291        node: TsNode<'_>,
292        source: &[u8],
293        path: &Path,
294        file_id: &str,
295        stem: &str,
296        class_id: &str,
297        fragment: &mut ExtractionFragment,
298    ) {
299        let name_node = match node.child_by_field_name("name") {
300            Some(n) => n,
301            None => return,
302        };
303        let method_name = jv_text(source, &name_node);
304        let method_id = make_id(&[stem, &method_name, "()"]);
305        let method_label = format!("{}()", method_name);
306        let mut method_node = jv_node(method_id.clone(), method_label, "method", path);
307        method_node.docstring = jv_javadoc(&node, source);
308        method_node.source_location = Some(format!("{}:{}", node.start_byte(), node.end_byte()));
309        add_node_if_missing(fragment, method_node);
310        fragment.edges.push(Edge {
311            source: class_id.to_string(),
312            target: method_id.clone(),
313            relation: "method".to_string(),
314            confidence: "EXTRACTED".to_string(),
315            source_file: Some(path.to_string_lossy().to_string()),
316            weight: 1.0,
317            context: None,
318        });
319        fragment.edges.push(Edge {
320            source: file_id.to_string(),
321            target: method_id.clone(),
322            relation: "contains".to_string(),
323            confidence: "EXTRACTED".to_string(),
324            source_file: Some(path.to_string_lossy().to_string()),
325            weight: 1.0,
326            context: None,
327        });
328
329        // Annotations → attribute context
330        // `modifiers` is not a named field in tree-sitter-java; scan children by kind
331        let mods_opt =
332            (0..node.child_count()).find_map(|i| node.child(i).filter(|c| c.kind() == "modifiers"));
333        if let Some(mods) = mods_opt {
334            for i in 0..mods.child_count() {
335                if let Some(annot) = mods.child(i) {
336                    if annot.kind() == "annotation" || annot.kind() == "marker_annotation" {
337                        if let Some(name_n) = annot.child_by_field_name("name") {
338                            let annot_name = jv_text(source, &name_n);
339                            let annot_id = make_id(&[&annot_name]);
340                            add_node_if_missing(
341                                fragment,
342                                jv_node(annot_id.clone(), annot_name, "code", path),
343                            );
344                            // Also emit full-text "@Name" node for compatibility
345                            let full_text = jv_text(source, &annot);
346                            let full_id = make_id(&[&full_text]);
347                            add_node_if_missing(
348                                fragment,
349                                jv_node(full_id, full_text, "code", path),
350                            );
351                            fragment.edges.push(Edge {
352                                source: method_id.clone(),
353                                target: annot_id,
354                                relation: "references".to_string(),
355                                confidence: "EXTRACTED".to_string(),
356                                source_file: Some(path.to_string_lossy().to_string()),
357                                weight: 1.0,
358                                context: Some("attribute".to_string()),
359                            });
360                        }
361                    }
362                }
363            }
364        }
365
366        // Return type
367        if let Some(ret_type) = node.child_by_field_name("type") {
368            Self::emit_type_ref(
369                &ret_type,
370                source,
371                path,
372                &method_id,
373                fragment,
374                "return_type",
375                stem,
376                true,
377            );
378        }
379
380        // Parameters → parameter_type context
381        if let Some(params) = node.child_by_field_name("parameters") {
382            for i in 0..params.child_count() {
383                if let Some(param) = params.child(i) {
384                    if param.kind() == "formal_parameter" || param.kind() == "spread_parameter" {
385                        if let Some(param_type) = param.child_by_field_name("type") {
386                            Self::emit_type_ref(
387                                &param_type,
388                                source,
389                                path,
390                                &method_id,
391                                fragment,
392                                "parameter_type",
393                                stem,
394                                false,
395                            );
396                        }
397                    }
398                }
399            }
400        }
401    }
402
403    #[allow(clippy::too_many_arguments)]
404    fn emit_type_ref(
405        type_node: &TsNode<'_>,
406        source: &[u8],
407        path: &Path,
408        method_id: &str,
409        fragment: &mut ExtractionFragment,
410        context: &str,
411        _stem: &str,
412        emit_generics: bool,
413    ) {
414        match type_node.kind() {
415            "type_identifier" => {
416                let name = jv_text(source, type_node);
417                let id = make_id(&[&name]);
418                add_node_if_missing(fragment, jv_node(id.clone(), name, "code", path));
419                fragment.edges.push(Edge {
420                    source: method_id.to_string(),
421                    target: id,
422                    relation: "references".to_string(),
423                    confidence: "EXTRACTED".to_string(),
424                    source_file: Some(path.to_string_lossy().to_string()),
425                    weight: 1.0,
426                    context: Some(context.to_string()),
427                });
428            }
429            "generic_type" => {
430                // tree-sitter-java generic_type has no named fields; scan children by kind
431                let name_n = (0..type_node.child_count())
432                    .find_map(|i| type_node.child(i).filter(|c| c.kind() == "type_identifier"));
433                if let Some(name_n) = name_n {
434                    let name = jv_text(source, &name_n);
435                    let id = make_id(&[&name]);
436                    add_node_if_missing(fragment, jv_node(id.clone(), name, "code", path));
437                    fragment.edges.push(Edge {
438                        source: method_id.to_string(),
439                        target: id,
440                        relation: "references".to_string(),
441                        confidence: "EXTRACTED".to_string(),
442                        source_file: Some(path.to_string_lossy().to_string()),
443                        weight: 1.0,
444                        context: Some(context.to_string()),
445                    });
446                }
447                // Emit generic args from type_arguments child
448                if emit_generics {
449                    let args = (0..type_node.child_count())
450                        .find_map(|i| type_node.child(i).filter(|c| c.kind() == "type_arguments"));
451                    if let Some(args) = args {
452                        for i in 0..args.child_count() {
453                            if let Some(arg) = args.child(i) {
454                                if arg.kind() == "type_identifier" {
455                                    let name = jv_text(source, &arg);
456                                    let id = make_id(&[&name]);
457                                    add_node_if_missing(
458                                        fragment,
459                                        jv_node(id.clone(), name, "code", path),
460                                    );
461                                    fragment.edges.push(Edge {
462                                        source: method_id.to_string(),
463                                        target: id,
464                                        relation: "references".to_string(),
465                                        confidence: "EXTRACTED".to_string(),
466                                        source_file: Some(path.to_string_lossy().to_string()),
467                                        weight: 1.0,
468                                        context: Some("generic_arg".to_string()),
469                                    });
470                                }
471                            }
472                        }
473                    }
474                }
475            }
476            _ => {}
477        }
478    }
479}
480
481impl LanguageExtractor for TsJavaExtractor {
482    fn file_extensions(&self) -> Vec<&'static str> {
483        vec!["java"]
484    }
485    fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
486        Self::extract(source, path)
487    }
488    fn resolve_imports(&self, _imports: &[ImportNode]) -> Vec<Edge> {
489        vec![]
490    }
491    fn collect_type_refs(&self, _fragment: &mut ExtractionFragment) {}
492}