amql_engine/query/

mod.rs

//! Unified query pipeline: code + annotations -> matched results.
//!
//! Queries dynamically parse code within a scope, join code elements with
//! annotation sidecars via binding keys, and match unified nodes against
//! CSS-like selectors.

mod node;
mod rewrite;

use crate::code_cache::CodeCache;
use crate::error::AqlError;
use crate::matcher;
use crate::matcher::Matchable;
use crate::resolver::ResolverRegistry;
use crate::selector::parse_selector;
use crate::store::AnnotationStore;
use crate::types::{AttrName, Binding, CodeElementName, RelativePath, Scope, TagName};
use node::{build_unified_nodes, flatten_annotations_recursive};
use rustc_hash::FxHashMap;
use serde::Serialize;
use serde_json::Value as JsonValue;

use crate::resolver::SourceLocation;

/// Synthetic attribute name injected into unified nodes to carry the annotation tag.
const ANN_TAG_ATTR: &str = "ann_tag";

/// Enriched query result.
#[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 QueryResult {
    /// The matched code element.
    pub code_element: CodeElementSummary,
    /// Annotations bound to this code element.
    pub annotations: Vec<AnnotationSummary>,
}

/// Summary of a matched code element.
#[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 CodeElementSummary {
    /// Element kind (e.g. "function", "struct").
    pub tag: TagName,
    /// Name of the code element.
    pub name: CodeElementName,
    /// Relative path to the source file.
    pub file: RelativePath,
    /// Full source location (file, line, column).
    pub source: SourceLocation,
    /// Attributes extracted from the source code.
    #[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>,
}

/// Summary of an annotation attached to a code element.
#[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 AnnotationSummary {
    /// Annotation tag (e.g. "controller", "react-hook").
    pub tag: TagName,
    /// Annotation attributes from the sidecar file.
    #[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 that linked this annotation to its code element.
    pub binding: Binding,
}

/// Main query entry point.
///
/// Pipeline:
/// 1. Parse selector. If tag is not in CODE_TAGS, rewrite to `[ann_tag="{tag}"]` + original attrs.
/// 2. `code_cache.ensure_scope(scope, resolvers)` — lazy parse code files.
/// 3. Refresh annotation sidecars in scope.
/// 4. For each file in scope: flatten CodeElement tree, look up annotations, build UnifiedNode list.
/// 5. Match unified nodes against selector via `filter_by_selector`.
/// 6. Build QueryResult from matched nodes (deduplicate by source location).
#[must_use = "running a query is useless without inspecting the results"]
#[cfg_attr(dylint_lib = "amql_lints", allow(no_mut_params))] // mtime-cache pattern: caches refresh from disk
pub fn unified_query(
    selector_str: &str,
    scope: &Scope,
    code_cache: &mut CodeCache,
    store: &mut AnnotationStore,
    resolvers: &ResolverRegistry,
    opts: Option<&crate::query_options::QueryOptions>,
) -> Result<Vec<QueryResult>, AqlError> {
    // Step 1: Parse and potentially rewrite selector
    let code_tags = resolvers.all_code_tags();
    let rewritten = rewrite::rewrite_selector(selector_str, &code_tags);
    let selector = parse_selector(&rewritten)?;

    // Step 2: Ensure code files in scope are parsed and cached
    code_cache.ensure_scope(scope, resolvers);

    // Step 3: Refresh annotation sidecars in scope
    let _ = store.refresh_scope(scope);

    // Step 4: Build unified nodes
    let elements = code_cache.elements_in_scope(scope);
    let ann_by_file = store.annotations_in_scope(scope);
    let ann_map: FxHashMap<&RelativePath, Vec<&crate::store::Annotation>> =
        ann_by_file.into_iter().collect();

    let mut unified_nodes = Vec::new();

    for (rel_path, root) in &elements {
        let file_annotations = ann_map.get(*rel_path).cloned().unwrap_or_default();
        // Flatten all annotations (including children) for binding lookup
        let flat_anns = flatten_annotations_recursive(&file_annotations);
        build_unified_nodes(root, &flat_anns, &mut unified_nodes);
    }

    // Step 5: Match (using indexed parent lookup for combinator support)
    let matchable_refs: Vec<&dyn Matchable> =
        unified_nodes.iter().map(|n| n as &dyn Matchable).collect();
    let parent_indices: Vec<Option<usize>> = unified_nodes.iter().map(|n| n.parent_idx).collect();
    let matched_indices =
        matcher::filter_by_selector_indexed(&matchable_refs, &parent_indices, &selector);

    // Step 6: Build results, deduplicating by source location
    let mut seen: FxHashMap<String, usize> = FxHashMap::default();
    let mut results: Vec<QueryResult> = Vec::new();

    for idx in matched_indices {
        let node = &unified_nodes[idx];

        let dedup_key = format!(
            "{}:{}:{}",
            node.source.file, node.source.line, node.source.column
        );

        if let Some(&idx) = seen.get(&dedup_key) {
            // Same code element, different annotation — merge
            if let Some(ann) = node.annotation {
                results[idx].annotations.push(AnnotationSummary {
                    tag: ann.tag.clone(),
                    attrs: ann.attrs.clone(),
                    binding: ann.binding.clone(),
                });
            }
        } else {
            let name = node
                .merged_attrs
                .get("name")
                .and_then(|v| v.as_str())
                .unwrap_or("");

            // Collect code-only attrs (exclude ann_tag which is synthetic)
            let mut code_attrs = node.merged_attrs.clone();
            code_attrs.remove(ANN_TAG_ATTR);

            let annotations = match node.annotation {
                Some(ann) => vec![AnnotationSummary {
                    tag: ann.tag.clone(),
                    attrs: ann.attrs.clone(),
                    binding: ann.binding.clone(),
                }],
                None => vec![],
            };

            let idx = results.len();
            seen.insert(dedup_key, idx);
            results.push(QueryResult {
                code_element: CodeElementSummary {
                    tag: node.code_tag.clone(),
                    name: CodeElementName::from(name),
                    file: node.source.file.clone(),
                    source: node.source.clone(),
                    attrs: code_attrs,
                },
                annotations,
            });
        }
    }

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

    Ok(results)
}

/// Run multiple selectors in parallel against the same scope.
///
/// Loads and parses files once (shared), then matches each selector via rayon.
/// Returns results grouped by selector (one `Vec<QueryResult>` per input selector).
#[must_use = "batch query results should be inspected"]
pub fn batch_query(
    selectors: &[&str],
    scope: &Scope,
    code_cache: &mut CodeCache,
    store: &mut AnnotationStore,
    resolvers: &ResolverRegistry,
    opts: Option<&crate::query_options::QueryOptions>,
) -> Result<Vec<Vec<QueryResult>>, AqlError> {
    if selectors.is_empty() {
        return Ok(vec![]);
    }

    // Shared setup: parse and cache scope once
    let code_tags = resolvers.all_code_tags();
    code_cache.ensure_scope(scope, resolvers);
    let _ = store.refresh_scope(scope);

    // Build unified nodes once
    let elements = code_cache.elements_in_scope(scope);
    let ann_by_file = store.annotations_in_scope(scope);
    let ann_map: FxHashMap<&RelativePath, Vec<&crate::store::Annotation>> =
        ann_by_file.into_iter().collect();

    let mut unified_nodes = Vec::new();

    for (rel_path, root) in &elements {
        let file_annotations = ann_map.get(*rel_path).cloned().unwrap_or_default();
        let flat_anns = flatten_annotations_recursive(&file_annotations);
        build_unified_nodes(root, &flat_anns, &mut unified_nodes);
    }

    let matchable_refs: Vec<&dyn Matchable> =
        unified_nodes.iter().map(|n| n as &dyn Matchable).collect();
    let parent_indices: Vec<Option<usize>> = unified_nodes.iter().map(|n| n.parent_idx).collect();

    // Parse and rewrite all selectors upfront
    let parsed: Vec<_> = selectors
        .iter()
        .map(|s| {
            let rewritten = rewrite::rewrite_selector(s, &code_tags);
            parse_selector(&rewritten)
        })
        .collect::<Result<Vec<_>, _>>()?;

    // Match each selector against the shared unified nodes
    let all_matched: Vec<Vec<usize>> = parsed
        .iter()
        .map(|sel| matcher::filter_by_selector_indexed(&matchable_refs, &parent_indices, sel))
        .collect();

    // Build results for each selector, then apply opts to each set
    let results: Vec<Vec<QueryResult>> = all_matched
        .into_iter()
        .map(|matched_indices| build_results(&unified_nodes, matched_indices))
        .collect();

    let results = match opts {
        Some(o) => results
            .into_iter()
            .map(|set| crate::query_options::apply_to_query_results(set, o))
            .collect(),
        None => results,
    };

    Ok(results)
}

/// Build deduplicated QueryResult list from matched node indices.
fn build_results(
    unified_nodes: &[node::UnifiedNode<'_>],
    matched_indices: Vec<usize>,
) -> Vec<QueryResult> {
    let mut seen: FxHashMap<String, usize> = FxHashMap::default();
    let mut results: Vec<QueryResult> = Vec::new();

    for idx in matched_indices {
        let node = &unified_nodes[idx];
        let dedup_key = format!(
            "{}:{}:{}",
            node.source.file, node.source.line, node.source.column
        );

        if let Some(&result_idx) = seen.get(&dedup_key) {
            if let Some(ann) = node.annotation {
                results[result_idx].annotations.push(AnnotationSummary {
                    tag: ann.tag.clone(),
                    attrs: ann.attrs.clone(),
                    binding: ann.binding.clone(),
                });
            }
        } else {
            let name = node
                .merged_attrs
                .get("name")
                .and_then(|v| v.as_str())
                .unwrap_or("");

            let mut code_attrs = node.merged_attrs.clone();
            code_attrs.remove(ANN_TAG_ATTR);

            let annotations = match node.annotation {
                Some(ann) => vec![AnnotationSummary {
                    tag: ann.tag.clone(),
                    attrs: ann.attrs.clone(),
                    binding: ann.binding.clone(),
                }],
                None => vec![],
            };

            let result_idx = results.len();
            seen.insert(dedup_key, result_idx);
            results.push(QueryResult {
                code_element: CodeElementSummary {
                    tag: node.code_tag.clone(),
                    name: CodeElementName::from(name),
                    file: node.source.file.clone(),
                    source: node.source.clone(),
                    attrs: code_attrs,
                },
                annotations,
            });
        }
    }

    results
}