lens_core/benchmark/
dataset_loader.rs

1//! General Dataset Loading Utilities
2//!
3//! Provides utilities for loading various dataset formats and managing
4//! corpus indexing for benchmarking.
5
6use super::types::{GoldenQuery, QueryType, LoadingError};
7use anyhow::{anyhow, Result};
8use serde_json;
9use std::path::{Path, PathBuf};
10use std::collections::HashMap;
11use tokio::fs;
12use tracing::{info, warn, debug};
13
14/// General dataset loader for various formats
15pub struct DatasetLoader {
16    base_path: PathBuf,
17}
18
19impl DatasetLoader {
20    /// Create new dataset loader
21    pub fn new<P: AsRef<Path>>(base_path: P) -> Self {
22        Self {
23            base_path: base_path.as_ref().to_path_buf(),
24        }
25    }
26
27    /// Load golden queries from JSON file (legacy format)
28    pub async fn load_golden_queries<P: AsRef<Path>>(&self, file_path: P) -> Result<Vec<GoldenQuery>> {
29        let full_path = self.base_path.join(file_path);
30        let content = fs::read_to_string(&full_path).await
31            .map_err(|e| anyhow!(LoadingError::IoError { source: e }))?;
32
33        // Try different JSON formats
34        
35        // Format 1: Array of GoldenQuery objects
36        if let Ok(queries) = serde_json::from_str::<Vec<GoldenQuery>>(&content) {
37            info!("📄 Loaded {} golden queries from {}", queries.len(), full_path.display());
38            return Ok(queries);
39        }
40
41        // Format 2: Simple array of query objects (legacy)
42        if let Ok(simple_queries) = serde_json::from_str::<Vec<serde_json::Value>>(&content) {
43            info!("📄 Converting {} legacy queries from {}", simple_queries.len(), full_path.display());
44            return self.convert_legacy_queries(simple_queries);
45        }
46
47        Err(anyhow!(LoadingError::InvalidFormat { 
48            reason: format!("Unrecognized JSON format in {}", full_path.display())
49        }))
50    }
51
52    /// Convert legacy query format to GoldenQuery
53    fn convert_legacy_queries(&self, legacy_queries: Vec<serde_json::Value>) -> Result<Vec<GoldenQuery>> {
54        let mut queries = Vec::new();
55
56        for (index, legacy_query) in legacy_queries.into_iter().enumerate() {
57            let query_obj = legacy_query.as_object()
58                .ok_or_else(|| anyhow!(LoadingError::InvalidFormat {
59                    reason: format!("Query {} is not an object", index)
60                }))?;
61
62            let query_str = query_obj.get("query")
63                .and_then(|v| v.as_str())
64                .unwrap_or("")
65                .to_string();
66
67            let expected_files = query_obj.get("expected_files")
68                .and_then(|v| v.as_array())
69                .map(|arr| arr.iter()
70                    .filter_map(|f| f.as_str().map(|s| s.to_string()))
71                    .collect::<Vec<_>>())
72                .unwrap_or_else(Vec::new);
73
74            let query_type = query_obj.get("query_type")
75                .and_then(|v| v.as_str())
76                .and_then(|s| self.parse_query_type(s))
77                .unwrap_or(QueryType::Identifier);
78
79            let language = query_obj.get("language")
80                .and_then(|v| v.as_str())
81                .map(|s| s.to_string());
82
83            queries.push(GoldenQuery {
84                query: query_str,
85                expected_files,
86                query_type,
87                metadata: HashMap::new(),
88                language,
89                confidence: None,
90            });
91        }
92
93        Ok(queries)
94    }
95
96    /// Parse query type from string
97    fn parse_query_type(&self, type_str: &str) -> Option<QueryType> {
98        match type_str.to_lowercase().as_str() {
99            "exact_match" | "exact" => Some(QueryType::ExactMatch),
100            "identifier" | "id" => Some(QueryType::Identifier),
101            "structural" | "struct" => Some(QueryType::Structural),
102            "semantic" | "nlp" => Some(QueryType::Semantic),
103            "reference" | "ref" => Some(QueryType::Reference),
104            "definition" | "def" => Some(QueryType::Definition),
105            _ => None,
106        }
107    }
108
109    /// Discover corpus files for indexing
110    pub async fn discover_corpus_files<P: AsRef<Path>>(&self, corpus_path: P) -> Result<Vec<PathBuf>> {
111        let corpus_dir = self.base_path.join(corpus_path);
112        
113        if !corpus_dir.exists() {
114            return Err(anyhow!(LoadingError::DatasetNotFound {
115                path: corpus_dir.to_string_lossy().to_string()
116            }));
117        }
118
119        let mut files = Vec::new();
120        self.walk_directory(&corpus_dir, &mut files).await?;
121        
122        // Filter for relevant source files
123        let source_files: Vec<PathBuf> = files.into_iter()
124            .filter(|path| self.is_source_file(path))
125            .collect();
126
127        info!("📁 Discovered {} source files in {}", source_files.len(), corpus_dir.display());
128        Ok(source_files)
129    }
130
131    /// Recursively walk directory to find files
132    fn walk_directory<'a>(&'a self, dir: &'a Path, files: &'a mut Vec<PathBuf>) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + 'a>> {
133        Box::pin(async move {
134            let mut entries = fs::read_dir(dir).await?;
135            
136            while let Some(entry) = entries.next_entry().await? {
137                let path = entry.path();
138                
139                if path.is_dir() {
140                    // Skip common build/cache directories
141                    let dir_name = path.file_name()
142                        .and_then(|n| n.to_str())
143                        .unwrap_or("");
144                    
145                    if !self.should_skip_directory(dir_name) {
146                        self.walk_directory(&path, files).await?;
147                    }
148                } else {
149                    files.push(path);
150                }
151            }
152            
153            Ok(())
154        })
155    }
156
157    /// Check if directory should be skipped during corpus discovery
158    fn should_skip_directory(&self, dir_name: &str) -> bool {
159        matches!(dir_name, 
160            "target" | "node_modules" | ".git" | ".svn" | ".hg" | 
161            "build" | "dist" | "out" | ".next" | ".nuxt" |
162            "__pycache__" | ".pytest_cache" | "coverage" |
163            ".DS_Store" | "Thumbs.db"
164        )
165    }
166
167    /// Check if file is a source file worth indexing
168    fn is_source_file(&self, path: &Path) -> bool {
169        let extension = path.extension()
170            .and_then(|ext| ext.to_str())
171            .unwrap_or("")
172            .to_lowercase();
173
174        matches!(extension.as_str(),
175            "rs" | "py" | "js" | "ts" | "jsx" | "tsx" | 
176            "java" | "kt" | "scala" | "go" | "c" | "cpp" | 
177            "cc" | "cxx" | "h" | "hpp" | "cs" | "php" | 
178            "rb" | "swift" | "m" | "mm" | "dart" | "elm" |
179            "clj" | "cljs" | "hs" | "ml" | "fs" | "pl" | 
180            "r" | "jl" | "lua" | "sh" | "bash" | "zsh" |
181            "sql" | "graphql" | "proto" | "thrift"
182        )
183    }
184
185    /// Generate corpus statistics
186    pub async fn generate_corpus_stats<P: AsRef<Path>>(&self, corpus_files: &[PathBuf]) -> Result<CorpusStats> {
187        let mut stats = CorpusStats {
188            total_files: corpus_files.len(),
189            total_lines: 0,
190            total_size_bytes: 0,
191            language_distribution: HashMap::new(),
192        };
193
194        for file_path in corpus_files {
195            if let Ok(metadata) = fs::metadata(file_path).await {
196                stats.total_size_bytes += metadata.len();
197            }
198
199            if let Ok(content) = fs::read_to_string(file_path).await {
200                let line_count = content.lines().count();
201                stats.total_lines += line_count;
202
203                // Count by language (based on extension)
204                if let Some(ext) = file_path.extension().and_then(|e| e.to_str()) {
205                    *stats.language_distribution.entry(ext.to_string()).or_insert(0) += 1;
206                }
207            }
208        }
209
210        Ok(stats)
211    }
212
213    /// Validate corpus file exists and is accessible
214    pub async fn validate_corpus_file<P: AsRef<Path>>(&self, file_path: P) -> bool {
215        let full_path = self.base_path.join(file_path);
216        
217        match fs::metadata(&full_path).await {
218            Ok(metadata) => {
219                if metadata.is_file() && metadata.len() > 0 {
220                    debug!("✅ Corpus file validated: {}", full_path.display());
221                    true
222                } else {
223                    debug!("❌ Invalid corpus file: {}", full_path.display());
224                    false
225                }
226            }
227            Err(_) => {
228                debug!("❌ Corpus file not found: {}", full_path.display());
229                false
230            }
231        }
232    }
233
234    /// Get file content for corpus validation
235    pub async fn get_file_content<P: AsRef<Path>>(&self, file_path: P) -> Result<String> {
236        let full_path = self.base_path.join(file_path);
237        
238        fs::read_to_string(&full_path).await
239            .map_err(|e| anyhow!(LoadingError::IoError { source: e }))
240    }
241}
242
243/// Corpus statistics
244#[derive(Debug, Clone)]
245pub struct CorpusStats {
246    pub total_files: usize,
247    pub total_lines: usize,
248    pub total_size_bytes: u64,
249    pub language_distribution: HashMap<String, usize>,
250}
251
252impl CorpusStats {
253    /// Get formatted size string
254    pub fn formatted_size(&self) -> String {
255        const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
256        let mut size = self.total_size_bytes as f64;
257        let mut unit_index = 0;
258        
259        while size >= 1024.0 && unit_index < UNITS.len() - 1 {
260            size /= 1024.0;
261            unit_index += 1;
262        }
263        
264        format!("{:.1} {}", size, UNITS[unit_index])
265    }
266
267    /// Get most common languages (top N)
268    pub fn top_languages(&self, n: usize) -> Vec<(String, usize)> {
269        let mut sorted_langs: Vec<_> = self.language_distribution.iter()
270            .map(|(lang, count)| (lang.clone(), *count))
271            .collect();
272        
273        sorted_langs.sort_by(|a, b| b.1.cmp(&a.1));
274        sorted_langs.into_iter().take(n).collect()
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281    use tempfile::TempDir;
282    use tokio::fs::File;
283    use tokio::io::AsyncWriteExt;
284
285    #[tokio::test]
286    async fn test_dataset_loader_creation() {
287        let temp_dir = TempDir::new().unwrap();
288        let loader = DatasetLoader::new(temp_dir.path());
289        
290        assert_eq!(loader.base_path, temp_dir.path());
291    }
292
293    #[tokio::test]
294    async fn test_query_type_parsing() {
295        let loader = DatasetLoader::new(".");
296        
297        assert_eq!(loader.parse_query_type("identifier"), Some(QueryType::Identifier));
298        assert_eq!(loader.parse_query_type("exact_match"), Some(QueryType::ExactMatch));
299        assert_eq!(loader.parse_query_type("structural"), Some(QueryType::Structural));
300        assert_eq!(loader.parse_query_type("invalid"), None);
301    }
302
303    #[tokio::test]
304    async fn test_source_file_detection() {
305        let loader = DatasetLoader::new(".");
306        
307        assert!(loader.is_source_file(Path::new("test.rs")));
308        assert!(loader.is_source_file(Path::new("script.py")));
309        assert!(loader.is_source_file(Path::new("component.tsx")));
310        assert!(!loader.is_source_file(Path::new("data.json")));
311        assert!(!loader.is_source_file(Path::new("image.png")));
312    }
313
314    #[tokio::test]
315    async fn test_directory_skip_logic() {
316        let loader = DatasetLoader::new(".");
317        
318        assert!(loader.should_skip_directory("target"));
319        assert!(loader.should_skip_directory("node_modules"));
320        assert!(loader.should_skip_directory(".git"));
321        assert!(!loader.should_skip_directory("src"));
322        assert!(!loader.should_skip_directory("lib"));
323    }
324
325    #[tokio::test]
326    async fn test_legacy_query_conversion() {
327        let temp_dir = TempDir::new().unwrap();
328        let loader = DatasetLoader::new(temp_dir.path());
329
330        // Create test data - each item should be a JSON object, not an array
331        let legacy_data = serde_json::json!({
332            "query": "function test",
333            "expected_files": ["test.js"],
334            "query_type": "identifier"
335        });
336
337        let legacy_queries = vec![legacy_data];
338        let converted = loader.convert_legacy_queries(legacy_queries).unwrap();
339        
340        assert_eq!(converted.len(), 1);
341        assert_eq!(converted[0].query, "function test");
342        assert_eq!(converted[0].query_type, QueryType::Identifier);
343    }
344
345    #[tokio::test]
346    async fn test_corpus_stats_formatting() {
347        let stats = CorpusStats {
348            total_files: 100,
349            total_lines: 10000,
350            total_size_bytes: 2_048_576, // 2MB
351            language_distribution: {
352                let mut map = HashMap::new();
353                map.insert("rs".to_string(), 50);
354                map.insert("py".to_string(), 30);
355                map.insert("js".to_string(), 20);
356                map
357            },
358        };
359
360        assert_eq!(stats.formatted_size(), "2.0 MB");
361        
362        let top_langs = stats.top_languages(2);
363        assert_eq!(top_langs.len(), 2);
364        assert_eq!(top_langs[0], ("rs".to_string(), 50));
365        assert_eq!(top_langs[1], ("py".to_string(), 30));
366    }
367}
lens_core/benchmark/dataset_loader.rs

lens_core/benchmark/
dataset_loader.rs
