xlog_logic/

diagnostics.rs

//! Stable diagnostic records for rule provenance and query proof traces.

use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};

use xlog_core::{symbol, ScalarType};

use crate::ast::{AggOp, ArithExpr, Atom, BodyLiteral, CompOp, Program, Rule, Term};
use crate::expand::generated_function_variable_source;

/// Origin class for a rule visible through diagnostic introspection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RuleSourceKind {
    /// Authored in the parsed XLOG source.
    Source,
    /// Generated by a compiler or rewrite pass.
    Generated,
    /// Mined or induced by an external rule search path.
    Mined,
    /// Imported from a module boundary.
    Imported,
    /// Injected into a live runtime session.
    RuntimeInjected,
}

impl RuleSourceKind {
    /// Stable lowercase string used by CLI and pyxlog JSON/dict outputs.
    pub fn as_str(self) -> &'static str {
        match self {
            RuleSourceKind::Source => "source",
            RuleSourceKind::Generated => "generated",
            RuleSourceKind::Mined => "mined",
            RuleSourceKind::Imported => "imported",
            RuleSourceKind::RuntimeInjected => "runtime_injected",
        }
    }
}

/// Public rule provenance record.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuleProvenance {
    /// Stable rule identifier within the diagnostic report.
    pub rule_id: String,
    /// Formatted rule head, e.g. `reach(X, Y)`.
    pub head: String,
    /// Rule origin class.
    pub source_kind: RuleSourceKind,
    /// Source span when available from the parser.
    pub source_span: Option<String>,
    /// Stable hash of the source or generation trace.
    pub generation_trace_hash: Option<String>,
    /// Body relation ids that support this rule.
    pub support_relation_ids: Vec<String>,
    /// Counterexample relation ids when supplied by an induction path.
    pub counterexample_relation_ids: Vec<String>,
}

/// Public query proof trace record.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QueryProofTrace {
    /// Stable query identifier.
    pub query_id: String,
    /// Formatted source query atom.
    pub query: String,
    /// Internal answer relation for the query.
    pub answer_relation: String,
    /// Rule ids that can derive the query predicate.
    pub rule_ids: Vec<String>,
    /// Source facts named by the direct proof frontier.
    pub source_facts: Vec<String>,
    /// Rejected or negated alternatives encountered in the trace.
    pub rejected_alternatives: Vec<String>,
}

/// Build provenance for source rules and optional generated rewrite rules.
pub fn rule_provenance(
    program: &Program,
    generated_program: Option<&Program>,
) -> Vec<RuleProvenance> {
    let mut out = Vec::new();
    let mut source_keys = BTreeSet::new();

    for (idx, rule) in program.rules.iter().enumerate() {
        source_keys.insert(rule_key(rule));
        out.push(rule_record(idx, rule, RuleSourceKind::Source));
    }

    if let Some(generated) = generated_program {
        let mut generated_idx = 0usize;
        for rule in &generated.rules {
            if source_keys.contains(&rule_key(rule)) {
                continue;
            }
            out.push(rule_record(generated_idx, rule, RuleSourceKind::Generated));
            generated_idx += 1;
        }
    }

    out
}

/// Build source and generated rule-provenance records for a program.
pub fn build_rule_provenance(
    program: &Program,
    generated_predicates: &[String],
) -> Vec<RuleProvenance> {
    let mut out = rule_provenance(program, None);
    for (idx, predicate) in generated_predicates.iter().enumerate() {
        out.push(RuleProvenance {
            rule_id: format!("rule:generated:{}:{}", idx, predicate),
            head: predicate.clone(),
            source_kind: RuleSourceKind::Generated,
            source_span: None,
            generation_trace_hash: Some(stable_hash(&format!("generated:{}", predicate))),
            support_relation_ids: vec![predicate.clone()],
            counterexample_relation_ids: Vec::new(),
        });
    }
    out
}

/// Build direct proof traces for source queries.
pub fn query_proof_traces(
    program: &Program,
    provenance: &[RuleProvenance],
) -> Vec<QueryProofTrace> {
    let mut rule_ids_by_head: BTreeMap<String, Vec<String>> = BTreeMap::new();
    for entry in provenance {
        rule_ids_by_head
            .entry(head_predicate(&entry.head).to_string())
            .or_default()
            .push(entry.rule_id.clone());
    }

    program
        .queries
        .iter()
        .enumerate()
        .map(|(idx, query)| {
            let query_pred = query.atom.predicate.clone();
            let deriving_rules: Vec<&Rule> = program
                .rules
                .iter()
                .filter(|rule| !rule.is_fact() && rule.head.predicate == query_pred)
                .collect();
            let rule_ids = rule_ids_by_head
                .get(&query_pred)
                .cloned()
                .unwrap_or_default();
            let source_facts = source_facts_for_rules(program, &deriving_rules);
            let rejected_alternatives = deriving_rules
                .iter()
                .flat_map(|rule| {
                    rule.body.iter().filter_map(|lit| match lit {
                        BodyLiteral::Negated(atom) => Some(format!("not {}", format_atom(atom))),
                        _ => None,
                    })
                })
                .collect::<Vec<_>>();

            QueryProofTrace {
                query_id: format!("query:source:{}:{}", idx, format_atom(&query.atom)),
                query: format_atom(&query.atom),
                answer_relation: format!("__xlog_query_{}", idx),
                rule_ids,
                source_facts,
                rejected_alternatives,
            }
        })
        .collect()
}

/// Build coarse proof traces that name source rules and source facts per query.
pub fn build_query_proof_traces(program: &Program) -> Vec<QueryProofTrace> {
    let provenance = rule_provenance(program, None);
    query_proof_traces(program, &provenance)
}

/// Build source-facing provenance and proof traces from authored and normalized programs.
///
/// Normalized support relations and rejected alternatives are merged only while rule and
/// query order still align with the authored program. Compiler-generated predicates and
/// variables are removed or mapped back to their authored function-local names.
pub fn source_diagnostics(
    source_program: &Program,
    analysis_program: &Program,
    rewritten_program: Option<&Program>,
) -> (Vec<RuleProvenance>, Vec<QueryProofTrace>) {
    let mut provenance = rule_provenance(source_program, None);
    let rules_align = analysis_program.rules.len() >= source_program.rules.len()
        && source_program
            .rules
            .iter()
            .zip(&analysis_program.rules)
            .all(|(source, normalized)| {
                source.head.predicate == normalized.head.predicate
                    && source.head.arity() == normalized.head.arity()
            });
    if rules_align {
        let normalized = rule_provenance(analysis_program, None);
        for (source, normalized) in provenance.iter_mut().zip(normalized) {
            source.support_relation_ids = source
                .support_relation_ids
                .iter()
                .chain(&normalized.support_relation_ids)
                .filter(|name| !name.starts_with("__"))
                .cloned()
                .collect::<BTreeSet<_>>()
                .into_iter()
                .collect();
        }
    }

    if let Some(rewritten_program) = rewritten_program {
        provenance.extend(
            rule_provenance(analysis_program, Some(rewritten_program))
                .into_iter()
                .filter(|entry| entry.source_kind == RuleSourceKind::Generated),
        );
    }

    let mut proof_traces = query_proof_traces(source_program, &provenance);
    let queries_align = source_program.queries.len() == analysis_program.queries.len()
        && source_program
            .queries
            .iter()
            .zip(&analysis_program.queries)
            .all(|(source, normalized)| {
                source.atom.predicate == normalized.atom.predicate
                    && source.atom.arity() == normalized.atom.arity()
            });
    if rules_align && queries_align {
        let normalized_provenance = rule_provenance(analysis_program, None);
        let normalized_traces = query_proof_traces(analysis_program, &normalized_provenance);
        let function_variable_sources =
            generated_function_variable_sources(source_program, analysis_program);
        let authored_facts = source_program
            .facts()
            .map(|rule| format!("{}.", format_atom(&rule.head)))
            .collect::<BTreeSet<_>>();
        for (source, normalized) in proof_traces.iter_mut().zip(normalized_traces) {
            source.source_facts = source
                .source_facts
                .iter()
                .chain(&normalized.source_facts)
                .filter(|fact| authored_facts.contains(*fact) && !fact.starts_with("__"))
                .cloned()
                .collect::<BTreeSet<_>>()
                .into_iter()
                .collect();
            source.rejected_alternatives = source
                .rejected_alternatives
                .iter()
                .cloned()
                .chain(normalized.rejected_alternatives.iter().map(|alternative| {
                    source_format_normalized_alternative(alternative, &function_variable_sources)
                }))
                .filter(|alternative| {
                    !alternative
                        .strip_prefix("not ")
                        .unwrap_or(alternative)
                        .starts_with("__")
                })
                .collect::<BTreeSet<_>>()
                .into_iter()
                .collect();
        }
    }
    (provenance, proof_traces)
}

/// Map generated predicate-function variables to their authored local names.
pub fn generated_function_variable_sources(
    source_program: &Program,
    analysis_program: &Program,
) -> HashMap<String, String> {
    let authored_variables = program_variable_names(source_program);
    let mut function_locals = source_program
        .functions
        .iter()
        .filter_map(|function| {
            let crate::ast::FuncBody::Predicate { result, body } = &function.body else {
                return None;
            };
            let parameters = function
                .params
                .iter()
                .map(|parameter| parameter.name.as_str())
                .collect::<HashSet<_>>();
            let mut locals = body
                .iter()
                .flat_map(BodyLiteral::variables)
                .filter(|name| !parameters.contains(name))
                .map(ToOwned::to_owned)
                .collect::<HashSet<_>>();
            if !parameters.contains(result.as_str()) {
                locals.insert(result.clone());
            }
            Some((function.name.clone(), locals))
        })
        .collect::<Vec<_>>();
    function_locals.sort_by_key(|(function, _)| std::cmp::Reverse(function.len()));

    program_variable_names(analysis_program)
        .into_iter()
        .filter(|name| !authored_variables.contains(name))
        .filter_map(|name| {
            let source_name = function_locals.iter().find_map(|(function, locals)| {
                let source_name = generated_function_variable_source(&name, function)?;
                locals
                    .contains(source_name)
                    .then(|| source_name.to_string())
            })?;
            Some((name, source_name))
        })
        .collect()
}

fn program_variable_names(program: &Program) -> HashSet<String> {
    let mut variables = HashSet::new();
    for rule in &program.rules {
        variables.extend(rule.head.variables().into_iter().map(ToOwned::to_owned));
        variables.extend(
            rule.body
                .iter()
                .flat_map(BodyLiteral::variables)
                .map(ToOwned::to_owned),
        );
    }
    for constraint in &program.constraints {
        variables.extend(
            constraint
                .body
                .iter()
                .flat_map(BodyLiteral::variables)
                .map(ToOwned::to_owned),
        );
    }
    for query in &program.queries {
        variables.extend(query.atom.variables().into_iter().map(ToOwned::to_owned));
    }
    variables
}

/// Replace generated predicate-function variable tokens with authored local names.
pub fn source_format_normalized_alternative(
    alternative: &str,
    function_variable_sources: &HashMap<String, String>,
) -> String {
    let characters = alternative.chars().collect::<Vec<_>>();
    let mut output = String::with_capacity(alternative.len());
    let mut index = 0;
    let mut quoted = false;
    let mut escaped = false;
    while index < characters.len() {
        let character = characters[index];
        if quoted {
            output.push(character);
            if escaped {
                escaped = false;
            } else if character == '\\' {
                escaped = true;
            } else if character == '"' {
                quoted = false;
            }
            index += 1;
            continue;
        }
        if character == '"' {
            quoted = true;
            output.push(character);
            index += 1;
            continue;
        }
        if character.is_ascii_alphanumeric() || character == '_' {
            let start = index;
            while index < characters.len()
                && (characters[index].is_ascii_alphanumeric() || characters[index] == '_')
            {
                index += 1;
            }
            let token = characters[start..index].iter().collect::<String>();
            output.push_str(
                function_variable_sources
                    .get(&token)
                    .map(String::as_str)
                    .unwrap_or(&token),
            );
            continue;
        }
        output.push(character);
        index += 1;
    }
    output
}

fn rule_record(idx: usize, rule: &Rule, source_kind: RuleSourceKind) -> RuleProvenance {
    let head = format_atom(&rule.head);
    let prefix = source_kind.as_str();
    RuleProvenance {
        rule_id: format!("rule:{}:{}:{}", prefix, idx, stable_hash(&rule_key(rule))),
        head,
        source_kind,
        source_span: Some(format!("rule_index:{}", idx)),
        generation_trace_hash: Some(stable_hash(&rule_key(rule))),
        support_relation_ids: support_relation_ids(rule),
        counterexample_relation_ids: Vec::new(),
    }
}

fn support_relation_ids(rule: &Rule) -> Vec<String> {
    rule.body_predicates()
        .into_iter()
        .map(str::to_string)
        .collect::<BTreeSet<_>>()
        .into_iter()
        .collect()
}

fn source_facts_for_rules(program: &Program, rules: &[&Rule]) -> Vec<String> {
    let wanted: BTreeSet<String> = rules
        .iter()
        .flat_map(|rule| {
            rule.body
                .iter()
                .filter_map(|lit| lit.atom().map(|atom| atom.predicate.clone()))
        })
        .collect();

    let mut facts = BTreeSet::new();
    for fact in program.facts() {
        if wanted.contains(&fact.head.predicate) {
            facts.insert(format!("{}.", format_atom(&fact.head)));
        }
    }
    facts.into_iter().collect()
}

fn rule_key(rule: &Rule) -> String {
    let mut out = format_atom(&rule.head);
    if !rule.body.is_empty() {
        let body = rule
            .body
            .iter()
            .map(format_body_literal)
            .collect::<Vec<_>>()
            .join(", ");
        out.push_str(" :- ");
        out.push_str(&body);
    }
    out
}

fn head_predicate(head: &str) -> &str {
    head.split_once('(').map(|(pred, _)| pred).unwrap_or(head)
}

/// Format an atom in source-like syntax.
pub fn format_atom(atom: &Atom) -> String {
    let args = atom
        .terms
        .iter()
        .map(format_term)
        .collect::<Vec<_>>()
        .join(", ");
    format!("{}({})", atom.predicate, args)
}

/// Format an integrity-constraint body in source-like syntax.
pub fn format_constraint_body(body: &[BodyLiteral]) -> String {
    let literals = body
        .iter()
        .map(format_body_literal)
        .collect::<Vec<_>>()
        .join(", ");
    format!(":- {}.", literals)
}

fn format_body_literal(lit: &BodyLiteral) -> String {
    match lit {
        BodyLiteral::Positive(atom) => format_atom(atom),
        BodyLiteral::Negated(atom) => format!("not {}", format_atom(atom)),
        BodyLiteral::Epistemic(lit) => format_epistemic_literal(lit),
        BodyLiteral::Comparison(comparison) => format!(
            "{} {} {}",
            format_term(&comparison.left),
            format_comp_op(comparison.op),
            format_term(&comparison.right)
        ),
        BodyLiteral::IsExpr(is_expr) => {
            format!("{} is {}", is_expr.target, format_arith_expr(&is_expr.expr))
        }
        BodyLiteral::Univ(univ) => {
            format!(
                "{} =.. {}",
                format_term(&univ.term),
                format_term(&univ.parts)
            )
        }
    }
}

fn format_epistemic_literal(lit: &crate::ast::EpistemicLiteral) -> String {
    let op = match lit.op {
        crate::ast::EpistemicOp::Know => "know",
        crate::ast::EpistemicOp::Possible => "possible",
    };
    if lit.negated {
        format!("not {} {}", op, format_atom(&lit.atom))
    } else {
        format!("{} {}", op, format_atom(&lit.atom))
    }
}

/// Format a term in source-like syntax.
pub fn format_term(term: &Term) -> String {
    match term {
        Term::Variable(name) => name.clone(),
        Term::Anonymous => "_".to_string(),
        Term::Integer(value) => value.to_string(),
        Term::Float(value) => value.to_string(),
        Term::String(value) => format!("\"{}\"", escape_string_contents(value)),
        Term::Symbol(id) => symbol::resolve(*id),
        Term::List(items) => {
            let values = items.iter().map(format_term).collect::<Vec<_>>().join(", ");
            format!("[{}]", values)
        }
        Term::Cons { head, tail } => {
            format!("[{} | {}]", format_term(head), format_term(tail))
        }
        Term::Compound { functor, args } => {
            let values = args.iter().map(format_term).collect::<Vec<_>>().join(", ");
            format!("{}({})", functor, values)
        }
        Term::PredRef(name) => name.clone(),
        Term::Aggregate(agg) => format!("{}({})", format_agg_op(agg.op), agg.variable),
    }
}

fn escape_string_contents(value: &str) -> String {
    let mut escaped = String::with_capacity(value.len());
    for character in value.chars() {
        match character {
            '\\' => escaped.push_str("\\\\"),
            '"' => escaped.push_str("\\\""),
            '\n' => escaped.push_str("\\n"),
            '\r' => escaped.push_str("\\r"),
            '\t' => escaped.push_str("\\t"),
            character if character.is_control() => escaped.extend(character.escape_default()),
            character => escaped.push(character),
        }
    }
    escaped
}

fn format_arith_expr(expr: &ArithExpr) -> String {
    match expr {
        ArithExpr::Variable(name) => name.clone(),
        ArithExpr::Integer(value) => value.to_string(),
        ArithExpr::Float(value) => value.to_string(),
        ArithExpr::Add(left, right) => {
            format!(
                "({} + {})",
                format_arith_expr(left),
                format_arith_expr(right)
            )
        }
        ArithExpr::Sub(left, right) => {
            format!(
                "({} - {})",
                format_arith_expr(left),
                format_arith_expr(right)
            )
        }
        ArithExpr::Mul(left, right) => {
            format!(
                "({} * {})",
                format_arith_expr(left),
                format_arith_expr(right)
            )
        }
        ArithExpr::Div(left, right) => {
            format!(
                "({} / {})",
                format_arith_expr(left),
                format_arith_expr(right)
            )
        }
        ArithExpr::Mod(left, right) => {
            format!(
                "({} % {})",
                format_arith_expr(left),
                format_arith_expr(right)
            )
        }
        ArithExpr::Abs(value) => format!("abs({})", format_arith_expr(value)),
        ArithExpr::Min(left, right) => {
            format!(
                "min({}, {})",
                format_arith_expr(left),
                format_arith_expr(right)
            )
        }
        ArithExpr::Max(left, right) => {
            format!(
                "max({}, {})",
                format_arith_expr(left),
                format_arith_expr(right)
            )
        }
        ArithExpr::Pow(left, right) => {
            format!(
                "pow({}, {})",
                format_arith_expr(left),
                format_arith_expr(right)
            )
        }
        ArithExpr::Cast(value, ty) => {
            format!(
                "cast({}, {})",
                format_arith_expr(value),
                format_scalar_type(*ty)
            )
        }
        ArithExpr::FuncCall { name, args } => {
            let values = args
                .iter()
                .map(format_arith_expr)
                .collect::<Vec<_>>()
                .join(", ");
            format!("{}({})", name, values)
        }
        ArithExpr::Conditional {
            cond_left,
            cond_op,
            cond_right,
            then_expr,
            else_expr,
        } => format!(
            "if {} {} {} then {} else {}",
            format_arith_expr(cond_left),
            format_comp_op(*cond_op),
            format_arith_expr(cond_right),
            format_arith_expr(then_expr),
            format_arith_expr(else_expr)
        ),
    }
}

pub(crate) fn format_scalar_type(typ: ScalarType) -> &'static str {
    match typ {
        ScalarType::U32 => "u32",
        ScalarType::U64 => "u64",
        ScalarType::I32 => "i32",
        ScalarType::I64 => "i64",
        ScalarType::F32 => "f32",
        ScalarType::F64 => "f64",
        ScalarType::Bool => "bool",
        ScalarType::Symbol => "symbol",
    }
}

fn format_comp_op(op: CompOp) -> &'static str {
    match op {
        CompOp::Eq => "==",
        CompOp::Ne => "!=",
        CompOp::Lt => "<",
        CompOp::Le => "<=",
        CompOp::Gt => ">",
        CompOp::Ge => ">=",
    }
}

fn format_agg_op(op: AggOp) -> &'static str {
    match op {
        AggOp::Count => "count",
        AggOp::Sum => "sum",
        AggOp::Min => "min",
        AggOp::Max => "max",
        AggOp::LogSumExp => "logsumexp",
    }
}

fn stable_hash(value: &str) -> String {
    let mut hash = 0xcbf29ce484222325u64;
    for byte in value.as_bytes() {
        hash ^= u64::from(*byte);
        hash = hash.wrapping_mul(0x100000001b3);
    }
    format!("{:016x}", hash)
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use super::{
        format_constraint_body, generated_function_variable_sources,
        source_format_normalized_alternative,
    };

    #[test]
    fn constraint_formatter_uses_source_syntax_for_function_expressions() {
        let program = crate::parse_program(
            "func get_parent(Child) = Parent :- parent(Child, Parent).\n\
             :- Parent is cast(get_parent(1), u64).",
        )
        .expect("parse constraint");

        assert_eq!(
            format_constraint_body(&program.constraints[0].body),
            ":- Parent is cast(get_parent(1), u64)."
        );
    }

    #[test]
    fn generated_variable_mapping_handles_overlapping_function_names_and_underscores() {
        let source = crate::parse_program(
            "pred first(i32, i32).\n\
             pred second(i32, i32).\n\
             pred answer(i32, i32).\n\
             first(1, 2).\n\
             second(1, 3).\n\
             func get(X) = Parent_Value :- first(X, Parent_Value).\n\
             func get_parent(X) = Value_With_Underscore :- second(X, Value_With_Underscore).\n\
             answer(A, B) :- A is get(1), B is get_parent(1).\n\
             ?- answer(A, B).",
        )
        .expect("parse predicate functions");
        let expanded =
            crate::expand_program_functions(&source, 1_000).expect("expand predicate functions");
        let mapping = generated_function_variable_sources(&source, &expanded);
        assert!(mapping.values().any(|name| name == "Parent_Value"));
        assert!(mapping.values().any(|name| name == "Value_With_Underscore"));
    }

    #[test]
    fn normalized_alternative_mapping_does_not_rewrite_quoted_tokens() {
        let generated = "__XLOG_FUNCTION_VISIBLE_Result_Value_7";
        let mapping = HashMap::from([(generated.to_string(), "Result_Value".to_string())]);
        assert_eq!(
            source_format_normalized_alternative(
                &format!("not blocked({generated}, \"{generated}\")"),
                &mapping,
            ),
            format!("not blocked(Result_Value, \"{generated}\")")
        );
    }

    #[test]
    fn normalized_alternative_mapping_preserves_escaped_string_boundaries() {
        let generated = "__XLOG_FUNCTION_VISIBLE_Result_Value_7";
        let mapping = HashMap::from([(generated.to_string(), "Result_Value".to_string())]);
        let mut literal = format!("quoted \" {generated} ");
        literal.push('\\');
        literal.push('\n');
        literal.push('\r');
        literal.push('\t');
        literal.push('\u{0007}');
        let atom = crate::ast::Atom {
            predicate: "blocked".to_string(),
            terms: vec![crate::ast::Term::String(literal)],
        };
        let formatted = super::format_atom(&atom);

        assert_eq!(
            source_format_normalized_alternative(&formatted, &mapping),
            formatted
        );
        assert!(formatted.contains("\\\\"));
        assert!(formatted.contains("\\n"));
        assert!(formatted.contains("\\r"));
        assert!(formatted.contains("\\t"));
        assert!(formatted.contains("\\u{7}"));
        assert!(!formatted.contains('\n'));
        assert!(!formatted.contains('\r'));
        assert!(!formatted.contains('\t'));
        assert!(formatted.contains(&format!("\\\" {generated} \\\\\\n")));
    }
}