1use crate::trace::{CallNode, CallTree, TraceDirection};
2use crate::tree::{NodeType, ReferenceTree, TreeNode};
3use crate::SearchResult;
4use colored::*;
5
6pub struct TreeFormatter {
8 max_width: usize,
9}
10
11impl TreeFormatter {
12 pub fn new() -> Self {
14 colored::control::set_override(true);
16 Self { max_width: 80 }
17 }
18
19 pub fn with_width(max_width: usize) -> Self {
21 Self { max_width }
22 }
23
24 pub fn format_result(&self, result: &SearchResult) -> String {
26 let mut output = String::new();
27
28 if !result.translation_entries.is_empty() {
30 output.push_str(&format!("{}\n", "=== Translation Files ===".bold()));
31 for entry in &result.translation_entries {
32 output.push_str(&format!(
33 "{}:{}:{}: {}\n",
34 entry.file.display(),
35 entry.line,
36 entry.key.yellow().bold(),
37 format!("\"{}\"", entry.value).green().bold()
38 ));
39 }
40 output.push('\n');
41 }
42
43 if !result.code_references.is_empty() {
45 output.push_str(&format!("{}\n", "=== Code References ===".bold()));
46 for code_ref in &result.code_references {
47 let highlighted_context =
49 self.highlight_key_in_context(&code_ref.context, &code_ref.key_path);
50 output.push_str(&format!(
51 "{}:{}:{}\n",
52 code_ref.file.display(),
53 code_ref.line,
54 highlighted_context
55 ));
56 }
57 }
58
59 output
60 }
61
62 fn highlight_key_in_context(&self, context: &str, key: &str) -> String {
64 context.replace(key, &key.bold().to_string())
66 }
67
68 pub fn format(&self, tree: &ReferenceTree) -> String {
70 let mut output = String::new();
71 self.format_node(&tree.root, &mut output, "", true, true);
72 output
73 }
74
75 pub fn format_trace_tree(&self, tree: &CallTree, direction: TraceDirection) -> String {
76 match direction {
77 TraceDirection::Forward => self.format_forward_tree(tree),
78 TraceDirection::Backward => self.format_backward_tree(tree),
79 }
80 }
81
82 fn format_forward_tree(&self, tree: &CallTree) -> String {
83 let mut output = String::new();
84 Self::format_call_node(&tree.root, &mut output, "", true, true);
85 output
86 }
87
88 fn format_backward_tree(&self, tree: &CallTree) -> String {
89 let mut output = String::new();
90 let mut paths = Vec::new();
101 Self::collect_backward_paths(&tree.root, vec![], &mut paths);
102
103 for path in paths {
104 let mut display_path = path.clone();
116 display_path.reverse();
117
118 let mut chain = display_path
119 .iter()
120 .map(|node| {
121 format!(
122 "{} ({}:{})",
123 node.def.name.bold(),
124 node.def.file.display(),
125 node.def.line
126 )
127 })
128 .collect::<Vec<_>>()
129 .join(" -> ");
130
131 if let Some(first) = display_path.first() {
133 if first.truncated {
134 chain = format!("{} -> {}", "[depth limit reached]".red(), chain);
135 }
136 }
137
138 output.push_str(&chain);
139 output.push('\n');
140 }
141
142 if output.is_empty() {
143 output.push_str(&format!(
145 "{} (No incoming calls found)\n",
146 tree.root.def.name
147 ));
148 }
149
150 output
151 }
152
153 fn collect_backward_paths<'a>(
154 node: &'a CallNode,
155 mut current_path: Vec<&'a CallNode>,
156 paths: &mut Vec<Vec<&'a CallNode>>,
157 ) {
158 current_path.push(node);
159
160 if node.children.is_empty() {
161 if node.truncated {
165 }
169 paths.push(current_path);
170 } else {
171 for child in &node.children {
172 Self::collect_backward_paths(child, current_path.clone(), paths);
173 }
174 }
175 }
176
177 fn format_call_node(
178 node: &CallNode,
179 output: &mut String,
180 prefix: &str,
181 is_last: bool,
182 is_root: bool,
183 ) {
184 if !is_root {
185 output.push_str(prefix);
186 output.push_str(if is_last { "└─> " } else { "├─> " });
187 }
188
189 let content = format!(
190 "{} ({}:{})",
191 node.def.name.bold(),
192 node.def.file.display(),
193 node.def.line
194 );
195 output.push_str(&content);
196
197 if node.truncated {
198 output.push_str(&" [depth limit reached]".red().to_string());
199 }
200
201 output.push('\n');
202
203 let child_count = node.children.len();
204 for (i, child) in node.children.iter().enumerate() {
205 let is_last_child = i == child_count - 1;
206 let child_prefix = if is_root {
207 String::new()
208 } else {
209 format!("{}{} ", prefix, if is_last { " " } else { "│" })
210 };
211 Self::format_call_node(child, output, &child_prefix, is_last_child, false);
212 }
213 }
214
215 fn format_node(
217 &self,
218 node: &TreeNode,
219 output: &mut String,
220 prefix: &str,
221 is_last: bool,
222 is_root: bool,
223 ) {
224 if !is_root {
226 output.push_str(prefix);
227 output.push_str(if is_last { "└─> " } else { "├─> " });
228 }
229
230 let content = self.format_content(node);
232 output.push_str(&content);
233
234 if let Some(location) = &node.location {
236 let location_str = format!(" ({}:{})", location.file.display(), location.line);
237 output.push_str(&location_str);
238 }
239
240 output.push('\n');
241
242 let child_count = node.children.len();
244 for (i, child) in node.children.iter().enumerate() {
245 let is_last_child = i == child_count - 1;
246 let child_prefix = if is_root {
247 String::new()
248 } else {
249 format!("{}{} ", prefix, if is_last { " " } else { "│" })
250 };
251
252 self.format_node(child, output, &child_prefix, is_last_child, false);
253 }
254 }
255
256 fn format_content(&self, node: &TreeNode) -> String {
258 match node.node_type {
259 NodeType::Root => {
260 format!("'{}' (search query)", node.content)
261 }
262 NodeType::Translation => {
263 self.truncate(&node.content, self.max_width - 20)
265 }
266 NodeType::KeyPath => {
267 format!("Key: {}", node.content)
268 }
269 NodeType::CodeRef => {
270 let truncated = self.truncate(node.content.trim(), self.max_width - 30);
272 format!("Code: {}", truncated)
273 }
274 }
275 }
276
277 fn truncate(&self, s: &str, max_len: usize) -> String {
279 if s.len() <= max_len {
280 s.to_string()
281 } else {
282 format!("{}...", &s[..max_len.saturating_sub(3)])
283 }
284 }
285}
286
287impl Default for TreeFormatter {
288 fn default() -> Self {
289 Self::new()
290 }
291}
292
293#[cfg(test)]
294mod tests {
295 use super::*;
296 use crate::tree::{Location, TreeNode};
297 use std::path::PathBuf;
298
299 #[test]
300 fn test_formatter_creation() {
301 let formatter = TreeFormatter::new();
302 assert_eq!(formatter.max_width, 80);
303 }
304
305 #[test]
306 fn test_formatter_with_custom_width() {
307 let formatter = TreeFormatter::with_width(120);
308 assert_eq!(formatter.max_width, 120);
309 }
310
311 #[test]
312 fn test_format_empty_tree() {
313 let tree = ReferenceTree::with_search_text("test".to_string());
314 let formatter = TreeFormatter::new();
315 let output = formatter.format(&tree);
316
317 assert!(output.contains("'test'"));
318 assert!(output.contains("search query"));
319 }
320
321 #[test]
322 fn test_format_tree_with_translation() {
323 let mut root = TreeNode::new(NodeType::Root, "add new".to_string());
324 let translation = TreeNode::with_location(
325 NodeType::Translation,
326 "invoice.labels.add_new: 'add new'".to_string(),
327 Location::new(PathBuf::from("en.yml"), 4),
328 );
329 root.add_child(translation);
330
331 let tree = ReferenceTree::new(root);
332 let formatter = TreeFormatter::new();
333 let output = formatter.format(&tree);
334
335 assert!(output.contains("'add new'"));
336 assert!(output.contains("invoice.labels.add_new"));
337 assert!(output.contains("en.yml:4"));
338 assert!(output.contains("└─>") || output.contains("├─>"));
339 }
340
341 #[test]
342 fn test_format_complete_tree() {
343 let mut root = TreeNode::new(NodeType::Root, "add new".to_string());
344
345 let mut translation = TreeNode::with_location(
346 NodeType::Translation,
347 "invoice.labels.add_new: 'add new'".to_string(),
348 Location::new(PathBuf::from("en.yml"), 4),
349 );
350
351 let mut key_path = TreeNode::new(NodeType::KeyPath, "invoice.labels.add_new".to_string());
352
353 let code_ref = TreeNode::with_location(
354 NodeType::CodeRef,
355 "I18n.t('invoice.labels.add_new')".to_string(),
356 Location::new(PathBuf::from("invoices.ts"), 14),
357 );
358
359 key_path.add_child(code_ref);
360 translation.add_child(key_path);
361 root.add_child(translation);
362
363 let tree = ReferenceTree::new(root);
364 let formatter = TreeFormatter::new();
365 let output = formatter.format(&tree);
366
367 assert!(output.contains("'add new'"));
369 assert!(output.contains("invoice.labels.add_new"));
370 assert!(output.contains("Key:"));
371 assert!(output.contains("Code:"));
372 assert!(output.contains("I18n.t"));
373 assert!(output.contains("en.yml:4"));
374 assert!(output.contains("invoices.ts:14"));
375 }
376
377 #[test]
378 fn test_format_multiple_children() {
379 let mut root = TreeNode::new(NodeType::Root, "test".to_string());
380
381 let child1 = TreeNode::with_location(
382 NodeType::Translation,
383 "key1: 'value1'".to_string(),
384 Location::new(PathBuf::from("file1.yml"), 1),
385 );
386
387 let child2 = TreeNode::with_location(
388 NodeType::Translation,
389 "key2: 'value2'".to_string(),
390 Location::new(PathBuf::from("file2.yml"), 2),
391 );
392
393 root.add_child(child1);
394 root.add_child(child2);
395
396 let tree = ReferenceTree::new(root);
397 let formatter = TreeFormatter::new();
398 let output = formatter.format(&tree);
399
400 assert!(output.contains("key1"));
402 assert!(output.contains("key2"));
403 assert!(output.contains("file1.yml:1"));
404 assert!(output.contains("file2.yml:2"));
405
406 assert!(output.contains("├─>"));
408 assert!(output.contains("└─>"));
409 }
410
411 #[test]
412 fn test_truncate_long_content() {
413 let formatter = TreeFormatter::with_width(50);
414 let long_string = "a".repeat(100);
415 let truncated = formatter.truncate(&long_string, 20);
416
417 assert!(truncated.len() <= 20);
418 assert!(truncated.ends_with("..."));
419 }
420
421 #[test]
422 fn test_truncate_short_content() {
423 let formatter = TreeFormatter::new();
424 let short_string = "short";
425 let result = formatter.truncate(short_string, 20);
426
427 assert_eq!(result, "short");
428 }
429
430 #[test]
431 fn test_format_content_root() {
432 let formatter = TreeFormatter::new();
433 let node = TreeNode::new(NodeType::Root, "test query".to_string());
434 let content = formatter.format_content(&node);
435
436 assert!(content.contains("test query"));
437 assert!(content.contains("search query"));
438 }
439
440 #[test]
441 fn test_format_content_key_path() {
442 let formatter = TreeFormatter::new();
443 let node = TreeNode::new(NodeType::KeyPath, "invoice.labels.add_new".to_string());
444 let content = formatter.format_content(&node);
445
446 assert!(content.contains("Key:"));
447 assert!(content.contains("invoice.labels.add_new"));
448 }
449
450 #[test]
451 fn test_format_content_code_ref() {
452 let formatter = TreeFormatter::new();
453 let node = TreeNode::new(
454 NodeType::CodeRef,
455 " I18n.t('invoice.labels.add_new') ".to_string(),
456 );
457 let content = formatter.format_content(&node);
458
459 assert!(content.contains("Code:"));
460 assert!(content.contains("I18n.t"));
461 assert!(!content.starts_with(" "));
463 }
464
465 #[test]
466 fn test_format_deep_nesting() {
467 let mut root = TreeNode::new(NodeType::Root, "test".to_string());
468 let mut level1 = TreeNode::new(NodeType::Translation, "level1".to_string());
469 let mut level2 = TreeNode::new(NodeType::KeyPath, "level2".to_string());
470 let level3 = TreeNode::new(NodeType::CodeRef, "level3".to_string());
471
472 level2.add_child(level3);
473 level1.add_child(level2);
474 root.add_child(level1);
475
476 let tree = ReferenceTree::new(root);
477 let formatter = TreeFormatter::new();
478 let output = formatter.format(&tree);
479
480 let lines: Vec<&str> = output.lines().collect();
482 assert!(lines.len() >= 4);
483
484 assert!(lines[2].starts_with(' ') || lines[2].starts_with('│'));
486 }
487}