context_creator/core/
context_builder.rs

1//! Context creation functionality for LLM consumption
2
3use crate::core::cache::FileCache;
4use crate::core::walker::FileInfo;
5use crate::utils::file_ext::FileType;
6use anyhow::Result;
7use std::collections::HashMap;
8use std::path::Path;
9use std::sync::Arc;
10
11/// Options for generating context for LLM consumption
12#[derive(Debug, Clone)]
13pub struct ContextOptions {
14    /// Maximum tokens allowed in the output
15    pub max_tokens: Option<usize>,
16    /// Include file tree in output
17    pub include_tree: bool,
18    /// Include token count statistics
19    pub include_stats: bool,
20    /// Group files by type
21    pub group_by_type: bool,
22    /// Sort files by priority
23    pub sort_by_priority: bool,
24    /// Template for file headers
25    pub file_header_template: String,
26    /// Template for the document header
27    pub doc_header_template: String,
28    /// Include table of contents
29    pub include_toc: bool,
30    /// Enable enhanced context with file metadata
31    pub enhanced_context: bool,
32}
33
34impl ContextOptions {
35    /// Create ContextOptions from CLI config
36    pub fn from_config(config: &crate::cli::Config) -> Result<Self> {
37        Ok(ContextOptions {
38            max_tokens: config.get_effective_context_tokens(),
39            include_tree: true,
40            include_stats: true,
41            group_by_type: false,
42            sort_by_priority: true,
43            file_header_template: "## {path}".to_string(),
44            doc_header_template: "# Code Context: {directory}".to_string(),
45            include_toc: true,
46            enhanced_context: config.enhanced_context,
47        })
48    }
49}
50
51impl Default for ContextOptions {
52    fn default() -> Self {
53        ContextOptions {
54            max_tokens: None,
55            include_tree: true,
56            include_stats: true,
57            group_by_type: false,
58            sort_by_priority: true,
59            file_header_template: "## {path}".to_string(),
60            doc_header_template: "# Code Context: {directory}".to_string(),
61            include_toc: true,
62            enhanced_context: false,
63        }
64    }
65}
66
67/// Estimate the total size of the markdown output
68fn estimate_output_size(files: &[FileInfo], options: &ContextOptions, cache: &FileCache) -> usize {
69    let mut size = 0;
70
71    // Document header
72    if !options.doc_header_template.is_empty() {
73        size += options.doc_header_template.len() + 50; // Extra for replacements and newlines
74    }
75
76    // Statistics section
77    if options.include_stats {
78        size += 500; // Estimated size for stats
79        size += files.len() * 50; // For file type listing
80    }
81
82    // File tree
83    if options.include_tree {
84        size += 100; // Headers
85        size += files.len() * 100; // Estimated per-file in tree
86    }
87
88    // Table of contents
89    if options.include_toc {
90        size += 50; // Header
91        size += files.len() * 100; // Per-file TOC entry
92    }
93
94    // File contents
95    for file in files {
96        // Header template
97        size +=
98            options.file_header_template.len() + file.relative_path.to_string_lossy().len() + 20;
99
100        // File content + code fence
101        if let Ok(content) = cache.get_or_load(&file.path) {
102            size += content.len() + 20; // Content + fence markers
103        } else {
104            size += file.size as usize; // Fallback to file size
105        }
106    }
107
108    // Add 20% buffer for formatting and unexpected overhead
109    size + (size / 5)
110}
111
112/// Generate markdown from a list of files
113pub fn generate_markdown(
114    files: Vec<FileInfo>,
115    options: ContextOptions,
116    cache: Arc<FileCache>,
117) -> Result<String> {
118    // Pre-allocate string with estimated capacity
119    let estimated_size = estimate_output_size(&files, &options, &cache);
120    let mut output = String::with_capacity(estimated_size);
121
122    // Add document header
123    if !options.doc_header_template.is_empty() {
124        let header = options.doc_header_template.replace("{directory}", ".");
125        output.push_str(&header);
126        output.push_str("\n\n");
127    }
128
129    // Add statistics if requested
130    if options.include_stats {
131        let stats = generate_statistics(&files);
132        output.push_str(&stats);
133        output.push_str("\n\n");
134    }
135
136    // Add file tree if requested
137    if options.include_tree {
138        let tree = generate_file_tree(&files, &options);
139        output.push_str("## File Structure\n\n");
140        output.push_str("```\n");
141        output.push_str(&tree);
142        output.push_str("```\n\n");
143    }
144
145    // Sort files if requested
146    let mut files = files;
147    if options.sort_by_priority {
148        files.sort_by(|a, b| {
149            b.priority
150                .partial_cmp(&a.priority)
151                .unwrap_or(std::cmp::Ordering::Equal)
152                .then_with(|| a.relative_path.cmp(&b.relative_path))
153        });
154    }
155
156    // Add table of contents if requested
157    if options.include_toc {
158        output.push_str("## Table of Contents\n\n");
159        for file in &files {
160            let anchor = path_to_anchor(&file.relative_path);
161            output.push_str(&format!(
162                "- [{path}](#{anchor})\n",
163                path = file.relative_path.display(),
164                anchor = anchor
165            ));
166        }
167        output.push('\n');
168    }
169
170    // Group files if requested
171    if options.group_by_type {
172        let grouped = group_files_by_type(files);
173        for (file_type, group_files) in grouped {
174            output.push_str(&format!("## {} Files\n\n", file_type_display(&file_type)));
175            for file in group_files {
176                append_file_content(&mut output, &file, &options, &cache)?;
177            }
178        }
179    } else {
180        // Add all files
181        for file in files {
182            append_file_content(&mut output, &file, &options, &cache)?;
183        }
184    }
185
186    Ok(output)
187}
188
189/// Append a single file's content to the output
190fn append_file_content(
191    output: &mut String,
192    file: &FileInfo,
193    options: &ContextOptions,
194    cache: &FileCache,
195) -> Result<()> {
196    // Read file content from cache
197    let content = match cache.get_or_load(&file.path) {
198        Ok(content) => content,
199        Err(e) => {
200            eprintln!(
201                "Warning: Could not read file {}: {}",
202                file.path.display(),
203                e
204            );
205            return Ok(());
206        }
207    };
208
209    // Add file header with optional metadata
210    let path_with_metadata = if options.enhanced_context {
211        format!(
212            "{} ({}, {})",
213            file.relative_path.display(),
214            format_size(file.size),
215            file_type_display(&file.file_type)
216        )
217    } else {
218        file.relative_path.display().to_string()
219    };
220
221    let header = options
222        .file_header_template
223        .replace("{path}", &path_with_metadata);
224    output.push_str(&header);
225    output.push_str("\n\n");
226
227    // Add semantic information if available
228    if !file.imports.is_empty() {
229        output.push_str("Imports: ");
230        let import_names: Vec<String> = file
231            .imports
232            .iter()
233            .map(|p| {
234                let filename = p.file_name().and_then(|n| n.to_str()).unwrap_or("");
235
236                // For Python __init__.py files, use the parent directory name
237                if filename == "__init__.py" {
238                    p.parent()
239                        .and_then(|parent| parent.file_name())
240                        .and_then(|n| n.to_str())
241                        .unwrap_or("unknown")
242                        .to_string()
243                } else {
244                    // Remove common extensions
245                    filename
246                        .strip_suffix(".py")
247                        .or_else(|| filename.strip_suffix(".rs"))
248                        .or_else(|| filename.strip_suffix(".js"))
249                        .or_else(|| filename.strip_suffix(".ts"))
250                        .unwrap_or(filename)
251                        .to_string()
252                }
253            })
254            .collect();
255        output.push_str(&format!("{}\n\n", import_names.join(", ")));
256    }
257
258    if !file.imported_by.is_empty() {
259        output.push_str("Imported by: ");
260        let imported_by_names: Vec<String> = file
261            .imported_by
262            .iter()
263            .map(|p| {
264                p.file_name()
265                    .and_then(|n| n.to_str())
266                    .unwrap_or_else(|| p.to_str().unwrap_or("unknown"))
267                    .to_string()
268            })
269            .collect();
270        output.push_str(&format!("{}\n\n", imported_by_names.join(", ")));
271    }
272
273    if !file.function_calls.is_empty() {
274        output.push_str("Function calls: ");
275        let function_names: Vec<String> = file
276            .function_calls
277            .iter()
278            .map(|fc| {
279                if let Some(module) = &fc.module {
280                    format!("{}.{}", module, fc.name)
281                } else {
282                    fc.name.clone()
283                }
284            })
285            .collect();
286        output.push_str(&format!("{}\n\n", function_names.join(", ")));
287    }
288
289    if !file.type_references.is_empty() {
290        output.push_str("Type references: ");
291        let type_names: Vec<String> = file
292            .type_references
293            .iter()
294            .map(|tr| {
295                if let Some(module) = &tr.module {
296                    format!("{}.{}", module, tr.name)
297                } else {
298                    tr.name.clone()
299                }
300            })
301            .collect();
302        output.push_str(&format!("{}\n\n", type_names.join(", ")));
303    }
304
305    // Add language hint for syntax highlighting
306    let language = get_language_hint(&file.file_type);
307    output.push_str(&format!("```{language}\n"));
308    output.push_str(&content);
309    if !content.ends_with('\n') {
310        output.push('\n');
311    }
312    output.push_str("```\n\n");
313
314    Ok(())
315}
316
317/// Generate statistics about the files
318fn generate_statistics(files: &[FileInfo]) -> String {
319    let total_files = files.len();
320    let total_size: u64 = files.iter().map(|f| f.size).sum();
321
322    // Count by file type
323    let mut type_counts: HashMap<FileType, usize> = HashMap::new();
324    for file in files {
325        *type_counts.entry(file.file_type.clone()).or_insert(0) += 1;
326    }
327
328    // Pre-allocate with estimated capacity
329    let mut stats = String::with_capacity(500 + type_counts.len() * 50);
330    stats.push_str("## Statistics\n\n");
331    stats.push_str(&format!("- Total files: {total_files}\n"));
332    stats.push_str(&format!(
333        "- Total size: {} bytes\n",
334        format_size(total_size)
335    ));
336    stats.push_str("\n### Files by type:\n");
337
338    let mut types: Vec<_> = type_counts.into_iter().collect();
339    types.sort_by_key(|(_, count)| std::cmp::Reverse(*count));
340
341    for (file_type, count) in types {
342        stats.push_str(&format!("- {}: {}\n", file_type_display(&file_type), count));
343    }
344
345    stats
346}
347
348/// Generate a file tree representation
349fn generate_file_tree(files: &[FileInfo], options: &ContextOptions) -> String {
350    use std::collections::{BTreeMap, HashMap};
351
352    #[derive(Default)]
353    struct TreeNode {
354        files: Vec<String>,
355        dirs: BTreeMap<String, TreeNode>,
356    }
357
358    let mut root = TreeNode::default();
359
360    // Create a lookup map from relative path to FileInfo for metadata
361    let file_lookup: HashMap<String, &FileInfo> = files
362        .iter()
363        .map(|f| (f.relative_path.to_string_lossy().to_string(), f))
364        .collect();
365
366    // Build tree structure
367    for file in files {
368        let parts: Vec<_> = file
369            .relative_path
370            .components()
371            .map(|c| c.as_os_str().to_string_lossy().to_string())
372            .collect();
373
374        let mut current = &mut root;
375        for (i, part) in parts.iter().enumerate() {
376            if i == parts.len() - 1 {
377                // File
378                current.files.push(part.clone());
379            } else {
380                // Directory
381                current = current.dirs.entry(part.clone()).or_default();
382            }
383        }
384    }
385
386    // Render tree
387    #[allow(clippy::too_many_arguments)]
388    fn render_tree(
389        node: &TreeNode,
390        prefix: &str,
391        _is_last: bool,
392        current_path: &str,
393        file_lookup: &HashMap<String, &FileInfo>,
394        options: &ContextOptions,
395    ) -> String {
396        // Pre-allocate with estimated size
397        let estimated_size = (node.dirs.len() + node.files.len()) * 100;
398        let mut output = String::with_capacity(estimated_size);
399
400        // Render directories
401        let dir_count = node.dirs.len();
402        for (i, (name, child)) in node.dirs.iter().enumerate() {
403            let is_last_dir = i == dir_count - 1 && node.files.is_empty();
404            let connector = if is_last_dir {
405                "└── "
406            } else {
407                "├── "
408            };
409            let extension = if is_last_dir { "    " } else { "│   " };
410
411            output.push_str(&format!("{prefix}{connector}{name}/\n"));
412            let child_path = if current_path.is_empty() {
413                name.clone()
414            } else {
415                format!("{current_path}/{name}")
416            };
417            output.push_str(&render_tree(
418                child,
419                &format!("{prefix}{extension}"),
420                is_last_dir,
421                &child_path,
422                file_lookup,
423                options,
424            ));
425        }
426
427        // Render files
428        let file_count = node.files.len();
429        for (i, name) in node.files.iter().enumerate() {
430            let is_last_file = i == file_count - 1;
431            let connector = if is_last_file {
432                "└── "
433            } else {
434                "├── "
435            };
436
437            let file_path = if current_path.is_empty() {
438                name.clone()
439            } else {
440                format!("{current_path}/{name}")
441            };
442
443            // Include metadata if enhanced context is enabled
444            let display_name = if options.enhanced_context {
445                if let Some(file_info) = file_lookup.get(&file_path) {
446                    format!(
447                        "{} ({}, {})",
448                        name,
449                        format_size(file_info.size),
450                        file_type_display(&file_info.file_type)
451                    )
452                } else {
453                    name.clone()
454                }
455            } else {
456                name.clone()
457            };
458
459            output.push_str(&format!("{prefix}{connector}{display_name}\n"));
460        }
461
462        output
463    }
464
465    // Pre-allocate output string
466    let mut output = String::with_capacity(files.len() * 100 + 10);
467    output.push_str(".\n");
468    output.push_str(&render_tree(&root, "", true, "", &file_lookup, options));
469    output
470}
471
472/// Group files by their type
473fn group_files_by_type(files: Vec<FileInfo>) -> Vec<(FileType, Vec<FileInfo>)> {
474    let mut groups: HashMap<FileType, Vec<FileInfo>> = HashMap::new();
475
476    for file in files {
477        groups.entry(file.file_type.clone()).or_default().push(file);
478    }
479
480    let mut result: Vec<_> = groups.into_iter().collect();
481    result.sort_by_key(|(file_type, _)| file_type_priority(file_type));
482    result
483}
484
485/// Get display name for file type
486fn file_type_display(file_type: &FileType) -> &'static str {
487    match file_type {
488        FileType::Rust => "Rust",
489        FileType::Python => "Python",
490        FileType::JavaScript => "JavaScript",
491        FileType::TypeScript => "TypeScript",
492        FileType::Go => "Go",
493        FileType::Java => "Java",
494        FileType::Cpp => "C++",
495        FileType::C => "C",
496        FileType::CSharp => "C#",
497        FileType::Ruby => "Ruby",
498        FileType::Php => "PHP",
499        FileType::Swift => "Swift",
500        FileType::Kotlin => "Kotlin",
501        FileType::Scala => "Scala",
502        FileType::Haskell => "Haskell",
503        FileType::Dart => "Dart",
504        FileType::Lua => "Lua",
505        FileType::R => "R",
506        FileType::Julia => "Julia",
507        FileType::Elixir => "Elixir",
508        FileType::Elm => "Elm",
509        FileType::Markdown => "Markdown",
510        FileType::Json => "JSON",
511        FileType::Yaml => "YAML",
512        FileType::Toml => "TOML",
513        FileType::Xml => "XML",
514        FileType::Html => "HTML",
515        FileType::Css => "CSS",
516        FileType::Text => "Text",
517        FileType::Other => "Other",
518    }
519}
520
521/// Get language hint for syntax highlighting
522fn get_language_hint(file_type: &FileType) -> &'static str {
523    match file_type {
524        FileType::Rust => "rust",
525        FileType::Python => "python",
526        FileType::JavaScript => "javascript",
527        FileType::TypeScript => "typescript",
528        FileType::Go => "go",
529        FileType::Java => "java",
530        FileType::Cpp => "cpp",
531        FileType::C => "c",
532        FileType::CSharp => "csharp",
533        FileType::Ruby => "ruby",
534        FileType::Php => "php",
535        FileType::Swift => "swift",
536        FileType::Kotlin => "kotlin",
537        FileType::Scala => "scala",
538        FileType::Haskell => "haskell",
539        FileType::Dart => "dart",
540        FileType::Lua => "lua",
541        FileType::R => "r",
542        FileType::Julia => "julia",
543        FileType::Elixir => "elixir",
544        FileType::Elm => "elm",
545        FileType::Markdown => "markdown",
546        FileType::Json => "json",
547        FileType::Yaml => "yaml",
548        FileType::Toml => "toml",
549        FileType::Xml => "xml",
550        FileType::Html => "html",
551        FileType::Css => "css",
552        FileType::Text => "text",
553        FileType::Other => "",
554    }
555}
556
557/// Get priority for file type ordering
558fn file_type_priority(file_type: &FileType) -> u8 {
559    match file_type {
560        FileType::Rust => 1,
561        FileType::Python => 2,
562        FileType::JavaScript => 3,
563        FileType::TypeScript => 3,
564        FileType::Go => 4,
565        FileType::Java => 5,
566        FileType::Cpp => 6,
567        FileType::C => 7,
568        FileType::CSharp => 8,
569        FileType::Ruby => 9,
570        FileType::Php => 10,
571        FileType::Swift => 11,
572        FileType::Kotlin => 12,
573        FileType::Scala => 13,
574        FileType::Haskell => 14,
575        FileType::Dart => 15,
576        FileType::Lua => 16,
577        FileType::R => 17,
578        FileType::Julia => 18,
579        FileType::Elixir => 19,
580        FileType::Elm => 20,
581        FileType::Markdown => 21,
582        FileType::Json => 22,
583        FileType::Yaml => 23,
584        FileType::Toml => 24,
585        FileType::Xml => 25,
586        FileType::Html => 26,
587        FileType::Css => 27,
588        FileType::Text => 28,
589        FileType::Other => 29,
590    }
591}
592
593/// Convert path to anchor-friendly string
594fn path_to_anchor(path: &Path) -> String {
595    path.display()
596        .to_string()
597        .replace(['/', '\\', '.', ' '], "-")
598        .to_lowercase()
599}
600
601/// Format file size in human-readable format
602fn format_size(size: u64) -> String {
603    const UNITS: &[&str] = &["B", "KB", "MB", "GB"];
604    let mut size = size as f64;
605    let mut unit_index = 0;
606
607    while size >= 1024.0 && unit_index < UNITS.len() - 1 {
608        size /= 1024.0;
609        unit_index += 1;
610    }
611
612    if unit_index == 0 {
613        format!("{} {}", size as u64, UNITS[unit_index])
614    } else {
615        format!("{:.2} {}", size, UNITS[unit_index])
616    }
617}
618
619#[cfg(test)]
620mod tests {
621    use super::*;
622    use std::path::PathBuf;
623
624    fn create_test_cache() -> Arc<FileCache> {
625        Arc::new(FileCache::new())
626    }
627
628    #[test]
629    fn test_format_size() {
630        assert_eq!(format_size(512), "512 B");
631        assert_eq!(format_size(1024), "1.00 KB");
632        assert_eq!(format_size(1536), "1.50 KB");
633        assert_eq!(format_size(1048576), "1.00 MB");
634    }
635
636    #[test]
637    fn test_path_to_anchor() {
638        assert_eq!(path_to_anchor(Path::new("src/main.rs")), "src-main-rs");
639        assert_eq!(path_to_anchor(Path::new("test file.txt")), "test-file-txt");
640    }
641
642    #[test]
643    fn test_file_type_display() {
644        assert_eq!(file_type_display(&FileType::Rust), "Rust");
645        assert_eq!(file_type_display(&FileType::Python), "Python");
646    }
647
648    #[test]
649    fn test_generate_statistics() {
650        let files = vec![
651            FileInfo {
652                path: PathBuf::from("test1.rs"),
653                relative_path: PathBuf::from("test1.rs"),
654                size: 100,
655                file_type: FileType::Rust,
656                priority: 1.0,
657                imports: Vec::new(),
658                imported_by: Vec::new(),
659                function_calls: Vec::new(),
660                type_references: Vec::new(),
661            },
662            FileInfo {
663                path: PathBuf::from("test2.py"),
664                relative_path: PathBuf::from("test2.py"),
665                size: 200,
666                file_type: FileType::Python,
667                priority: 0.9,
668                imports: Vec::new(),
669                imported_by: Vec::new(),
670                function_calls: Vec::new(),
671                type_references: Vec::new(),
672            },
673        ];
674
675        let stats = generate_statistics(&files);
676        assert!(stats.contains("Total files: 2"));
677        assert!(stats.contains("Total size: 300 B"));
678        assert!(stats.contains("Rust: 1"));
679        assert!(stats.contains("Python: 1"));
680    }
681
682    #[test]
683    fn test_generate_statistics_empty() {
684        let files = vec![];
685        let stats = generate_statistics(&files);
686        assert!(stats.contains("Total files: 0"));
687        assert!(stats.contains("Total size: 0 B"));
688    }
689
690    #[test]
691    fn test_generate_statistics_large_files() {
692        let files = vec![
693            FileInfo {
694                path: PathBuf::from("large.rs"),
695                relative_path: PathBuf::from("large.rs"),
696                size: 2_000_000, // 2MB
697                file_type: FileType::Rust,
698                priority: 1.0,
699                imports: Vec::new(),
700                imported_by: Vec::new(),
701                function_calls: Vec::new(),
702                type_references: Vec::new(),
703            },
704            FileInfo {
705                path: PathBuf::from("huge.py"),
706                relative_path: PathBuf::from("huge.py"),
707                size: 50_000_000, // 50MB
708                file_type: FileType::Python,
709                priority: 0.9,
710                imports: Vec::new(),
711                imported_by: Vec::new(),
712                function_calls: Vec::new(),
713                type_references: Vec::new(),
714            },
715        ];
716
717        let stats = generate_statistics(&files);
718        assert!(stats.contains("Total files: 2"));
719        assert!(stats.contains("MB bytes")); // Just check that it's in MB
720        assert!(stats.contains("Python: 1"));
721        assert!(stats.contains("Rust: 1"));
722    }
723
724    #[test]
725    fn test_generate_file_tree_with_grouping() {
726        let files = vec![
727            FileInfo {
728                path: PathBuf::from("src/main.rs"),
729                relative_path: PathBuf::from("src/main.rs"),
730                size: 1000,
731                file_type: FileType::Rust,
732                priority: 1.5,
733                imports: Vec::new(),
734                imported_by: Vec::new(),
735                function_calls: Vec::new(),
736                type_references: Vec::new(),
737            },
738            FileInfo {
739                path: PathBuf::from("src/lib.rs"),
740                relative_path: PathBuf::from("src/lib.rs"),
741                size: 2000,
742                file_type: FileType::Rust,
743                priority: 1.2,
744                imports: Vec::new(),
745                imported_by: Vec::new(),
746                function_calls: Vec::new(),
747                type_references: Vec::new(),
748            },
749            FileInfo {
750                path: PathBuf::from("tests/test.rs"),
751                relative_path: PathBuf::from("tests/test.rs"),
752                size: 500,
753                file_type: FileType::Rust,
754                priority: 0.8,
755                imports: Vec::new(),
756                imported_by: Vec::new(),
757                function_calls: Vec::new(),
758                type_references: Vec::new(),
759            },
760        ];
761
762        let options = ContextOptions::default();
763        let tree = generate_file_tree(&files, &options);
764        assert!(tree.contains("src/"));
765        assert!(tree.contains("tests/"));
766        assert!(tree.contains("main.rs"));
767        assert!(tree.contains("lib.rs"));
768        assert!(tree.contains("test.rs"));
769    }
770
771    #[test]
772    fn test_context_options_from_config() {
773        use crate::cli::Config;
774        use tempfile::TempDir;
775
776        let temp_dir = TempDir::new().unwrap();
777        let config = Config {
778            paths: Some(vec![temp_dir.path().to_path_buf()]),
779            max_tokens: Some(100000),
780            ..Config::default()
781        };
782
783        let options = ContextOptions::from_config(&config).unwrap();
784        assert_eq!(options.max_tokens, Some(100000));
785        assert!(options.include_tree);
786        assert!(options.include_stats);
787        assert!(!options.group_by_type); // Default is false according to implementation
788    }
789
790    #[test]
791    fn test_generate_markdown_structure_headers() {
792        let files = vec![];
793
794        let options = ContextOptions {
795            max_tokens: None,
796            include_tree: true,
797            include_stats: true,
798            group_by_type: true,
799            sort_by_priority: true,
800            file_header_template: "## {path}".to_string(),
801            doc_header_template: "# Code Context".to_string(),
802            include_toc: true,
803            enhanced_context: false,
804        };
805
806        let cache = create_test_cache();
807        let markdown = generate_markdown(files, options, cache).unwrap();
808
809        // Check that main structure is present even with no files
810        assert!(markdown.contains("# Code Context"));
811        assert!(markdown.contains("## Statistics"));
812    }
813
814    #[test]
815    fn test_enhanced_tree_generation_with_metadata() {
816        use crate::core::walker::FileInfo;
817        use crate::utils::file_ext::FileType;
818        use std::path::PathBuf;
819
820        let files = vec![
821            FileInfo {
822                path: PathBuf::from("src/main.rs"),
823                relative_path: PathBuf::from("src/main.rs"),
824                size: 145,
825                file_type: FileType::Rust,
826                priority: 1.5,
827                imports: Vec::new(),
828                imported_by: Vec::new(),
829                function_calls: Vec::new(),
830                type_references: Vec::new(),
831            },
832            FileInfo {
833                path: PathBuf::from("src/lib.rs"),
834                relative_path: PathBuf::from("src/lib.rs"),
835                size: 89,
836                file_type: FileType::Rust,
837                priority: 1.2,
838                imports: Vec::new(),
839                imported_by: Vec::new(),
840                function_calls: Vec::new(),
841                type_references: Vec::new(),
842            },
843        ];
844
845        let options = ContextOptions {
846            max_tokens: None,
847            include_tree: true,
848            include_stats: true,
849            group_by_type: false,
850            sort_by_priority: true,
851            file_header_template: "## {path}".to_string(),
852            doc_header_template: "# Code Context".to_string(),
853            include_toc: true,
854            enhanced_context: true,
855        };
856
857        let cache = create_test_cache();
858        let markdown = generate_markdown(files, options, cache).unwrap();
859
860        // Should include file sizes and types in tree
861        assert!(markdown.contains("main.rs (145 B, Rust)"));
862        assert!(markdown.contains("lib.rs (89 B, Rust)"));
863    }
864
865    #[test]
866    fn test_enhanced_file_headers_with_metadata() {
867        use crate::core::walker::FileInfo;
868        use crate::utils::file_ext::FileType;
869        use std::path::PathBuf;
870
871        let files = vec![FileInfo {
872            path: PathBuf::from("src/main.rs"),
873            relative_path: PathBuf::from("src/main.rs"),
874            size: 145,
875            file_type: FileType::Rust,
876            priority: 1.5,
877            imports: Vec::new(),
878            imported_by: Vec::new(),
879            function_calls: Vec::new(),
880            type_references: Vec::new(),
881        }];
882
883        let options = ContextOptions {
884            max_tokens: None,
885            include_tree: true,
886            include_stats: true,
887            group_by_type: false,
888            sort_by_priority: true,
889            file_header_template: "## {path}".to_string(),
890            doc_header_template: "# Code Context".to_string(),
891            include_toc: true,
892            enhanced_context: true,
893        };
894
895        let cache = create_test_cache();
896        let markdown = generate_markdown(files, options, cache).unwrap();
897
898        // Should include metadata in file headers
899        assert!(markdown.contains("## src/main.rs (145 B, Rust)"));
900    }
901
902    #[test]
903    fn test_basic_mode_unchanged() {
904        use crate::core::walker::FileInfo;
905        use crate::utils::file_ext::FileType;
906        use std::path::PathBuf;
907
908        let files = vec![FileInfo {
909            path: PathBuf::from("src/main.rs"),
910            relative_path: PathBuf::from("src/main.rs"),
911            size: 145,
912            file_type: FileType::Rust,
913            priority: 1.5,
914            imports: Vec::new(),
915            imported_by: Vec::new(),
916            function_calls: Vec::new(),
917            type_references: Vec::new(),
918        }];
919
920        let options = ContextOptions {
921            max_tokens: None,
922            include_tree: true,
923            include_stats: true,
924            group_by_type: false,
925            sort_by_priority: true,
926            file_header_template: "## {path}".to_string(),
927            doc_header_template: "# Code Context".to_string(),
928            include_toc: true,
929            enhanced_context: false,
930        };
931
932        let cache = create_test_cache();
933        let markdown = generate_markdown(files, options, cache).unwrap();
934
935        // Should NOT include metadata - backward compatibility
936        assert!(markdown.contains("## src/main.rs"));
937        assert!(!markdown.contains("## src/main.rs (145 B, Rust)"));
938        assert!(markdown.contains("main.rs") && !markdown.contains("main.rs (145 B, Rust)"));
939    }
940}
context_creator/core/context_builder.rs

context_creator/core/
context_builder.rs
