agcodex_ast/
lib.rs

1//! AGCodex AST Module - Tree-sitter based code intelligence
2//!
3//! This crate provides comprehensive AST parsing, analysis, and compaction
4//! for 50+ programming languages using tree-sitter.
5
6pub mod compactor;
7pub mod error;
8pub mod language_registry;
9pub mod parser_cache;
10pub mod semantic_index;
11pub mod types;
12
13pub use compactor::AstCompactor;
14pub use compactor::CompressionLevel;
15pub use error::AstError;
16pub use error::AstResult;
17pub use language_registry::Language;
18pub use language_registry::LanguageRegistry;
19pub use parser_cache::ParserCache;
20pub use semantic_index::SemanticIndex;
21pub use semantic_index::Symbol;
22pub use semantic_index::SymbolKind;
23pub use types::AstNode;
24pub use types::AstNodeKind;
25pub use types::ParsedAst;
26pub use types::SourceLocation;
27
28use std::path::Path;
29use std::sync::Arc;
30use tokio::sync::RwLock;
31
32/// Main AST engine for AGCodex
33#[derive(Debug)]
34pub struct AstEngine {
35    registry: Arc<LanguageRegistry>,
36    cache: Arc<RwLock<ParserCache>>,
37    compactor: Arc<AstCompactor>,
38    semantic_index: Arc<RwLock<SemanticIndex>>,
39}
40
41impl AstEngine {
42    /// Create a new AST engine with specified compression level
43    pub fn new(compression_level: CompressionLevel) -> Self {
44        Self {
45            registry: Arc::new(LanguageRegistry::new()),
46            cache: Arc::new(RwLock::new(ParserCache::new(1024 * 1024 * 100))), // 100MB default
47            compactor: Arc::new(AstCompactor::new(compression_level)),
48            semantic_index: Arc::new(RwLock::new(SemanticIndex::new())),
49        }
50    }
51
52    /// Parse a file and return its AST
53    pub async fn parse_file(&self, path: &Path) -> AstResult<ParsedAst> {
54        // Check cache first
55        {
56            let mut cache = self.cache.write().await;
57            if let Some(ast) = cache.get(path) {
58                return Ok((*ast).clone());
59            }
60        }
61
62        // Detect language and parse
63        let language = self.registry.detect_language(path)?;
64        let content = tokio::fs::read_to_string(path)
65            .await
66            .map_err(|e| AstError::IoError(e.to_string()))?;
67
68        let parsed = self.registry.parse(&language, &content)?;
69
70        // Cache the result
71        {
72            let mut cache = self.cache.write().await;
73            cache.insert(path.to_path_buf(), parsed.clone());
74        }
75
76        // Update semantic index
77        {
78            let mut index = self.semantic_index.write().await;
79            index.index_ast(path, &parsed)?;
80        }
81
82        Ok(parsed)
83    }
84
85    /// Compact code using AI Distiller-style compression
86    pub async fn compact_code(&self, path: &Path) -> AstResult<String> {
87        let ast = self.parse_file(path).await?;
88        self.compactor.compact(&ast)
89    }
90
91    /// Search for symbols across the codebase
92    pub async fn search_symbols(&self, query: &str) -> AstResult<Vec<Symbol>> {
93        let index = self.semantic_index.read().await;
94        Ok(index.search(query))
95    }
96
97    /// Get call graph for a function
98    pub async fn get_call_graph(&self, path: &Path, function_name: &str) -> AstResult<Vec<Symbol>> {
99        let index = self.semantic_index.read().await;
100        Ok(index.get_call_graph(path, function_name))
101    }
102
103    /// Clear the cache
104    pub async fn clear_cache(&self) {
105        let mut cache = self.cache.write().await;
106        cache.clear();
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use std::fs;
114    use tempfile::tempdir;
115
116    #[tokio::test]
117    async fn test_ast_engine_basic() {
118        let engine = AstEngine::new(CompressionLevel::Medium);
119        let dir = tempdir().unwrap();
120        let file_path = dir.path().join("test.rs");
121
122        fs::write(
123            &file_path,
124            "fn main() {\n    println!(\"Hello, world!\");\n}",
125        )
126        .unwrap();
127
128        let ast = engine.parse_file(&file_path).await.unwrap();
129        assert!(!ast.root_node.kind.is_empty());
130    }
131
132    #[tokio::test]
133    async fn test_compression_levels() {
134        let light_engine = AstEngine::new(CompressionLevel::Light);
135        let hard_engine = AstEngine::new(CompressionLevel::Hard);
136
137        let dir = tempdir().unwrap();
138        let file_path = dir.path().join("test.py");
139
140        let code = r#"
141def calculate_fibonacci(n):
142    """Calculate fibonacci number"""
143    if n <= 1:
144        return n
145    else:
146        return calculate_fibonacci(n-1) + calculate_fibonacci(n-2)
147
148class MathOperations:
149    def __init__(self):
150        self.cache = {}
151    
152    def add(self, a, b):
153        return a + b
154    
155    def multiply(self, a, b):
156        return a * b
157"#;
158
159        fs::write(&file_path, code).unwrap();
160
161        let light_compact = light_engine.compact_code(&file_path).await.unwrap();
162        let hard_compact = hard_engine.compact_code(&file_path).await.unwrap();
163
164        // Hard compression should be shorter
165        assert!(hard_compact.len() < light_compact.len());
166    }
167}
agcodex_ast/lib.rs

agcodex_ast/
lib.rs
