1use std::collections::HashSet;
28use std::fs;
29use std::path::{Path, PathBuf};
30use std::process::Command as StdCommand;
31
32use clap::Args;
33use regex::Regex;
34use serde::{Deserialize, Serialize};
35use tldr_core::walker::walk_project;
36
37use crate::output::OutputFormat;
38
39use super::error::{DaemonError, DaemonResult};
40use super::ipc::{check_socket_alive, send_command};
41use super::types::DaemonCommand;
42
43#[derive(Debug, Clone, Args)]
49pub struct WarmArgs {
50 #[arg(default_value = ".")]
52 pub path: PathBuf,
53
54 #[arg(long, short = 'b')]
56 pub background: bool,
57 }
59
60#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct WarmOutput {
67 pub status: String,
69 pub files: usize,
71 pub edges: usize,
73 pub languages: Vec<String>,
75 pub cache_path: PathBuf,
77}
78
79#[derive(Debug, Clone, Serialize, Deserialize)]
81pub struct CallEdge {
82 pub from_file: PathBuf,
83 pub from_func: String,
84 pub to_file: PathBuf,
85 pub to_func: String,
86}
87
88#[derive(Debug, Clone, Serialize, Deserialize)]
90pub struct CallGraphCache {
91 pub edges: Vec<CallEdge>,
92 pub languages: Vec<String>,
93 pub timestamp: i64,
94}
95
96impl WarmArgs {
101 pub fn run(&self, format: OutputFormat, quiet: bool) -> anyhow::Result<()> {
103 let runtime = tokio::runtime::Runtime::new()?;
105 runtime.block_on(self.run_async(format, quiet))
106 }
107
108 async fn run_async(&self, format: OutputFormat, quiet: bool) -> anyhow::Result<()> {
110 let project = self.path.canonicalize().unwrap_or_else(|_| {
112 std::env::current_dir()
113 .unwrap_or_else(|_| PathBuf::from("."))
114 .join(&self.path)
115 });
116
117 if self.background {
118 self.run_background(&project, format, quiet).await
120 } else {
121 if check_socket_alive(&project).await {
123 self.run_via_daemon(&project, format, quiet).await
124 } else {
125 self.run_foreground(&project, format, quiet).await
127 }
128 }
129 }
130
131 async fn run_background(
133 &self,
134 project: &Path,
135 format: OutputFormat,
136 quiet: bool,
137 ) -> anyhow::Result<()> {
138 let exe = std::env::current_exe()?;
140 let mut cmd = StdCommand::new(exe);
141 cmd.arg("warm").arg(project.to_str().unwrap_or("."));
142
143 #[cfg(unix)]
147 {
148 use std::os::unix::process::CommandExt;
149 cmd.process_group(0);
150 }
151
152 #[cfg(windows)]
154 {
155 use std::os::windows::process::CommandExt;
156 const CREATE_NO_WINDOW: u32 = 0x08000000;
157 const DETACHED_PROCESS: u32 = 0x00000008;
158 cmd.creation_flags(CREATE_NO_WINDOW | DETACHED_PROCESS);
159 }
160
161 cmd.spawn()?;
162
163 if !quiet {
165 match format {
166 OutputFormat::Json | OutputFormat::Compact => {
167 let output = serde_json::json!({
168 "status": "ok",
169 "message": "Warming cache in background..."
170 });
171 println!("{}", serde_json::to_string_pretty(&output)?);
172 }
173 OutputFormat::Text | OutputFormat::Sarif | OutputFormat::Dot => {
174 println!("Warming cache in background...");
175 }
176 }
177 }
178
179 Ok(())
180 }
181
182 async fn run_via_daemon(
184 &self,
185 project: &Path,
186 format: OutputFormat,
187 quiet: bool,
188 ) -> anyhow::Result<()> {
189 let cmd = DaemonCommand::Warm {
190 language: None, };
192
193 let response = send_command(project, &cmd)
194 .await
195 .map_err(|e| anyhow::anyhow!("Failed to send warm command to daemon: {}", e))?;
196
197 if !quiet {
198 match format {
199 OutputFormat::Json | OutputFormat::Compact => {
200 println!("{}", serde_json::to_string_pretty(&response)?);
201 }
202 OutputFormat::Text | OutputFormat::Sarif | OutputFormat::Dot => {
203 println!("Warm command sent to daemon");
204 }
205 }
206 }
207
208 Ok(())
209 }
210
211 async fn run_foreground(
213 &self,
214 project: &Path,
215 format: OutputFormat,
216 quiet: bool,
217 ) -> anyhow::Result<()> {
218 if !quiet {
219 match format {
220 OutputFormat::Text | OutputFormat::Sarif | OutputFormat::Dot => {
221 println!("Warming call graph cache...");
222 }
223 _ => {}
224 }
225 }
226
227 let tldr_dir = project.join(".tldr");
229 fs::create_dir_all(&tldr_dir)?;
230
231 let ignore_path = project.join(".tldrignore");
233 if !ignore_path.exists() {
234 fs::write(
235 &ignore_path,
236 "# TLDR ignore file\n\
237 .git/\n\
238 node_modules/\n\
239 __pycache__/\n\
240 target/\n\
241 build/\n\
242 dist/\n\
243 .venv/\n\
244 venv/\n\
245 *.pyc\n\
246 *.pyo\n",
247 )?;
248 }
249
250 let languages = detect_languages(project)?;
252
253 let (files, edges) = build_call_graph(project, &languages)?;
255
256 let cache_dir = tldr_dir.join("cache");
258 fs::create_dir_all(&cache_dir)?;
259 let cache_path = cache_dir.join("call_graph.json");
260
261 let cache = CallGraphCache {
262 edges: edges.clone(),
263 languages: languages.clone(),
264 timestamp: chrono::Utc::now().timestamp(),
265 };
266
267 fs::write(&cache_path, serde_json::to_string_pretty(&cache)?)?;
268
269 let output = WarmOutput {
271 status: "ok".to_string(),
272 files,
273 edges: edges.len(),
274 languages,
275 cache_path: PathBuf::from(".tldr/cache/call_graph.json"),
276 };
277
278 match format {
280 OutputFormat::Json | OutputFormat::Compact => {
281 println!("{}", serde_json::to_string_pretty(&output)?);
282 }
283 OutputFormat::Text | OutputFormat::Sarif | OutputFormat::Dot => {
284 println!(
285 "Indexed {} files, found {} edges",
286 output.files, output.edges
287 );
288 println!("Languages: {}", output.languages.join(", "));
289 println!("Cache written to: {}", output.cache_path.display());
290 }
291 }
292
293 Ok(())
294 }
295}
296
297const SKIP_DIRS: &[&str] = &[
303 "node_modules",
304 "__pycache__",
305 "target",
306 "build",
307 "dist",
308 "venv",
309 ".venv",
310];
311
312fn load_tldrignore(project: &Path) -> HashSet<String> {
315 let mut patterns = HashSet::new();
316 let ignore_path = project.join(".tldrignore");
317 if let Ok(content) = fs::read_to_string(&ignore_path) {
318 for line in content.lines() {
319 let trimmed = line.trim();
320 if trimmed.is_empty() || trimmed.starts_with('#') {
321 continue;
322 }
323 let name = trimmed.trim_end_matches('/');
325 if !name.is_empty() {
326 patterns.insert(name.to_string());
327 }
328 }
329 }
330 patterns
331}
332
333fn should_skip_component(component: &str, ignore_patterns: &HashSet<String>) -> bool {
335 component.starts_with('.')
336 || SKIP_DIRS.contains(&component)
337 || ignore_patterns.contains(component)
338}
339
340fn path_has_ignored_component(
347 path: &Path,
348 project: &Path,
349 ignore_patterns: &HashSet<String>,
350) -> bool {
351 let rel = path.strip_prefix(project).unwrap_or(path);
352 rel.components().any(|c| {
353 c.as_os_str()
354 .to_str()
355 .map(|s| should_skip_component(s, ignore_patterns))
356 .unwrap_or(false)
357 })
358}
359
360fn detect_languages(project: &Path) -> anyhow::Result<Vec<String>> {
362 let mut languages = HashSet::new();
363 let ignore_patterns = load_tldrignore(project);
364
365 for entry in walk_project(project)
369 .filter(|e| !path_has_ignored_component(e.path(), project, &ignore_patterns))
370 {
371 if entry.file_type().map(|ft| ft.is_file()).unwrap_or(false) {
372 if let Some(ext) = entry.path().extension() {
373 let ext_str = ext.to_string_lossy().to_lowercase();
374 match ext_str.as_str() {
375 "py" => {
376 languages.insert("python".to_string());
377 }
378 "ts" | "tsx" => {
379 languages.insert("typescript".to_string());
380 }
381 "js" | "jsx" => {
382 languages.insert("javascript".to_string());
383 }
384 "rs" => {
385 languages.insert("rust".to_string());
386 }
387 "go" => {
388 languages.insert("go".to_string());
389 }
390 "java" => {
391 languages.insert("java".to_string());
392 }
393 "rb" => {
394 languages.insert("ruby".to_string());
395 }
396 "cpp" | "cc" | "cxx" | "hpp" | "h" => {
397 languages.insert("cpp".to_string());
398 }
399 "c" => {
400 languages.insert("c".to_string());
401 }
402 _ => {}
403 }
404 }
405 }
406 }
407
408 let mut result: Vec<_> = languages.into_iter().collect();
409 result.sort();
410
411 if result.is_empty() {
412 result.push("unknown".to_string());
413 }
414
415 Ok(result)
416}
417
418fn build_call_graph(
422 project: &Path,
423 languages: &[String],
424) -> anyhow::Result<(usize, Vec<CallEdge>)> {
425 let mut file_count = 0;
426 let mut edges = Vec::new();
427
428 let extensions: HashSet<&str> = languages
430 .iter()
431 .flat_map(|lang| match lang.as_str() {
432 "python" => vec!["py"],
433 "typescript" => vec!["ts", "tsx"],
434 "javascript" => vec!["js", "jsx"],
435 "rust" => vec!["rs"],
436 "go" => vec!["go"],
437 "java" => vec!["java"],
438 "ruby" => vec!["rb"],
439 "cpp" => vec!["cpp", "cc", "cxx", "hpp", "h"],
440 "c" => vec!["c", "h"],
441 _ => vec![],
442 })
443 .collect();
444
445 let ignore_patterns = load_tldrignore(project);
449 for entry in walk_project(project)
450 .filter(|e| !path_has_ignored_component(e.path(), project, &ignore_patterns))
451 {
452 if entry.file_type().map(|ft| ft.is_file()).unwrap_or(false) {
453 let path = entry.path();
454 if let Some(ext) = path.extension() {
455 let ext_str = ext.to_string_lossy().to_lowercase();
456 if extensions.contains(ext_str.as_str()) {
457 file_count += 1;
458
459 if let Ok(content) = fs::read_to_string(path) {
461 let file_edges = extract_call_edges(path, &content, &ext_str);
462 edges.extend(file_edges);
463 }
464 }
465 }
466 }
467 }
468
469 Ok((file_count, edges))
470}
471
472fn extract_call_edges(file_path: &std::path::Path, content: &str, lang: &str) -> Vec<CallEdge> {
477 let mut edges = Vec::new();
478 let mut current_func: Option<String> = None;
479
480 let func_pattern = match lang {
482 "py" => Regex::new(r"^\s*def\s+(\w+)\s*\(").ok(),
483 "ts" | "tsx" | "js" | "jsx" => {
484 Regex::new(r"(?:function\s+(\w+)|(\w+)\s*(?::\s*\w+)?\s*=\s*(?:async\s+)?(?:function|\([^)]*\)\s*=>))").ok()
485 }
486 "rs" => Regex::new(r"^\s*(?:pub\s+)?(?:async\s+)?fn\s+(\w+)").ok(),
487 "go" => Regex::new(r"^\s*func\s+(?:\([^)]+\)\s+)?(\w+)\s*\(").ok(),
488 "java" => Regex::new(r"^\s*(?:public|private|protected)?\s*(?:static)?\s*\w+\s+(\w+)\s*\(").ok(),
489 "rb" => Regex::new(r"^\s*def\s+(\w+)").ok(),
490 _ => None,
491 };
492
493 let call_pattern = Regex::new(r"\b(\w+)\s*\(").ok();
495
496 for line in content.lines() {
497 if let Some(ref pattern) = func_pattern {
499 if let Some(caps) = pattern.captures(line) {
500 current_func = caps
502 .iter()
503 .skip(1)
504 .flatten()
505 .next()
506 .map(|m| m.as_str().to_string());
507 }
508 }
509
510 if let (Some(ref current), Some(ref pattern)) = (¤t_func, &call_pattern) {
512 for caps in pattern.captures_iter(line) {
513 if let Some(call) = caps.get(1) {
514 let call_name = call.as_str();
515 if !is_builtin_or_keyword(call_name) && call_name != current {
517 edges.push(CallEdge {
518 from_file: file_path.to_path_buf(),
519 from_func: current.clone(),
520 to_file: file_path.to_path_buf(), to_func: call_name.to_string(),
522 });
523 }
524 }
525 }
526 }
527 }
528
529 edges
530}
531
532fn is_builtin_or_keyword(name: &str) -> bool {
534 let common_builtins = [
535 "if",
536 "else",
537 "for",
538 "while",
539 "return",
540 "print",
541 "len",
542 "str",
543 "int",
544 "float",
545 "bool",
546 "list",
547 "dict",
548 "set",
549 "tuple",
550 "range",
551 "enumerate",
552 "zip",
553 "map",
554 "filter",
555 "sorted",
556 "reversed",
557 "sum",
558 "min",
559 "max",
560 "abs",
561 "round",
562 "type",
563 "isinstance",
564 "issubclass",
565 "hasattr",
566 "getattr",
567 "setattr",
568 "delattr",
569 "open",
570 "close",
571 "read",
572 "write",
573 "append",
574 "extend",
575 "insert",
576 "remove",
577 "pop",
578 "clear",
579 "copy",
580 "update",
581 "get",
582 "keys",
583 "values",
584 "items",
585 "join",
586 "split",
587 "strip",
588 "replace",
589 "format",
590 "console",
591 "log",
592 "require",
593 "import",
594 "export",
595 "const",
596 "let",
597 "var",
598 "new",
599 "this",
600 "self",
601 "super",
602 "class",
603 "struct",
604 "impl",
605 "trait",
606 "pub",
607 "fn",
608 "async",
609 "await",
610 "match",
611 "Some",
612 "None",
613 "Ok",
614 "Err",
615 "Vec",
616 "String",
617 "Box",
618 "Arc",
619 "Rc",
620 "Mutex",
621 "Result",
622 "Option",
623 ];
624
625 common_builtins.contains(&name)
626}
627
628pub async fn cmd_warm(args: WarmArgs) -> DaemonResult<WarmOutput> {
630 let project = args.path.canonicalize().unwrap_or_else(|_| {
632 std::env::current_dir()
633 .unwrap_or_else(|_| PathBuf::from("."))
634 .join(&args.path)
635 });
636
637 let languages = detect_languages(&project)
639 .map_err(|e| DaemonError::Io(std::io::Error::other(e.to_string())))?;
640
641 let (files, edges) = build_call_graph(&project, &languages)
643 .map_err(|e| DaemonError::Io(std::io::Error::other(e.to_string())))?;
644
645 let cache_dir = project.join(".tldr/cache");
647 fs::create_dir_all(&cache_dir).map_err(DaemonError::Io)?;
648 let cache_path = cache_dir.join("call_graph.json");
649
650 let cache = CallGraphCache {
651 edges: edges.clone(),
652 languages: languages.clone(),
653 timestamp: chrono::Utc::now().timestamp(),
654 };
655
656 fs::write(&cache_path, serde_json::to_string_pretty(&cache)?)?;
657
658 Ok(WarmOutput {
659 status: "ok".to_string(),
660 files,
661 edges: edges.len(),
662 languages,
663 cache_path: PathBuf::from(".tldr/cache/call_graph.json"),
664 })
665}
666
667#[cfg(test)]
672mod tests {
673 use super::*;
674 use tempfile::TempDir;
675
676 #[test]
677 fn test_warm_args_default() {
678 let args = WarmArgs {
679 path: PathBuf::from("."),
680 background: false,
681 };
682
683 assert_eq!(args.path, PathBuf::from("."));
684 assert!(!args.background);
685 }
686
687 #[test]
688 fn test_warm_args_with_options() {
689 let args = WarmArgs {
690 path: PathBuf::from("/test/project"),
691 background: true,
692 };
693
694 assert!(args.background);
695 }
696
697 #[test]
698 fn test_warm_output_serialization() {
699 let output = WarmOutput {
700 status: "ok".to_string(),
701 files: 150,
702 edges: 2500,
703 languages: vec!["python".to_string(), "typescript".to_string()],
704 cache_path: PathBuf::from(".tldr/cache/call_graph.json"),
705 };
706
707 let json = serde_json::to_string(&output).unwrap();
708 assert!(json.contains("ok"));
709 assert!(json.contains("150"));
710 assert!(json.contains("2500"));
711 assert!(json.contains("python"));
712 }
713
714 #[test]
715 fn test_detect_languages() {
716 let temp = TempDir::new().unwrap();
717 fs::write(temp.path().join("main.py"), "def main(): pass").unwrap();
718 fs::write(temp.path().join("app.ts"), "function main() {}").unwrap();
719
720 let languages = detect_languages(temp.path()).unwrap();
721 assert!(languages.contains(&"python".to_string()));
722 assert!(languages.contains(&"typescript".to_string()));
723 }
724
725 #[test]
726 fn test_detect_languages_empty() {
727 let temp = TempDir::new().unwrap();
728 let languages = detect_languages(temp.path()).unwrap();
729 assert_eq!(languages, vec!["unknown".to_string()]);
730 }
731
732 #[test]
733 fn test_build_call_graph_python() {
734 let temp = TempDir::new().unwrap();
735 fs::write(
736 temp.path().join("main.py"),
737 "def main():\n helper()\n\ndef helper():\n pass",
738 )
739 .unwrap();
740
741 let (files, edges) = build_call_graph(temp.path(), &["python".to_string()]).unwrap();
742
743 assert_eq!(files, 1);
744 assert!(!edges.is_empty());
745 assert!(edges
747 .iter()
748 .any(|e| e.from_func == "main" && e.to_func == "helper"));
749 }
750
751 #[test]
752 fn test_extract_call_edges_python() {
753 let content = "def foo():\n bar()\n baz(1, 2)\n";
754 let edges = extract_call_edges(std::path::Path::new("test.py"), content, "py");
755
756 assert!(edges
757 .iter()
758 .any(|e| e.from_func == "foo" && e.to_func == "bar"));
759 assert!(edges
760 .iter()
761 .any(|e| e.from_func == "foo" && e.to_func == "baz"));
762 }
763
764 #[test]
765 fn test_is_builtin_or_keyword() {
766 assert!(is_builtin_or_keyword("print"));
767 assert!(is_builtin_or_keyword("len"));
768 assert!(is_builtin_or_keyword("if"));
769 assert!(!is_builtin_or_keyword("my_function"));
770 }
771
772 #[test]
773 fn test_call_graph_cache_serialization() {
774 let cache = CallGraphCache {
775 edges: vec![CallEdge {
776 from_file: PathBuf::from("main.py"),
777 from_func: "main".to_string(),
778 to_file: PathBuf::from("utils.py"),
779 to_func: "helper".to_string(),
780 }],
781 languages: vec!["python".to_string()],
782 timestamp: 1234567890,
783 };
784
785 let json = serde_json::to_string(&cache).unwrap();
786 assert!(json.contains("main.py"));
787 assert!(json.contains("helper"));
788 assert!(json.contains("1234567890"));
789 }
790
791 mod proptest_warm {
796 use super::*;
797 use proptest::prelude::*;
798
799 fn arb_component() -> impl Strategy<Value = String> {
801 prop::string::string_regex("[a-zA-Z0-9_.][a-zA-Z0-9_.-]{0,15}").unwrap()
802 }
803
804 proptest! {
805 #[test]
807 fn skip_component_no_panic(component in ".*") {
808 let patterns = HashSet::new();
809 let _ = should_skip_component(&component, &patterns);
810 }
811
812 #[test]
814 fn hidden_dirs_always_skipped(name in "\\.[a-zA-Z0-9_]{1,20}") {
815 let patterns = HashSet::new();
816 prop_assert!(should_skip_component(&name, &patterns),
817 "'{}' starts with '.' but was not skipped", name);
818 }
819
820 #[test]
822 fn ignore_patterns_always_skipped(
823 name in arb_component(),
824 extra in prop::collection::hash_set(arb_component(), 0..5),
825 ) {
826 let mut patterns = extra;
827 patterns.insert(name.clone());
828 prop_assert!(should_skip_component(&name, &patterns),
829 "'{}' is in ignore set but was not skipped", name);
830 }
831
832 #[test]
835 fn detect_languages_no_panic(
836 files in prop::collection::vec(
837 (arb_component(), prop::sample::select(vec!["py", "ts", "rs", "go", "rb", "txt", ""])),
838 0..10
839 )
840 ) {
841 let temp = TempDir::new().unwrap();
842 for (name, ext) in &files {
843 let filename = if ext.is_empty() {
844 name.clone()
845 } else {
846 format!("{}.{}", name, ext)
847 };
848 let _ = fs::write(temp.path().join(&filename), "content");
849 }
850 let result = detect_languages(temp.path());
851 prop_assert!(result.is_ok(), "detect_languages should not fail");
852 let langs = result.unwrap();
853 prop_assert!(!langs.is_empty(), "should return at least 'unknown'");
854 }
855 }
856 }
857}