amql_engine/store/

mod.rs

//! Load, index, and query `.aqm` annotation sidecar files.
//!
//! The annotation store is a pure in-memory index of parsed annotation data.
//! It never parses source code — only AQL sidecar files.

#[cfg(feature = "fs")]
mod glob;
mod parse;

#[cfg(feature = "fs")]
pub use glob::{glob_annotation_files, glob_aql_files};

use crate::error::AqlError;
use crate::matcher::{self, Matchable};
use crate::selector::parse_selector;
use crate::sidecar::SidecarLocator;
use crate::types::{AttrName, Binding, ProjectRoot, RelativePath, Scope, TagName};
use rustc_hash::FxHashMap;
use serde::Serialize;
use serde_json::Value as JsonValue;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::SystemTime;

#[cfg(feature = "fs")]
use rayon::prelude::*;

/// A single annotation entry from an `.aqm` file.
#[cfg_attr(feature = "ts", derive(ts_rs::TS))]
#[cfg_attr(feature = "flow", derive(flowjs_rs::Flow))]
#[cfg_attr(feature = "ts", ts(export))]
#[cfg_attr(feature = "flow", flow(export))]
#[derive(Debug, Clone, Serialize)]
pub struct Annotation {
    /// Annotation tag (e.g. "controller", "react-hook").
    pub tag: TagName,
    /// User-defined attributes from the element's attributes.
    #[cfg_attr(
        feature = "ts",
        ts(as = "std::collections::HashMap<AttrName, serde_json::Value>")
    )]
    #[cfg_attr(
        feature = "flow",
        flow(as = "std::collections::HashMap<AttrName, serde_json::Value>")
    )]
    pub attrs: FxHashMap<AttrName, JsonValue>,
    /// Binding key for joining annotations to code elements.
    /// From the `bind` attribute on the element.
    pub binding: Binding,
    /// Relative path of the source file this annotation targets.
    pub file: RelativePath,
    /// Nested child annotations.
    pub children: Vec<Annotation>,
}

/// Input format for loading annotations: a file path and its XML content.
#[cfg_attr(feature = "ts", derive(ts_rs::TS))]
#[cfg_attr(feature = "flow", derive(flowjs_rs::Flow))]
#[cfg_attr(feature = "ts", ts(export))]
#[cfg_attr(feature = "flow", flow(export))]
#[derive(Debug, Clone, serde::Deserialize, Serialize)]
pub struct FileEntry {
    /// Relative file path.
    pub file: RelativePath,
    /// AQL XML content.
    pub xml: String,
}

/// Internal storage of file annotations with mtime for cache invalidation.
struct CachedFileEntry {
    annotations: Vec<Annotation>,
    #[cfg_attr(not(feature = "fs"), allow(dead_code))]
    mtime: SystemTime,
    #[cfg_attr(not(feature = "fs"), allow(dead_code))]
    sidecar_path: PathBuf,
}

/// In-memory index of annotation sidecar files.
pub struct AnnotationStore {
    index: FxHashMap<RelativePath, CachedFileEntry>,
    #[cfg_attr(not(feature = "fs"), allow(dead_code))]
    project_root: ProjectRoot,
    locator: Arc<dyn SidecarLocator>,
}

impl AnnotationStore {
    /// Create an empty store rooted at the given project directory.
    /// Uses colocated sidecar strategy by default.
    #[cfg(feature = "fs")]
    pub fn new(project_root: &Path) -> Self {
        Self {
            index: FxHashMap::default(),
            project_root: ProjectRoot::from(project_root),
            locator: Arc::new(crate::sidecar::ColocatedLocator),
        }
    }

    /// Create an empty in-memory store (no filesystem access).
    #[cfg(not(feature = "fs"))]
    pub fn new(project_root: &Path) -> Self {
        Self {
            index: FxHashMap::default(),
            project_root: ProjectRoot::from(project_root),
            locator: Arc::new(crate::sidecar::InMemoryLocator),
        }
    }

    /// Create an empty store with a custom sidecar locator strategy.
    pub fn with_locator(project_root: &Path, locator: Arc<dyn SidecarLocator>) -> Self {
        Self {
            index: FxHashMap::default(),
            project_root: ProjectRoot::from(project_root),
            locator,
        }
    }

    /// Return the sidecar locator used by this store.
    pub fn locator(&self) -> &dyn SidecarLocator {
        &*self.locator
    }

    /// Discover and load all sidecar files using the configured locator.
    /// Returns parse errors for files that failed to load.
    #[cfg(feature = "fs")]
    pub fn load_all_from_locator(&mut self) -> Vec<String> {
        let pairs = self.locator.discover(&self.project_root);
        let results: Vec<_> = pairs
            .par_iter()
            .filter_map(|(sidecar_path, rel_source)| {
                match (
                    std::fs::metadata(sidecar_path),
                    std::fs::read_to_string(sidecar_path),
                ) {
                    (Ok(meta), Ok(raw)) => {
                        let mtime = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
                        Some((
                            sidecar_path.clone(),
                            rel_source.clone(),
                            mtime,
                            parse::parse_sidecar(&raw, rel_source),
                        ))
                    }
                    _ => None,
                }
            })
            .collect();

        let mut errors = Vec::new();
        for (sidecar_path, rel_source, mtime, parse_result) in results {
            match parse_result {
                Ok(annotations) => {
                    self.index.insert(
                        rel_source,
                        CachedFileEntry {
                            annotations,
                            mtime,
                            sidecar_path,
                        },
                    );
                }
                Err(e) => errors.push(e),
            }
        }
        errors
    }

    /// Load from pre-resolved sidecar paths (backward-compatible).
    /// Returns parse errors for files that failed to load.
    #[cfg(feature = "fs")]
    pub fn load_all(&mut self, sidecar_paths: &[PathBuf]) -> Vec<String> {
        let project_root = &self.project_root;
        let results: Vec<_> = sidecar_paths
            .par_iter()
            .filter_map(|path| {
                let stem = path.file_stem()?.to_string_lossy().to_string();
                let parent = path.parent()?;
                let source_path = parent.join(&stem);
                let rel_source = crate::paths::relative(project_root, &source_path);

                match (std::fs::metadata(path), std::fs::read_to_string(path)) {
                    (Ok(meta), Ok(raw)) => {
                        let mtime = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
                        let parsed = parse::parse_sidecar(&raw, &rel_source);
                        Some((path.to_path_buf(), rel_source, mtime, parsed))
                    }
                    _ => None,
                }
            })
            .collect();

        let mut errors = Vec::new();
        for (sidecar_path, rel_source, mtime, parse_result) in results {
            match parse_result {
                Ok(annotations) => {
                    self.index.insert(
                        rel_source,
                        CachedFileEntry {
                            annotations,
                            mtime,
                            sidecar_path,
                        },
                    );
                }
                Err(e) => errors.push(e),
            }
        }
        errors
    }

    /// Load or reload a single sidecar file.
    /// Returns `Err` if the sidecar exists but fails to parse.
    /// Returns `Ok(())` if loaded successfully or the file was removed.
    #[cfg(feature = "fs")]
    pub fn load_file(&mut self, sidecar_path: &Path) -> Result<(), String> {
        let stem = sidecar_path
            .file_stem()
            .map(|s| s.to_string_lossy().to_string())
            .unwrap_or_default();
        let parent = sidecar_path.parent().unwrap_or(Path::new(""));
        let source_path = parent.join(&stem);

        let rel_source = crate::paths::relative(&self.project_root, &source_path);

        match (
            std::fs::metadata(sidecar_path),
            std::fs::read_to_string(sidecar_path),
        ) {
            (Ok(meta), Ok(raw)) => {
                let mtime = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
                let annotations = match parse::parse_sidecar(&raw, &rel_source) {
                    Ok(a) => a,
                    Err(e) => {
                        self.index.remove(&*rel_source);
                        return Err(e);
                    }
                };
                self.index.insert(
                    rel_source,
                    CachedFileEntry {
                        annotations,
                        mtime,
                        sidecar_path: sidecar_path.to_path_buf(),
                    },
                );
                Ok(())
            }
            _ => {
                self.index.remove(&*rel_source);
                Ok(())
            }
        }
    }

    /// Check mtime and reload if the sidecar changed.
    /// Entries with no sidecar path (extractor-only or load_xml) are skipped.
    #[cfg(feature = "fs")]
    pub fn refresh_if_stale(&mut self, rel_source: &RelativePath) -> bool {
        let entry = match self.index.get(rel_source) {
            Some(e) => e,
            None => return false,
        };

        let sidecar_path = entry.sidecar_path.clone();

        // Extractor-only or load_xml entries have no sidecar file to check
        if sidecar_path.as_os_str().is_empty() {
            return false;
        }

        match std::fs::metadata(&sidecar_path) {
            Ok(meta) => {
                let new_mtime = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
                if new_mtime > entry.mtime {
                    // Parse errors on refresh silently remove the entry
                    let _ = self.load_file(&sidecar_path);
                    return true;
                }
            }
            Err(_) => {
                self.index.remove(rel_source);
                return true;
            }
        }
        false
    }

    /// Check whether a file's sidecar is stale (mtime changed on disk).
    ///
    /// Returns `true` if the sidecar file on disk has a newer mtime than
    /// the cached entry, or if the sidecar file no longer exists.
    /// Returns `false` for extractor-only entries (no sidecar path).
    #[cfg(feature = "fs")]
    pub fn is_stale(&self, rel_source: &RelativePath) -> bool {
        let entry = match self.index.get(rel_source) {
            Some(e) => e,
            None => return false,
        };
        if entry.sidecar_path.as_os_str().is_empty() {
            return false;
        }
        match std::fs::metadata(&entry.sidecar_path) {
            Ok(meta) => {
                let disk_mtime = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
                disk_mtime > entry.mtime
            }
            Err(_) => true,
        }
    }

    /// Get cached code element tree for a file.
    ///
    /// Accepts a source path (`src/api.ts`) or sidecar path (`src/api.aqm`).
    /// Falls back to stem-prefix matching when an exact key is not found.
    pub fn get_file_annotations(&self, rel_source: &str) -> Vec<&Annotation> {
        // Exact match first
        if let Some(entry) = self.index.get(rel_source) {
            return entry.annotations.iter().collect();
        }

        // Strip .aqm or source extension to get the stem for matching
        let stem = rel_source
            .strip_suffix(".aqm")
            .or_else(|| rel_source.rfind('.').map(|dot| &rel_source[..dot]))
            .unwrap_or(rel_source);

        // Find the first key whose stem matches
        self.index
            .iter()
            .find(|(key, _)| {
                let key_str: &str = key.as_ref();
                let key_stem = match key_str.rfind('.') {
                    Some(dot) => &key_str[..dot],
                    None => key_str,
                };
                key_stem == stem
            })
            .map(|(_, entry)| entry.annotations.iter().collect())
            .unwrap_or_default()
    }

    /// Get all annotations across all files.
    pub fn get_all_annotations(&self) -> Vec<&Annotation> {
        self.index
            .values()
            .flat_map(|e| e.annotations.iter())
            .collect()
    }

    /// Return annotations from files whose relative path starts with `scope`.
    pub fn get_scope_annotations(&self, scope: &Scope) -> Vec<&Annotation> {
        self.index
            .iter()
            .filter(|(key, _)| {
                let k: &str = key.as_ref();
                k.starts_with(&**scope)
            })
            .flat_map(|(_, e)| e.annotations.iter())
            .collect()
    }

    /// Query annotations by selector string.
    /// If `file` is provided, only search that file's annotations.
    /// If `scope` is provided, only search annotations under that directory prefix.
    /// Supports combinator selectors for nested annotations (e.g. `parser > helper`).
    pub fn select(
        &self,
        selector_str: &str,
        file: Option<&str>,
        scope: Option<&str>,
        opts: Option<&crate::query_options::QueryOptions>,
    ) -> Result<Vec<Annotation>, AqlError> {
        let selector = parse_selector(selector_str)?;

        let annotations: Vec<&Annotation> = match file {
            Some(f) => self.get_file_annotations(f),
            None => match scope {
                Some(s) if !s.is_empty() => self.get_scope_annotations(&Scope::from(s)),
                _ => self.get_all_annotations(),
            },
        };

        let flat = flatten_annotations(&annotations);
        let matchable_refs: Vec<&dyn Matchable> =
            flat.iter().map(|n| n as &dyn Matchable).collect();
        let parent_indices: Vec<Option<usize>> = flat.iter().map(|n| n.parent_idx).collect();
        let matched_indices =
            matcher::filter_by_selector_indexed(&matchable_refs, &parent_indices, &selector);

        let results: Vec<Annotation> = matched_indices
            .into_iter()
            .map(|idx| flat[idx].ann.clone())
            .collect();

        let results = match opts {
            Some(o) => crate::query_options::apply_to_annotations(results, o),
            None => results,
        };

        Ok(results)
    }

    /// Refresh all sidecar files whose rel_source starts with `scope`.
    /// Also discovers newly added sidecar files via the locator.
    /// Returns parse errors encountered during discovery.
    /// If scope is empty, refreshes all files.
    #[cfg(feature = "fs")]
    pub fn refresh_scope(&mut self, scope: &Scope) -> Vec<String> {
        let mut errors = Vec::new();

        // Refresh existing entries
        let keys: Vec<RelativePath> = self
            .index
            .keys()
            .filter(|k| scope.is_empty() || k.starts_with(&**scope))
            .cloned()
            .collect();
        for key in keys {
            self.refresh_if_stale(&key);
        }

        // Discover new sidecar files not yet in the index
        let pairs = self.locator.discover(&self.project_root);
        for (sidecar_path, rel_source) in pairs {
            if !(scope.is_empty() || rel_source.starts_with(&**scope)) {
                continue;
            }
            if self.index.contains_key(&*rel_source) {
                continue;
            }
            // New file — load it
            if let (Ok(meta), Ok(raw)) = (
                std::fs::metadata(&sidecar_path),
                std::fs::read_to_string(&sidecar_path),
            ) {
                let mtime = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
                match parse::parse_sidecar(&raw, &rel_source) {
                    Ok(annotations) => {
                        self.index.insert(
                            rel_source,
                            CachedFileEntry {
                                annotations,
                                mtime,
                                sidecar_path,
                            },
                        );
                    }
                    Err(e) => errors.push(e),
                }
            }
        }
        errors
    }

    /// Get all annotations grouped by file within a scope prefix.
    pub fn annotations_in_scope(&self, scope: &Scope) -> Vec<(&RelativePath, Vec<&Annotation>)> {
        self.index
            .iter()
            .filter(|(rel, _)| scope.is_empty() || rel.starts_with(&**scope))
            .map(|(rel, entry)| (rel, entry.annotations.iter().collect()))
            .collect()
    }

    /// Number of loaded annotation files.
    pub fn file_count(&self) -> usize {
        self.index.len()
    }

    /// All source file paths that have annotations.
    pub fn annotated_files(&self) -> Vec<RelativePath> {
        self.index.keys().cloned().collect()
    }

    /// Load annotations from an XML string for a given relative source path.
    /// Unlike `load_file`, this does not touch the filesystem.
    /// Returns `Err` if the XML is malformed.
    pub fn load_xml(&mut self, rel_source: &RelativePath, xml: &str) -> Result<(), AqlError> {
        let annotations = parse::parse_sidecar(xml, rel_source).map_err(AqlError::from)?;
        self.index.insert(
            rel_source.clone(),
            CachedFileEntry {
                annotations,
                mtime: SystemTime::UNIX_EPOCH,
                sidecar_path: PathBuf::new(),
            },
        );
        Ok(())
    }

    /// Load annotations from extractor output, merging with existing sidecar data.
    ///
    /// For bindings that already exist in sidecar files, sidecar attributes take
    /// precedence — extractor attrs are only used as defaults.
    pub fn load_extractor_output(&mut self, annotations: Vec<Annotation>) {
        for ann in annotations {
            let rel = ann.file.clone();
            let entry = self.index.entry(rel).or_insert_with(|| CachedFileEntry {
                annotations: Vec::new(),
                mtime: SystemTime::UNIX_EPOCH,
                sidecar_path: PathBuf::new(),
            });

            // Check if a sidecar annotation already covers this binding+tag
            let existing = entry.annotations.iter_mut().find(|a| {
                !a.binding.is_empty()
                    && !ann.binding.is_empty()
                    && a.binding == ann.binding
                    && a.tag == ann.tag
            });

            match existing {
                Some(existing_ann) => {
                    // Sidecar overrides extractor: only fill in attrs the sidecar lacks
                    for (key, value) in &ann.attrs {
                        existing_ann
                            .attrs
                            .entry(key.clone())
                            .or_insert(value.clone());
                    }
                }
                None => {
                    entry.annotations.push(ann);
                }
            }
        }
    }

    /// Inject test data directly (for unit tests).
    #[cfg(test)]
    pub fn inject_test_data(&mut self, rel_source: &RelativePath, annotations: Vec<Annotation>) {
        let sidecar_rel = self.locator.sidecar_for(rel_source);
        let sidecar_path = self.project_root.join(AsRef::<Path>::as_ref(&sidecar_rel));

        self.index.insert(
            rel_source.clone(),
            CachedFileEntry {
                annotations,
                mtime: SystemTime::now(),
                sidecar_path,
            },
        );
    }
}

// ---------------------------------------------------------------------------
// Matchable wrapper for annotations (zero-copy: borrows from &Annotation)
// ---------------------------------------------------------------------------

struct AnnotationNode<'a> {
    ann: &'a Annotation,
    parent_idx: Option<usize>,
}

impl Matchable for AnnotationNode<'_> {
    fn tag(&self) -> &TagName {
        &self.ann.tag
    }
    fn attrs(&self) -> &FxHashMap<AttrName, JsonValue> {
        &self.ann.attrs
    }
    fn parent(&self) -> Option<&dyn Matchable> {
        None
    }
}

fn flatten_annotations<'a>(annotations: &[&'a Annotation]) -> Vec<AnnotationNode<'a>> {
    let mut result = Vec::new();

    fn walk<'a>(
        ann: &'a Annotation,
        parent_idx: Option<usize>,
        result: &mut Vec<AnnotationNode<'a>>,
    ) {
        let idx = result.len();
        result.push(AnnotationNode { ann, parent_idx });
        for child in &ann.children {
            walk(child, Some(idx), result);
        }
    }

    for ann in annotations {
        walk(ann, None, &mut result);
    }
    result
}

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

    #[test]
    fn loads_xml_from_string() {
        // Arrange
        let mut store = AnnotationStore::new(Path::new("/project"));
        let xml = r#"<controller method="POST" bind="handle_create" />"#;

        // Act
        store
            .load_xml(&RelativePath::from("src/api.ts"), xml)
            .unwrap();
        let results = store.select("controller", None, None, None).unwrap();

        // Assert
        assert_eq!(store.file_count(), 1, "should have one loaded file");
        assert_eq!(results.len(), 1, "should find one controller annotation");
        assert_eq!(results[0].tag, "controller", "tag should be controller");
        assert_eq!(
            results[0].file, "src/api.ts",
            "file should match rel_source"
        );
        assert_eq!(
            results[0].binding, "handle_create",
            "binding should come from bind attr"
        );
    }

    #[test]
    fn selects_annotations() {
        // Arrange
        let mut store = AnnotationStore::new(Path::new("/project"));
        store.inject_test_data(
            &RelativePath::from("src/api.ts"),
            vec![
                Annotation {
                    tag: TagName::from("controller"),
                    attrs: {
                        let mut m = FxHashMap::default();
                        m.insert(
                            AttrName::from("method"),
                            JsonValue::String("POST".to_string()),
                        );
                        m
                    },
                    binding: Binding::from(""),
                    file: RelativePath::from("src/api.ts"),
                    children: vec![],
                },
                Annotation {
                    tag: TagName::from("react-hook"),
                    attrs: FxHashMap::default(),
                    binding: Binding::from(""),
                    file: RelativePath::from("src/api.ts"),
                    children: vec![],
                },
            ],
        );
        store.inject_test_data(
            &RelativePath::from("src/b.ts"),
            vec![Annotation {
                tag: TagName::from("controller"),
                attrs: FxHashMap::default(),
                binding: Binding::from(""),
                file: RelativePath::from("src/b.ts"),
                children: vec![],
            }],
        );

        // Act
        let by_tag = store.select("controller", None, None, None).unwrap();
        let by_attr_match = store
            .select(r#"controller[method="POST"]"#, None, None, None)
            .unwrap();
        let by_attr_miss = store
            .select(r#"controller[method="GET"]"#, None, None, None)
            .unwrap();
        let scoped = store
            .select("controller", Some("src/api.ts"), None, None)
            .unwrap();

        // Assert
        assert_eq!(by_tag.len(), 2);
        assert!(by_tag.iter().all(|r| r.tag == "controller"));

        assert_eq!(by_attr_match.len(), 1);
        assert_eq!(by_attr_miss.len(), 0);

        assert_eq!(scoped.len(), 1);
        assert_eq!(scoped[0].file, "src/api.ts");
    }
}