amql_engine/

pattern.rs

//! Pattern extraction and delta diffing for annotation instances.
//!
//! Given all annotations of a tag, identifies the common structure (template)
//! and what is unique per instance (delta). Agents call once to get the pattern
//! plus per-instance differences, avoiding redundant source transmission.

use crate::store::{Annotation, AnnotationStore};
use crate::types::{AttrName, Binding, RelativePath, TagName};
use rustc_hash::{FxHashMap, FxHashSet};
use serde::Serialize;
use serde_json::Value as JsonValue;

/// Common structure across all annotations of a tag.
#[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 Pattern {
    /// The tag these annotations share.
    pub tag: TagName,
    /// Summary of each attribute across all instances.
    #[cfg_attr(
        feature = "ts",
        ts(as = "std::collections::HashMap<AttrName, AttrSummary>")
    )]
    #[cfg_attr(
        feature = "flow",
        flow(as = "std::collections::HashMap<AttrName, AttrSummary>")
    )]
    pub common_attrs: FxHashMap<AttrName, AttrSummary>,
    /// Representative source snippet (when code elements are available).
    pub template_source: Option<String>,
    /// Number of annotation instances analyzed.
    pub instance_count: usize,
}

/// Statistical summary of an attribute across annotation instances.
#[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 AttrSummary {
    /// Whether this attribute is present in every instance.
    pub always_present: bool,
    /// Number of distinct values observed.
    pub distinct_values: usize,
    /// Up to 3 sample values for context.
    pub sample_values: Vec<JsonValue>,
}

/// What is unique about a specific annotation relative to the pattern.
#[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 Delta {
    /// Binding key of this specific instance.
    pub binding: Binding,
    /// Attributes whose values differ from the most common value, or are unique to this instance.
    #[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 unique_attrs: FxHashMap<AttrName, JsonValue>,
    /// File containing this annotation.
    pub file: RelativePath,
}

/// Result of pattern extraction: the common pattern and per-instance deltas.
#[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 PatternResult {
    /// Common structure, if any annotations of the tag exist.
    pub pattern: Option<Pattern>,
    /// Per-instance differences from the pattern.
    pub deltas: Vec<Delta>,
}

const MAX_SAMPLE_VALUES: usize = 3;

/// Analyze all annotations of a tag and return the common pattern.
#[must_use = "pattern contains structural analysis of the tag"]
pub fn extract_pattern(store: &AnnotationStore, tag: &TagName) -> Option<Pattern> {
    let annotations = collect_by_tag(store, tag);
    if annotations.is_empty() {
        return None;
    }

    let common_attrs = build_attr_summaries(&annotations);

    Some(Pattern {
        tag: tag.clone(),
        common_attrs,
        template_source: None,
        instance_count: annotations.len(),
    })
}

/// Given a pattern and a specific annotation, return only what is unique.
pub fn diff_from_pattern(pattern: &Pattern, annotation: &Annotation) -> Delta {
    let mut unique_attrs = FxHashMap::default();

    for (attr, value) in &annotation.attrs {
        match pattern.common_attrs.get(attr) {
            Some(summary) => {
                // Attribute is unique if it has high cardinality (many distinct values)
                // or is not always present
                if summary.distinct_values > 1 {
                    unique_attrs.insert(attr.clone(), value.clone());
                }
            }
            None => {
                // Attribute not in pattern at all — definitely unique
                unique_attrs.insert(attr.clone(), value.clone());
            }
        }
    }

    Delta {
        binding: annotation.binding.clone(),
        unique_attrs,
        file: annotation.file.clone(),
    }
}

/// Extract the pattern and compute deltas for every instance of a tag.
#[must_use = "pattern and deltas contain per-instance analysis"]
pub fn pattern_with_deltas(
    store: &AnnotationStore,
    tag: &TagName,
) -> Option<(Pattern, Vec<Delta>)> {
    let pattern = extract_pattern(store, tag)?;
    let annotations = collect_by_tag(store, tag);
    let deltas = annotations
        .iter()
        .map(|ann| diff_from_pattern(&pattern, ann))
        .collect();
    Some((pattern, deltas))
}

/// Collect all annotations matching a tag (flat, including children).
fn collect_by_tag<'a>(store: &'a AnnotationStore, tag: &TagName) -> Vec<&'a Annotation> {
    let all = store.get_all_annotations();
    all.iter()
        .flat_map(|ann| collect_recursive(ann, tag))
        .collect()
}

fn collect_recursive<'a>(ann: &'a Annotation, tag: &TagName) -> Vec<&'a Annotation> {
    let mut out = Vec::new();
    if ann.tag == *tag {
        out.push(ann);
    }
    for child in &ann.children {
        out.extend(collect_recursive(child, tag));
    }
    out
}

/// Build per-attribute summaries across all annotations.
fn build_attr_summaries(annotations: &[&Annotation]) -> FxHashMap<AttrName, AttrSummary> {
    let total = annotations.len();

    // Collect all observed attribute names
    let mut attr_names: FxHashSet<AttrName> = FxHashSet::default();
    for ann in annotations {
        for key in ann.attrs.keys() {
            attr_names.insert(key.clone());
        }
    }

    let mut summaries = FxHashMap::default();

    for attr in attr_names {
        let mut present_count = 0usize;
        let mut distinct: FxHashSet<String> = FxHashSet::default();
        let mut samples: Vec<JsonValue> = Vec::new();

        for ann in annotations {
            if let Some(value) = ann.attrs.get(&attr) {
                present_count += 1;
                let serialized = value.to_string();
                if distinct.insert(serialized) && samples.len() < MAX_SAMPLE_VALUES {
                    samples.push(value.clone());
                }
            }
        }

        summaries.insert(
            attr,
            AttrSummary {
                always_present: present_count == total,
                distinct_values: distinct.len(),
                sample_values: samples,
            },
        );
    }

    summaries
}

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

    fn make_annotation(
        tag: &str,
        binding: &str,
        file: &str,
        attrs: Vec<(&str, &str)>,
    ) -> Annotation {
        let mut attr_map = FxHashMap::default();
        for (k, v) in attrs {
            attr_map.insert(AttrName::from(k), JsonValue::String(v.to_string()));
        }
        Annotation {
            tag: TagName::from(tag),
            attrs: attr_map,
            binding: Binding::from(binding),
            file: RelativePath::from(file),
            children: vec![],
        }
    }

    fn test_store() -> AnnotationStore {
        let mut store = AnnotationStore::new(Path::new("/project"));
        store.inject_test_data(
            &RelativePath::from("src/routes/users.ts"),
            vec![make_annotation(
                "route",
                "getUsers",
                "src/routes/users.ts",
                vec![("method", "GET"), ("path", "/users")],
            )],
        );
        store.inject_test_data(
            &RelativePath::from("src/routes/posts.ts"),
            vec![make_annotation(
                "route",
                "getPosts",
                "src/routes/posts.ts",
                vec![("method", "GET"), ("path", "/posts")],
            )],
        );
        store.inject_test_data(
            &RelativePath::from("src/routes/admin.ts"),
            vec![make_annotation(
                "route",
                "createUser",
                "src/routes/admin.ts",
                vec![("method", "POST"), ("path", "/admin/users")],
            )],
        );
        store
    }

    #[test]
    fn extract_pattern_finds_common_attrs() {
        // Arrange
        let store = test_store();
        let tag = TagName::from("route");

        // Act
        let pattern = extract_pattern(&store, &tag).unwrap();

        // Assert
        assert_eq!(pattern.instance_count, 3, "should find 3 route annotations");
        assert_eq!(pattern.tag, "route", "tag should match");

        let method_summary = pattern.common_attrs.get(&AttrName::from("method")).unwrap();
        assert!(
            method_summary.always_present,
            "method should be always present"
        );
        assert_eq!(
            method_summary.distinct_values, 2,
            "GET and POST are distinct"
        );

        let path_summary = pattern.common_attrs.get(&AttrName::from("path")).unwrap();
        assert!(path_summary.always_present, "path should be always present");
        assert_eq!(path_summary.distinct_values, 3, "all paths are distinct");
    }

    #[test]
    fn diff_from_pattern_identifies_unique_attrs() {
        // Arrange
        let store = test_store();
        let tag = TagName::from("route");
        let pattern = extract_pattern(&store, &tag).unwrap();
        let ann = make_annotation(
            "route",
            "getUsers",
            "src/routes/users.ts",
            vec![("method", "GET"), ("path", "/users")],
        );

        // Act
        let delta = diff_from_pattern(&pattern, &ann);

        // Assert
        assert_eq!(delta.binding, "getUsers", "binding should match");
        assert!(
            delta.unique_attrs.contains_key(&AttrName::from("method")),
            "method has >1 distinct value, so it is unique per instance"
        );
        assert!(
            delta.unique_attrs.contains_key(&AttrName::from("path")),
            "path has >1 distinct value, so it is unique per instance"
        );
    }

    #[test]
    fn pattern_with_deltas_returns_all() {
        // Arrange
        let store = test_store();
        let tag = TagName::from("route");

        // Act
        let (pattern, deltas) = pattern_with_deltas(&store, &tag).unwrap();

        // Assert
        assert_eq!(pattern.instance_count, 3, "should have 3 instances");
        assert_eq!(deltas.len(), 3, "should have 3 deltas");
    }

    #[test]
    fn missing_tag_returns_none() {
        // Arrange
        let store = test_store();
        let tag = TagName::from("nonexistent");

        // Act
        let result = extract_pattern(&store, &tag);

        // Assert
        assert!(result.is_none(), "should return None for unknown tag");
    }

    #[test]
    fn constant_attr_excluded_from_delta() {
        // Arrange — all annotations share the same method
        let mut store = AnnotationStore::new(Path::new("/project"));
        store.inject_test_data(
            &RelativePath::from("src/a.ts"),
            vec![make_annotation(
                "hook",
                "useA",
                "src/a.ts",
                vec![("type", "effect")],
            )],
        );
        store.inject_test_data(
            &RelativePath::from("src/b.ts"),
            vec![make_annotation(
                "hook",
                "useB",
                "src/b.ts",
                vec![("type", "effect")],
            )],
        );
        let tag = TagName::from("hook");
        let pattern = extract_pattern(&store, &tag).unwrap();
        let ann = make_annotation("hook", "useA", "src/a.ts", vec![("type", "effect")]);

        // Act
        let delta = diff_from_pattern(&pattern, &ann);

        // Assert
        assert!(
            delta.unique_attrs.is_empty(),
            "constant attr should not appear in delta"
        );
    }
}