vtcode_core/code/code_completion/context/
analyzer.rs1use super::CompletionContext;
2use crate::tools::tree_sitter::TreeSitterAnalyzer;
3pub struct ContextAnalyzer {
7 tree_sitter: TreeSitterAnalyzer,
8}
9
10impl ContextAnalyzer {
11 pub fn new() -> Self {
12 Self {
13 tree_sitter: TreeSitterAnalyzer::new().expect("Failed to initialize TreeSitter"),
14 }
15 }
16
17 pub fn analyze(&mut self, source: &str, line: usize, column: usize) -> CompletionContext {
19 let language = self.detect_language(source);
20 let prefix = self.extract_prefix(source, line, column);
21
22 let mut context = CompletionContext::new(line, column, prefix, language);
23 context.scope = self.extract_scope(source, line, column);
24 context.imports = self.extract_imports(source);
25 context.recent_symbols = self.extract_recent_symbols(source, line);
26
27 context
28 }
29
30 fn detect_language(&self, source: &str) -> String {
31 if let Some(language) = self.tree_sitter.detect_language_from_content(source) {
33 return match language {
34 crate::tools::tree_sitter::LanguageSupport::Rust => "rust".to_string(),
35 crate::tools::tree_sitter::LanguageSupport::Python => "python".to_string(),
36 crate::tools::tree_sitter::LanguageSupport::JavaScript => "javascript".to_string(),
37 crate::tools::tree_sitter::LanguageSupport::TypeScript => "typescript".to_string(),
38 crate::tools::tree_sitter::LanguageSupport::Go => "go".to_string(),
39 crate::tools::tree_sitter::LanguageSupport::Java => "java".to_string(),
40 crate::tools::tree_sitter::LanguageSupport::Swift => "swift".to_string(),
41 };
42 }
43
44 "rust".to_string()
46 }
47
48 fn extract_prefix(&self, source: &str, line: usize, column: usize) -> String {
49 let lines: Vec<&str> = source.lines().collect();
50 if line < lines.len() && column <= lines[line].len() {
51 lines[line][..column].to_string()
52 } else {
53 String::new()
54 }
55 }
56
57 fn extract_scope(&mut self, source: &str, line: usize, column: usize) -> Vec<String> {
58 let language = self.detect_language(source);
59
60 let lang_support = match language.as_str() {
62 "rust" => crate::tools::tree_sitter::LanguageSupport::Rust,
63 "python" => crate::tools::tree_sitter::LanguageSupport::Python,
64 "javascript" => crate::tools::tree_sitter::LanguageSupport::JavaScript,
65 "typescript" => crate::tools::tree_sitter::LanguageSupport::TypeScript,
66 "go" => crate::tools::tree_sitter::LanguageSupport::Go,
67 "java" => crate::tools::tree_sitter::LanguageSupport::Java,
68 _ => crate::tools::tree_sitter::LanguageSupport::Rust,
69 };
70
71 if let Ok(tree) = self.tree_sitter.parse(source, lang_support) {
73 let root_node = tree.root_node();
74 let mut scopes = Vec::new();
75
76 if let Some(node) = Self::find_node_at_position(root_node, line, column) {
78 let mut current = Some(node);
80 while let Some(n) = current {
81 let kind = n.kind();
82
83 if kind.contains("function") || kind.contains("method") {
85 scopes.push(format!("function:{}", kind));
86 } else if kind.contains("class") || kind.contains("struct") {
87 scopes.push(format!("class:{}", kind));
88 } else if kind.contains("module") || kind.contains("namespace") {
89 scopes.push(format!("module:{}", kind));
90 }
91
92 current = n.parent();
93 }
94 }
95
96 scopes
97 } else {
98 vec![]
99 }
100 }
101
102 fn extract_imports(&mut self, source: &str) -> Vec<String> {
104 let language = self.detect_language(source);
105
106 let lang_support = match language.as_str() {
108 "rust" => crate::tools::tree_sitter::LanguageSupport::Rust,
109 "python" => crate::tools::tree_sitter::LanguageSupport::Python,
110 "javascript" => crate::tools::tree_sitter::LanguageSupport::JavaScript,
111 "typescript" => crate::tools::tree_sitter::LanguageSupport::TypeScript,
112 "go" => crate::tools::tree_sitter::LanguageSupport::Go,
113 "java" => crate::tools::tree_sitter::LanguageSupport::Java,
114 _ => crate::tools::tree_sitter::LanguageSupport::Rust,
115 };
116
117 if let Ok(tree) = self.tree_sitter.parse(source, lang_support) {
118 let root_node = tree.root_node();
119 let mut imports = Vec::new();
120
121 Self::extract_imports_recursive(root_node, source, &lang_support, &mut imports);
123
124 imports
125 } else {
126 vec![]
127 }
128 }
129
130 fn extract_imports_recursive(
132 node: tree_sitter::Node,
133 source: &str,
134 language: &crate::tools::tree_sitter::LanguageSupport,
135 imports: &mut Vec<String>,
136 ) {
137 let kind = node.kind();
138
139 match language {
141 crate::tools::tree_sitter::LanguageSupport::Rust => {
142 if kind == "use_declaration" {
143 let import_text = &source[node.start_byte()..node.end_byte()];
144 imports.push(import_text.to_string());
145 }
146 }
147 crate::tools::tree_sitter::LanguageSupport::Python => {
148 if kind == "import_statement" || kind == "import_from_statement" {
149 let import_text = &source[node.start_byte()..node.end_byte()];
150 imports.push(import_text.to_string());
151 }
152 }
153 crate::tools::tree_sitter::LanguageSupport::JavaScript
154 | crate::tools::tree_sitter::LanguageSupport::TypeScript => {
155 if kind == "import_statement" {
156 let import_text = &source[node.start_byte()..node.end_byte()];
157 imports.push(import_text.to_string());
158 }
159 }
160 _ => {
161 if kind.contains("import") || kind.contains("require") {
163 let import_text = &source[node.start_byte()..node.end_byte()];
164 imports.push(import_text.to_string());
165 }
166 }
167 }
168
169 let mut cursor = node.walk();
171 for child in node.children(&mut cursor) {
172 Self::extract_imports_recursive(child, source, language, imports);
173 }
174 }
175
176 fn extract_recent_symbols(&mut self, source: &str, line: usize) -> Vec<String> {
178 let language = self.detect_language(source);
179
180 let lang_support = match language.as_str() {
182 "rust" => crate::tools::tree_sitter::LanguageSupport::Rust,
183 "python" => crate::tools::tree_sitter::LanguageSupport::Python,
184 "javascript" => crate::tools::tree_sitter::LanguageSupport::JavaScript,
185 "typescript" => crate::tools::tree_sitter::LanguageSupport::TypeScript,
186 "go" => crate::tools::tree_sitter::LanguageSupport::Go,
187 "java" => crate::tools::tree_sitter::LanguageSupport::Java,
188 _ => crate::tools::tree_sitter::LanguageSupport::Rust,
189 };
190
191 if let Ok(tree) = self.tree_sitter.parse(source, lang_support) {
193 let _root_node = tree.root_node();
194 let mut symbols = Vec::new();
195
196 if let Ok(extracted_symbols) =
198 self.tree_sitter
199 .extract_symbols(&tree, source, lang_support)
200 {
201 for symbol in extracted_symbols {
203 if symbol.position.row < line {
204 symbols.push(symbol.name);
205 }
206 }
207 }
208
209 if symbols.len() > 10 {
211 symbols[symbols.len() - 10..].to_vec()
212 } else {
213 symbols
214 }
215 } else {
216 vec![]
217 }
218 }
219
220 fn find_node_at_position<'a>(
222 node: tree_sitter::Node<'a>,
223 line: usize,
224 column: usize,
225 ) -> Option<tree_sitter::Node<'a>> {
226 let start_pos = node.start_position();
227 let end_pos = node.end_position();
228
229 if start_pos.row <= line && end_pos.row >= line {
231 if start_pos.row == line && start_pos.column > column {
232 return None;
233 }
234 if end_pos.row == line && end_pos.column < column {
235 return None;
236 }
237
238 let mut cursor = node.walk();
240 for child in node.children(&mut cursor) {
241 if let Some(found) = Self::find_node_at_position(child, line, column) {
242 return Some(found);
243 }
244 }
245
246 return Some(node);
248 }
249
250 None
251 }
252}
253
254impl Default for ContextAnalyzer {
255 fn default() -> Self {
256 Self::new()
257 }
258}