amql_engine/

query_options.rs

//! Result-set options: sort, limit, and offset for query and select operations.
//!
//! Sort field syntax: `"field"` (ascending) or `"-field"` (descending).
//! Valid fields for `unified_query` results: `line`, `name`, `file`, `attr:<name>`.
//! Valid fields for `store::select` results: `tag`, `file`, `binding`, `attr:<name>`.
//!
//! Agentic or LLM-based comparison belongs in `QueryMiddleware` — the consumer
//! injects a middleware that re-orders results after the engine produces them.

#[cfg(all(feature = "resolver", feature = "fs"))]
use crate::query::QueryResult;
use crate::store::Annotation;
use serde::{Deserialize, Serialize};

/// Options applied to query and select result sets after matching.
#[non_exhaustive]
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
pub struct QueryOptions {
    /// Return at most this many results.
    pub limit: Option<usize>,
    /// Skip this many results before applying `limit`.
    pub offset: Option<usize>,
    /// Sort field. Prefix with `-` for descending.
    /// For `query`: `line`, `name`, `file`, `attr:<name>`.
    /// For `select`: `tag`, `file`, `binding`, `attr:<name>`.
    pub sort_by: Option<String>,
}

impl QueryOptions {
    /// Create options with explicit limit, offset, and sort_by.
    pub fn new(limit: Option<usize>, offset: Option<usize>, sort_by: Option<String>) -> Self {
        Self {
            limit,
            offset,
            sort_by,
        }
    }
}

/// Apply `QueryOptions` to query results (sort → skip → take).
#[cfg(all(feature = "resolver", feature = "fs"))]
pub fn apply_to_query_results(
    mut results: Vec<QueryResult>,
    opts: &QueryOptions,
) -> Vec<QueryResult> {
    if let Some(ref field) = opts.sort_by {
        sort_query_results(&mut results, field);
    }
    let offset = opts.offset.unwrap_or(0);
    results
        .into_iter()
        .skip(offset)
        .take(opts.limit.unwrap_or(usize::MAX))
        .collect()
}

/// Apply `QueryOptions` to annotations (sort → skip → take).
pub fn apply_to_annotations(mut results: Vec<Annotation>, opts: &QueryOptions) -> Vec<Annotation> {
    if let Some(ref field) = opts.sort_by {
        sort_annotations(&mut results, field);
    }
    let offset = opts.offset.unwrap_or(0);
    results
        .into_iter()
        .skip(offset)
        .take(opts.limit.unwrap_or(usize::MAX))
        .collect()
}

#[cfg(all(feature = "resolver", feature = "fs"))]
fn sort_query_results(results: &mut [QueryResult], field: &str) {
    let (field, desc) = parse_field(field);
    results.sort_by(|a, b| {
        let ord = match field {
            "line" => a.code_element.source.line.cmp(&b.code_element.source.line),
            "name" => {
                let an: &str = &a.code_element.name;
                let bn: &str = &b.code_element.name;
                an.cmp(bn)
            }
            "file" => {
                let af: &str = &a.code_element.file;
                let bf: &str = &b.code_element.file;
                af.cmp(bf)
            }
            f if f.starts_with("attr:") => {
                let attr = &f["attr:".len()..];
                let attr_key = crate::types::AttrName::from(attr);
                let av = a
                    .code_element
                    .attrs
                    .get(&attr_key)
                    .map(attr_sort_key)
                    .unwrap_or_default();
                let bv = b
                    .code_element
                    .attrs
                    .get(&attr_key)
                    .map(attr_sort_key)
                    .unwrap_or_default();
                av.cmp(&bv)
            }
            _ => std::cmp::Ordering::Equal,
        };
        if desc {
            ord.reverse()
        } else {
            ord
        }
    });
}

fn sort_annotations(results: &mut [Annotation], field: &str) {
    let (field, desc) = parse_field(field);
    results.sort_by(|a, b| {
        let ord = match field {
            "tag" => {
                let at: &str = &a.tag;
                let bt: &str = &b.tag;
                at.cmp(bt)
            }
            "file" => {
                let af: &str = &a.file;
                let bf: &str = &b.file;
                af.cmp(bf)
            }
            "binding" => {
                let ab: &str = &a.binding;
                let bb: &str = &b.binding;
                ab.cmp(bb)
            }
            f if f.starts_with("attr:") => {
                let attr = &f["attr:".len()..];
                let attr_key = crate::types::AttrName::from(attr);
                let av = a
                    .attrs
                    .get(&attr_key)
                    .map(attr_sort_key)
                    .unwrap_or_default();
                let bv = b
                    .attrs
                    .get(&attr_key)
                    .map(attr_sort_key)
                    .unwrap_or_default();
                av.cmp(&bv)
            }
            _ => std::cmp::Ordering::Equal,
        };
        if desc {
            ord.reverse()
        } else {
            ord
        }
    });
}

/// Parse a sort field string: strip leading `-` for descending. Returns `(field, is_desc)`.
fn parse_field(s: &str) -> (&str, bool) {
    if let Some(f) = s.strip_prefix('-') {
        (f, true)
    } else {
        (s, false)
    }
}

/// Derive a stable string sort key from a JSON attribute value.
fn attr_sort_key(v: &serde_json::Value) -> String {
    match v {
        serde_json::Value::String(s) => s.clone(),
        serde_json::Value::Number(n) => {
            // Zero-pad numbers for correct lexicographic ordering (up to 20 digits).
            if let Some(i) = n.as_u64() {
                format!("{i:020}")
            } else if let Some(i) = n.as_i64() {
                // Negative: shift by i64::MIN to make unsigned-comparable.
                format!("{:020}", i.wrapping_sub(i64::MIN))
            } else {
                n.to_string()
            }
        }
        serde_json::Value::Bool(b) => b.to_string(),
        serde_json::Value::Null => String::new(),
        other => other.to_string(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::store::Annotation;
    use crate::types::{Binding, RelativePath, TagName};
    use rustc_hash::FxHashMap;

    fn make_ann(tag: &str, file: &str) -> Annotation {
        Annotation {
            tag: TagName::from(tag),
            attrs: FxHashMap::default(),
            binding: Binding::from(""),
            file: RelativePath::from(file),
            children: vec![],
        }
    }

    #[test]
    fn sorts_annotations_by_tag() {
        // Arrange
        let results = vec![
            make_ann("route", "b.ts"),
            make_ann("controller", "a.ts"),
            make_ann("middleware", "c.ts"),
        ];
        let opts = QueryOptions {
            sort_by: Some("tag".into()),
            ..Default::default()
        };

        // Act
        let results = apply_to_annotations(results, &opts);

        // Assert
        let tags: Vec<&str> = results.iter().map(|a| a.tag.as_ref()).collect();
        assert_eq!(
            tags,
            ["controller", "middleware", "route"],
            "should be sorted by tag asc"
        );
    }

    #[test]
    fn applies_limit_and_offset() {
        // Arrange
        let results: Vec<Annotation> = (0..5).map(|i| make_ann(&format!("t{i}"), "f.ts")).collect();
        let opts = QueryOptions {
            offset: Some(1),
            limit: Some(2),
            ..Default::default()
        };

        // Act
        let results = apply_to_annotations(results, &opts);

        // Assert
        assert_eq!(results.len(), 2, "exactly 2 results after offset+limit");
        let tag: &str = &results[0].tag;
        assert_eq!(
            tag, "t1",
            "first result is the second element (after offset 1)"
        );
    }

    #[test]
    fn sorts_descending_with_minus_prefix() {
        // Arrange
        let results = vec![
            make_ann("a", "f.ts"),
            make_ann("c", "f.ts"),
            make_ann("b", "f.ts"),
        ];
        let opts = QueryOptions {
            sort_by: Some("-tag".into()),
            ..Default::default()
        };

        // Act
        let results = apply_to_annotations(results, &opts);

        // Assert
        let tags: Vec<&str> = results.iter().map(|a| a.tag.as_ref()).collect();
        assert_eq!(tags, ["c", "b", "a"], "should be sorted by tag desc");
    }
}