amql_engine/

code_cache.rs

//! Per-file cache of parsed `CodeElement` trees with mtime invalidation.
//!
//! The code cache lazily parses source files within a given scope, caches
//! results keyed by relative path, and invalidates entries when the file's
//! mtime changes.

use crate::resolver::{CodeElement, ResolverRegistry};
use crate::types::{ProjectRoot, RelativePath, Scope};
use rayon::prelude::*;
use rustc_hash::FxHashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::time::SystemTime;

/// Cached parse result for a single source file.
struct CodeFileEntry {
    root: CodeElement,
    mtime: SystemTime,
}

/// Per-file cache of parsed `CodeElement` trees with mtime invalidation.
pub struct CodeCache {
    index: FxHashMap<RelativePath, CodeFileEntry>,
    project_root: ProjectRoot,
}

// Future: DependencyTracker hook — per-file mtime invalidation is sufficient
// for now. Disjoint-set tracking becomes relevant when resolvers follow
// cross-file dependencies (e.g. TypeScript re-exports).

impl CodeCache {
    /// Create an empty cache rooted at the given project directory.
    pub fn new(project_root: &Path) -> Self {
        Self {
            index: FxHashMap::default(),
            project_root: ProjectRoot::from(project_root),
        }
    }

    /// Lazily parse all code files within scope. Skips cached+fresh files.
    /// Uses rayon for parallel parsing of stale/missing files.
    pub fn ensure_scope(&mut self, scope: &Scope, resolvers: &ResolverRegistry) {
        let source_files = glob_source_files(&self.project_root, scope, resolvers);

        // Partition into files that need (re-)parsing
        let project_root = &self.project_root;
        let needs_parse: Vec<_> = source_files
            .into_iter()
            .filter(|abs_path| {
                let rel = crate::paths::relative(project_root, abs_path);
                if let Some(entry) = self.index.get(&*rel) {
                    // Check mtime — skip if fresh
                    match fs::metadata(abs_path) {
                        Ok(meta) => {
                            let mtime = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
                            mtime > entry.mtime
                        }
                        Err(_) => true,
                    }
                } else {
                    true // not cached
                }
            })
            .collect();

        if needs_parse.is_empty() {
            return;
        }

        let parsed: Vec<_> = needs_parse
            .par_iter()
            .filter_map(|abs_path| {
                let resolver = resolvers.get_for_file(abs_path)?;
                let root = resolver.resolve(abs_path).ok()?;
                let mtime = fs::metadata(abs_path)
                    .and_then(|m| m.modified())
                    .unwrap_or(SystemTime::UNIX_EPOCH);
                let rel = crate::paths::relative(project_root, abs_path);
                let root = relativize_source_files(root, &rel);
                Some((rel, CodeFileEntry { root, mtime }))
            })
            .collect();

        for (rel, entry) in parsed {
            self.index.insert(rel, entry);
        }
    }

    /// Check mtime and re-parse a single file if stale. Returns true if re-parsed.
    pub fn refresh_if_stale(
        &mut self,
        rel_path: &RelativePath,
        resolvers: &ResolverRegistry,
    ) -> bool {
        let abs_path = self.project_root.join(rel_path);

        if let Some(entry) = self.index.get(rel_path) {
            match fs::metadata(&abs_path) {
                Ok(meta) => {
                    let new_mtime = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
                    if new_mtime <= entry.mtime {
                        return false; // still fresh
                    }
                }
                Err(_) => {
                    self.index.remove(rel_path);
                    return true;
                }
            }
        }

        // Parse (or re-parse)
        let resolver = match resolvers.get_for_file(&abs_path) {
            Some(r) => r,
            None => return false,
        };
        match resolver.resolve(&abs_path) {
            Ok(root) => {
                let mtime = fs::metadata(&abs_path)
                    .and_then(|m| m.modified())
                    .unwrap_or(SystemTime::UNIX_EPOCH);
                let rel = rel_path.clone();
                let root = relativize_source_files(root, &rel);
                self.index.insert(rel, CodeFileEntry { root, mtime });
                true
            }
            Err(_) => false,
        }
    }

    /// Get cached code element tree for a file.
    pub fn get(&self, rel_path: &RelativePath) -> Option<&CodeElement> {
        self.index.get(rel_path).map(|e| &e.root)
    }

    /// Inject test data directly (for unit tests).
    #[cfg(test)]
    pub fn inject_test_data(&mut self, rel_path: &str, root: CodeElement) {
        self.index.insert(
            RelativePath::from(rel_path),
            CodeFileEntry {
                root,
                mtime: SystemTime::now(),
            },
        );
    }

    /// Get all cached (rel_path, CodeElement) pairs within a scope prefix.
    pub fn elements_in_scope(&self, scope: &Scope) -> Vec<(&RelativePath, &CodeElement)> {
        self.index
            .iter()
            .filter(|(rel, _)| scope.is_empty() || rel.starts_with(&**scope))
            .map(|(rel, entry)| (rel, &entry.root))
            .collect()
    }
}

/// Return a CodeElement tree with `source.file` set to the relative path.
fn relativize_source_files(element: CodeElement, rel: &RelativePath) -> CodeElement {
    CodeElement {
        source: crate::resolver::SourceLocation {
            file: rel.clone(),
            ..element.source
        },
        children: element
            .children
            .into_iter()
            .map(|c| relativize_source_files(c, rel))
            .collect(),
        ..element
    }
}

/// Glob source files within a scope, filtered to resolver-supported extensions.
/// Reuses `ignore::WalkBuilder` pattern from `glob_annotation_files`.
pub fn glob_source_files(
    project_root: &Path,
    scope: &Scope,
    resolvers: &ResolverRegistry,
) -> Vec<PathBuf> {
    let walk_root = if scope.is_empty() {
        project_root.to_path_buf()
    } else {
        let candidate = project_root.join(&**scope);
        if candidate.is_file() {
            // Single file scope — just return it if resolver supports it
            if resolvers.has_resolver_for(&candidate) {
                return vec![candidate];
            }
            return vec![];
        }
        candidate
    };

    if !walk_root.exists() {
        return vec![];
    }

    let results = Mutex::new(Vec::new());

    let walker = ignore::WalkBuilder::new(&walk_root)
        .hidden(false)
        .git_ignore(true)
        .git_global(true)
        .git_exclude(true)
        .filter_entry(|entry| {
            let name = entry.file_name().to_string_lossy();
            if entry.file_type().is_some_and(|ft| ft.is_dir()) {
                return !crate::paths::SKIP_DIRS.contains(&name.as_ref());
            }
            true
        })
        .build_parallel();

    walker.run(|| {
        Box::new(|entry| {
            if let Ok(entry) = entry {
                let path = entry.path();
                if path.is_file()
                    && resolvers.has_resolver_for(path)
                    && !path
                        .file_name()
                        .is_some_and(|n| n.to_string_lossy().ends_with(".aqm"))
                {
                    results.lock().unwrap().push(path.to_path_buf());
                }
            }
            ignore::WalkState::Continue
        })
    });

    results.into_inner().unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::resolver::RustResolver;

    fn project_root() -> PathBuf {
        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .unwrap()
            .parent()
            .unwrap()
            .to_path_buf()
    }

    fn make_resolvers() -> ResolverRegistry {
        let mut reg = ResolverRegistry::new();
        reg.register(Box::new(RustResolver));
        reg
    }

    #[test]
    fn globs_source_files() {
        // Arrange
        let root = project_root();
        let resolvers = make_resolvers();

        // Act
        let dir_files =
            glob_source_files(&root, &Scope::from("crates/amql-engine/src/"), &resolvers);
        let single_file = glob_source_files(
            &root,
            &Scope::from("crates/amql-engine/src/selector.rs"),
            &resolvers,
        );
        let all_files = glob_source_files(&root, &Scope::from(""), &resolvers);

        // Assert
        assert!(!dir_files.is_empty());
        assert!(dir_files
            .iter()
            .all(|p| p.extension().is_some_and(|e| e == "rs")));

        assert_eq!(single_file.len(), 1);
        assert!(single_file[0].ends_with("selector.rs"));

        assert!(
            all_files.len() > 5,
            "Should find many .rs files project-wide"
        );
    }

    #[test]
    fn caches_and_scopes() {
        // Arrange
        let root = project_root();
        let resolvers = make_resolvers();
        let mut cache = CodeCache::new(&root);

        // Act
        cache.ensure_scope(&Scope::from("crates/amql-engine/src/"), &resolvers);

        // Assert
        let selector = cache.get(&RelativePath::from("crates/amql-engine/src/selector.rs"));
        assert!(selector.is_some());
        assert_eq!(selector.unwrap().tag, "module");

        let all = cache.elements_in_scope(&Scope::from("crates/amql-engine/src/"));
        assert!(all.len() >= 5, "Should have multiple source files cached");

        let selector_only = cache.elements_in_scope(&Scope::from("crates/amql-engine/src/selector"));
        assert_eq!(selector_only.len(), 1);

        let count1 = cache.elements_in_scope(&Scope::from("")).len();
        cache.ensure_scope(&Scope::from("crates/amql-engine/src/"), &resolvers);
        let count2 = cache.elements_in_scope(&Scope::from("")).len();
        assert_eq!(count1, count2, "Second ensure_scope should be idempotent");
    }
}