mermaid_cli/models/
lazy_context.rs

1use anyhow::Result;
2use rustc_hash::FxHashMap;
3use std::path::{Path, PathBuf};
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::{Arc, Mutex};
6use tokio::sync::RwLock;
7
8/// A lazily-loaded project context that loads files on demand
9#[derive(Debug, Clone)]
10pub struct LazyProjectContext {
11    /// Root path of the project
12    pub root_path: String,
13    /// List of all file paths (loaded immediately)
14    pub file_paths: Arc<Vec<PathBuf>>,
15    /// Lazily loaded file contents (using FxHashMap for performance)
16    pub files: Arc<RwLock<FxHashMap<String, String>>>,
17    /// Running token count (updated as files are loaded)
18    pub token_count: Arc<AtomicUsize>,
19    /// Files that have been requested for loading
20    pub loading_queue: Arc<Mutex<Vec<PathBuf>>>,
21    /// Cache manager for persistent caching
22    pub cache: Option<Arc<crate::cache::CacheManager>>,
23}
24
25impl LazyProjectContext {
26    /// Create a new lazy project context with just file paths
27    pub fn new(root_path: String, file_paths: Vec<PathBuf>) -> Self {
28        let cache = crate::cache::CacheManager::new().ok().map(Arc::new);
29
30        Self {
31            root_path,
32            file_paths: Arc::new(file_paths),
33            files: Arc::new(RwLock::new(FxHashMap::default())),
34            token_count: Arc::new(AtomicUsize::new(0)),
35            loading_queue: Arc::new(Mutex::new(Vec::new())),
36            cache,
37        }
38    }
39
40    /// Get a file's content, loading it if necessary
41    pub async fn get_file(&self, path: &str) -> Result<Option<String>> {
42        // Check if already loaded
43        {
44            let files = self.files.read().await;
45            if let Some(content) = files.get(path) {
46                return Ok(Some(content.clone()));
47            }
48        }
49
50        // Not loaded, need to load it
51        let full_path = if path.starts_with(&self.root_path) {
52            PathBuf::from(path)
53        } else {
54            PathBuf::from(&self.root_path).join(path)
55        };
56
57        // Load the file
58        if full_path.exists() {
59            let content = tokio::fs::read_to_string(&full_path).await?;
60
61            // Count tokens
62            if let Some(ref cache) = self.cache {
63                if let Ok(tokens) = cache.get_or_compute_tokens(&full_path, &content, "cl100k_base")
64                {
65                    self.token_count.fetch_add(tokens, Ordering::Relaxed);
66                }
67            }
68
69            // Store in memory
70            let mut files = self.files.write().await;
71            files.insert(path.to_string(), content.clone());
72
73            Ok(Some(content))
74        } else {
75            Ok(None)
76        }
77    }
78
79    /// Load a batch of files in the background
80    pub async fn load_files_batch(&self, paths: Vec<String>) -> Result<()> {
81        use futures::future::join_all;
82
83        let futures = paths.into_iter().map(|path| {
84            let self_clone = self.clone();
85            async move {
86                let _ = self_clone.get_file(&path).await;
87            }
88        });
89
90        join_all(futures).await;
91        Ok(())
92    }
93
94    /// Get the list of all file paths (instant)
95    pub fn get_file_list(&self) -> Vec<String> {
96        self.file_paths
97            .iter()
98            .filter_map(|p| {
99                p.strip_prefix(&self.root_path)
100                    .ok()
101                    .and_then(|p| p.to_str())
102                    .map(|s| s.to_string())
103            })
104            .collect()
105    }
106
107    /// Get current loaded file count
108    pub async fn loaded_file_count(&self) -> usize {
109        self.files.read().await.len()
110    }
111
112    /// Get total file count
113    pub fn total_file_count(&self) -> usize {
114        self.file_paths.len()
115    }
116
117    /// Check if all files are loaded
118    pub async fn is_fully_loaded(&self) -> bool {
119        self.loaded_file_count().await >= self.total_file_count()
120    }
121
122    /// Convert to regular ProjectContext (for compatibility)
123    pub async fn to_project_context(&self) -> crate::models::ProjectContext {
124        let files = self.files.read().await;
125        let mut context = crate::models::ProjectContext::new(self.root_path.clone());
126        context.token_count = self.token_count.load(Ordering::Relaxed);
127
128        for (path, content) in files.iter() {
129            context.add_file(path.clone(), content.clone());
130        }
131
132        context
133    }
134}
135
136/// Priority files to load first for better UX
137pub fn get_priority_files(root_path: &str) -> Vec<String> {
138    vec![
139        "README.md",
140        "readme.md",
141        "README.rst",
142        "README.txt",
143        "CLAUDE.md", // Mermaid's own instructions file
144        "Cargo.toml",
145        "package.json",
146        "pyproject.toml",
147        "go.mod",
148        ".gitignore",
149        "LICENSE",
150    ]
151    .into_iter()
152    .filter_map(|f| {
153        let path = Path::new(root_path).join(f);
154        if path.exists() {
155            Some(f.to_string())
156        } else {
157            None
158        }
159    })
160    .collect()
161}
mermaid_cli/models/lazy_context.rs

mermaid_cli/models/
lazy_context.rs
