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 walkdir::WalkDir;
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 detect_languages(project: &Path) -> anyhow::Result<Vec<String>> {
342 let mut languages = HashSet::new();
343 let ignore_patterns = load_tldrignore(project);
344
345 for entry in WalkDir::new(project)
347 .follow_links(false)
348 .into_iter()
349 .filter_entry(|e| {
350 if e.file_type().is_dir() && e.depth() > 0 {
351 let name = e.file_name().to_string_lossy();
352 !should_skip_component(&name, &ignore_patterns)
353 } else {
354 true
355 }
356 })
357 .filter_map(|e| e.ok())
358 {
359 if entry.file_type().is_file() {
360 if let Some(ext) = entry.path().extension() {
361 let ext_str = ext.to_string_lossy().to_lowercase();
362 match ext_str.as_str() {
363 "py" => {
364 languages.insert("python".to_string());
365 }
366 "ts" | "tsx" => {
367 languages.insert("typescript".to_string());
368 }
369 "js" | "jsx" => {
370 languages.insert("javascript".to_string());
371 }
372 "rs" => {
373 languages.insert("rust".to_string());
374 }
375 "go" => {
376 languages.insert("go".to_string());
377 }
378 "java" => {
379 languages.insert("java".to_string());
380 }
381 "rb" => {
382 languages.insert("ruby".to_string());
383 }
384 "cpp" | "cc" | "cxx" | "hpp" | "h" => {
385 languages.insert("cpp".to_string());
386 }
387 "c" => {
388 languages.insert("c".to_string());
389 }
390 _ => {}
391 }
392 }
393 }
394 }
395
396 let mut result: Vec<_> = languages.into_iter().collect();
397 result.sort();
398
399 if result.is_empty() {
400 result.push("unknown".to_string());
401 }
402
403 Ok(result)
404}
405
406fn build_call_graph(
410 project: &Path,
411 languages: &[String],
412) -> anyhow::Result<(usize, Vec<CallEdge>)> {
413 let mut file_count = 0;
414 let mut edges = Vec::new();
415
416 let extensions: HashSet<&str> = languages
418 .iter()
419 .flat_map(|lang| match lang.as_str() {
420 "python" => vec!["py"],
421 "typescript" => vec!["ts", "tsx"],
422 "javascript" => vec!["js", "jsx"],
423 "rust" => vec!["rs"],
424 "go" => vec!["go"],
425 "java" => vec!["java"],
426 "ruby" => vec!["rb"],
427 "cpp" => vec!["cpp", "cc", "cxx", "hpp", "h"],
428 "c" => vec!["c", "h"],
429 _ => vec![],
430 })
431 .collect();
432
433 let ignore_patterns = load_tldrignore(project);
435 for entry in WalkDir::new(project)
436 .follow_links(false)
437 .into_iter()
438 .filter_entry(|e| {
439 if e.file_type().is_dir() && e.depth() > 0 {
440 let name = e.file_name().to_string_lossy();
441 !should_skip_component(&name, &ignore_patterns)
442 } else {
443 true
444 }
445 })
446 .filter_map(|e| e.ok())
447 {
448 if entry.file_type().is_file() {
449 let path = entry.path();
450 if let Some(ext) = path.extension() {
451 let ext_str = ext.to_string_lossy().to_lowercase();
452 if extensions.contains(ext_str.as_str()) {
453 file_count += 1;
454
455 if let Ok(content) = fs::read_to_string(path) {
457 let file_edges = extract_call_edges(path, &content, &ext_str);
458 edges.extend(file_edges);
459 }
460 }
461 }
462 }
463 }
464
465 Ok((file_count, edges))
466}
467
468fn extract_call_edges(file_path: &std::path::Path, content: &str, lang: &str) -> Vec<CallEdge> {
473 let mut edges = Vec::new();
474 let mut current_func: Option<String> = None;
475
476 let func_pattern = match lang {
478 "py" => Regex::new(r"^\s*def\s+(\w+)\s*\(").ok(),
479 "ts" | "tsx" | "js" | "jsx" => {
480 Regex::new(r"(?:function\s+(\w+)|(\w+)\s*(?::\s*\w+)?\s*=\s*(?:async\s+)?(?:function|\([^)]*\)\s*=>))").ok()
481 }
482 "rs" => Regex::new(r"^\s*(?:pub\s+)?(?:async\s+)?fn\s+(\w+)").ok(),
483 "go" => Regex::new(r"^\s*func\s+(?:\([^)]+\)\s+)?(\w+)\s*\(").ok(),
484 "java" => Regex::new(r"^\s*(?:public|private|protected)?\s*(?:static)?\s*\w+\s+(\w+)\s*\(").ok(),
485 "rb" => Regex::new(r"^\s*def\s+(\w+)").ok(),
486 _ => None,
487 };
488
489 let call_pattern = Regex::new(r"\b(\w+)\s*\(").ok();
491
492 for line in content.lines() {
493 if let Some(ref pattern) = func_pattern {
495 if let Some(caps) = pattern.captures(line) {
496 current_func = caps
498 .iter()
499 .skip(1)
500 .flatten()
501 .next()
502 .map(|m| m.as_str().to_string());
503 }
504 }
505
506 if let (Some(ref current), Some(ref pattern)) = (¤t_func, &call_pattern) {
508 for caps in pattern.captures_iter(line) {
509 if let Some(call) = caps.get(1) {
510 let call_name = call.as_str();
511 if !is_builtin_or_keyword(call_name) && call_name != current {
513 edges.push(CallEdge {
514 from_file: file_path.to_path_buf(),
515 from_func: current.clone(),
516 to_file: file_path.to_path_buf(), to_func: call_name.to_string(),
518 });
519 }
520 }
521 }
522 }
523 }
524
525 edges
526}
527
528fn is_builtin_or_keyword(name: &str) -> bool {
530 let common_builtins = [
531 "if",
532 "else",
533 "for",
534 "while",
535 "return",
536 "print",
537 "len",
538 "str",
539 "int",
540 "float",
541 "bool",
542 "list",
543 "dict",
544 "set",
545 "tuple",
546 "range",
547 "enumerate",
548 "zip",
549 "map",
550 "filter",
551 "sorted",
552 "reversed",
553 "sum",
554 "min",
555 "max",
556 "abs",
557 "round",
558 "type",
559 "isinstance",
560 "issubclass",
561 "hasattr",
562 "getattr",
563 "setattr",
564 "delattr",
565 "open",
566 "close",
567 "read",
568 "write",
569 "append",
570 "extend",
571 "insert",
572 "remove",
573 "pop",
574 "clear",
575 "copy",
576 "update",
577 "get",
578 "keys",
579 "values",
580 "items",
581 "join",
582 "split",
583 "strip",
584 "replace",
585 "format",
586 "console",
587 "log",
588 "require",
589 "import",
590 "export",
591 "const",
592 "let",
593 "var",
594 "new",
595 "this",
596 "self",
597 "super",
598 "class",
599 "struct",
600 "impl",
601 "trait",
602 "pub",
603 "fn",
604 "async",
605 "await",
606 "match",
607 "Some",
608 "None",
609 "Ok",
610 "Err",
611 "Vec",
612 "String",
613 "Box",
614 "Arc",
615 "Rc",
616 "Mutex",
617 "Result",
618 "Option",
619 ];
620
621 common_builtins.contains(&name)
622}
623
624pub async fn cmd_warm(args: WarmArgs) -> DaemonResult<WarmOutput> {
626 let project = args.path.canonicalize().unwrap_or_else(|_| {
628 std::env::current_dir()
629 .unwrap_or_else(|_| PathBuf::from("."))
630 .join(&args.path)
631 });
632
633 let languages = detect_languages(&project).map_err(|e| {
635 DaemonError::Io(std::io::Error::other(e.to_string()))
636 })?;
637
638 let (files, edges) = build_call_graph(&project, &languages).map_err(|e| {
640 DaemonError::Io(std::io::Error::other(e.to_string()))
641 })?;
642
643 let cache_dir = project.join(".tldr/cache");
645 fs::create_dir_all(&cache_dir).map_err(DaemonError::Io)?;
646 let cache_path = cache_dir.join("call_graph.json");
647
648 let cache = CallGraphCache {
649 edges: edges.clone(),
650 languages: languages.clone(),
651 timestamp: chrono::Utc::now().timestamp(),
652 };
653
654 fs::write(&cache_path, serde_json::to_string_pretty(&cache)?)?;
655
656 Ok(WarmOutput {
657 status: "ok".to_string(),
658 files,
659 edges: edges.len(),
660 languages,
661 cache_path: PathBuf::from(".tldr/cache/call_graph.json"),
662 })
663}
664
665#[cfg(test)]
670mod tests {
671 use super::*;
672 use tempfile::TempDir;
673
674 #[test]
675 fn test_warm_args_default() {
676 let args = WarmArgs {
677 path: PathBuf::from("."),
678 background: false,
679 };
680
681 assert_eq!(args.path, PathBuf::from("."));
682 assert!(!args.background);
683 }
684
685 #[test]
686 fn test_warm_args_with_options() {
687 let args = WarmArgs {
688 path: PathBuf::from("/test/project"),
689 background: true,
690 };
691
692 assert!(args.background);
693 }
694
695 #[test]
696 fn test_warm_output_serialization() {
697 let output = WarmOutput {
698 status: "ok".to_string(),
699 files: 150,
700 edges: 2500,
701 languages: vec!["python".to_string(), "typescript".to_string()],
702 cache_path: PathBuf::from(".tldr/cache/call_graph.json"),
703 };
704
705 let json = serde_json::to_string(&output).unwrap();
706 assert!(json.contains("ok"));
707 assert!(json.contains("150"));
708 assert!(json.contains("2500"));
709 assert!(json.contains("python"));
710 }
711
712 #[test]
713 fn test_detect_languages() {
714 let temp = TempDir::new().unwrap();
715 fs::write(temp.path().join("main.py"), "def main(): pass").unwrap();
716 fs::write(temp.path().join("app.ts"), "function main() {}").unwrap();
717
718 let languages = detect_languages(temp.path()).unwrap();
719 assert!(languages.contains(&"python".to_string()));
720 assert!(languages.contains(&"typescript".to_string()));
721 }
722
723 #[test]
724 fn test_detect_languages_empty() {
725 let temp = TempDir::new().unwrap();
726 let languages = detect_languages(temp.path()).unwrap();
727 assert_eq!(languages, vec!["unknown".to_string()]);
728 }
729
730 #[test]
731 fn test_build_call_graph_python() {
732 let temp = TempDir::new().unwrap();
733 fs::write(
734 temp.path().join("main.py"),
735 "def main():\n helper()\n\ndef helper():\n pass",
736 )
737 .unwrap();
738
739 let (files, edges) =
740 build_call_graph(temp.path(), &["python".to_string()]).unwrap();
741
742 assert_eq!(files, 1);
743 assert!(!edges.is_empty());
744 assert!(edges
746 .iter()
747 .any(|e| e.from_func == "main" && e.to_func == "helper"));
748 }
749
750 #[test]
751 fn test_extract_call_edges_python() {
752 let content = "def foo():\n bar()\n baz(1, 2)\n";
753 let edges = extract_call_edges(std::path::Path::new("test.py"), content, "py");
754
755 assert!(edges
756 .iter()
757 .any(|e| e.from_func == "foo" && e.to_func == "bar"));
758 assert!(edges
759 .iter()
760 .any(|e| e.from_func == "foo" && e.to_func == "baz"));
761 }
762
763 #[test]
764 fn test_is_builtin_or_keyword() {
765 assert!(is_builtin_or_keyword("print"));
766 assert!(is_builtin_or_keyword("len"));
767 assert!(is_builtin_or_keyword("if"));
768 assert!(!is_builtin_or_keyword("my_function"));
769 }
770
771 #[test]
772 fn test_call_graph_cache_serialization() {
773 let cache = CallGraphCache {
774 edges: vec![CallEdge {
775 from_file: PathBuf::from("main.py"),
776 from_func: "main".to_string(),
777 to_file: PathBuf::from("utils.py"),
778 to_func: "helper".to_string(),
779 }],
780 languages: vec!["python".to_string()],
781 timestamp: 1234567890,
782 };
783
784 let json = serde_json::to_string(&cache).unwrap();
785 assert!(json.contains("main.py"));
786 assert!(json.contains("helper"));
787 assert!(json.contains("1234567890"));
788 }
789
790 mod proptest_warm {
795 use super::*;
796 use proptest::prelude::*;
797
798 fn arb_component() -> impl Strategy<Value = String> {
800 prop::string::string_regex("[a-zA-Z0-9_.][a-zA-Z0-9_.-]{0,15}")
801 .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}