Skip to main content

forgekit_core/treesitter/
java.rs

1use crate::cfg::TestCfg;
2use crate::error::{ForgeError, Result};
3use crate::types::BlockId;
4
5use super::{CfgExtractor, FunctionInfo, SupportedLanguage};
6
7impl CfgExtractor {
8    /// Extract CFG from Java source code
9    pub fn extract_java(source: &str) -> Result<Vec<FunctionInfo>> {
10        use tree_sitter::Parser;
11        use tree_sitter_java;
12
13        let mut parser = Parser::new();
14        parser
15            .set_language(&tree_sitter_java::LANGUAGE.into())
16            .map_err(|e| {
17                ForgeError::DatabaseError(format!("Failed to set Java language: {:?}", e))
18            })?;
19
20        let tree = parser
21            .parse(source, None)
22            .ok_or_else(|| ForgeError::DatabaseError("Failed to parse Java code".to_string()))?;
23
24        let root = tree.root_node();
25        let mut functions = Vec::new();
26
27        Self::extract_java_functions(source, &root, &mut functions)?;
28
29        Ok(functions)
30    }
31
32    fn extract_java_functions(
33        source: &str,
34        node: &tree_sitter::Node,
35        functions: &mut Vec<FunctionInfo>,
36    ) -> Result<()> {
37        let kind = node.kind();
38
39        // Look for method declarations
40        if kind == "method_declaration" {
41            if let Some(func) = Self::parse_java_function(source, node)? {
42                functions.push(func);
43            }
44        }
45
46        // Recurse into children
47        let mut cursor = node.walk();
48        for child in node.children(&mut cursor) {
49            Self::extract_java_functions(source, &child, functions)?;
50        }
51
52        Ok(())
53    }
54
55    fn parse_java_function(source: &str, node: &tree_sitter::Node) -> Result<Option<FunctionInfo>> {
56        let start_byte = node.start_byte();
57        let end_byte = node.end_byte();
58
59        // Find method name
60        let mut name = "unknown".to_string();
61        let mut cursor = node.walk();
62        for child in node.children(&mut cursor) {
63            if child.kind() == "identifier" {
64                name = Self::node_text(source, &child);
65                break;
66            }
67        }
68
69        // Find method body (block)
70        let mut body = None;
71        let mut cursor = node.walk();
72        for child in node.children(&mut cursor) {
73            if child.kind() == "block" {
74                body = Some(child);
75                break;
76            }
77        }
78
79        let cfg = if let Some(body) = body {
80            Self::build_cfg_from_body(source, &body, SupportedLanguage::Java)?
81        } else {
82            // Abstract method without body
83            TestCfg::new(BlockId(0))
84        };
85
86        Ok(Some(FunctionInfo {
87            name,
88            start_byte,
89            end_byte,
90            cfg,
91        }))
92    }
93}