lace_sans_io/logic/lacegram/

validation.rs

//! Rule safety, predicate-role, and variable-kind validation.

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

use crate::types::{
    BodyAtom, OrderingExpression, Rule, Term, Variable, VariableKind,
    is_canonical_unsigned_decimal, is_record_reference_text,
};

use crate::logic::profile::{self, ArgKind, ProfilePredicate, Role};

use super::parser::{
    is_record_reference_variable_name, is_reserved_logic_word, is_text_variable_name, render_term,
};

pub(super) fn check_safety(rules: &[Rule]) -> Result<(), String> {
    for rule in rules {
        // Every variable in the head MUST appear in a positive non-builtin
        // body atom. Anonymous terms are not head terms.
        if rule.head_terms.iter().any(Term::is_anonymous) {
            return Err(format!(
                "anonymous term in head of predicate {}",
                rule.head_predicate
            ));
        }
        let pos_vars = collect_positive_nonbuiltin_vars(rule);
        let head_vars = collect_head_vars(rule);

        for hv in &head_vars {
            if !pos_vars.contains(hv) {
                return Err(format!(
                    "unsafe variable {:?} in head of predicate {}",
                    hv.0, rule.head_predicate
                ));
            }
        }

        // Variables in negated atoms and non-binding builtins
        // MUST appear in a positive non-builtin body atom. Variables that only
        // occur inside Cardinality are local to the cardinality test.
        for atom in &rule.body {
            let constrained = match atom {
                BodyAtom::Negated { terms, .. } => collect_vars(terms),
                BodyAtom::NotEqual(a, b) => {
                    if a.is_anonymous() || b.is_anonymous() {
                        return Err("anonymous term in != body atom".into());
                    }
                    let all_v = collect_terms_vars(&[a.clone(), b.clone()]);
                    if all_v.iter().any(|x| !pos_vars.contains(x)) {
                        return Err("unsafe variable in != body atom".into());
                    }
                    continue;
                }
                BodyAtom::Compare { left, right, .. } => {
                    validate_static_int_term(left, "Compare left")?;
                    validate_static_int_term(right, "Compare right")?;
                    if left.term().is_anonymous() || right.term().is_anonymous() {
                        return Err("anonymous term in Compare body atom".into());
                    }
                    collect_terms_vars(&[left.term().clone(), right.term().clone()])
                }
                BodyAtom::TextMatch { text, .. } => {
                    if text.is_anonymous() {
                        return Err("anonymous term in TextMatch body atom".into());
                    }
                    collect_vars(std::slice::from_ref(text))
                }
                BodyAtom::TextShape {
                    text, start, end, ..
                } => {
                    if text.is_anonymous() || start.is_anonymous() || end.is_anonymous() {
                        return Err("anonymous term in TextShape body atom".into());
                    }
                    collect_terms_vars(&[text.clone(), start.clone(), end.clone()])
                }
                BodyAtom::Cardinality {
                    predicate, terms, ..
                } => {
                    if !predicate_allows_cardinality(predicate, terms.len()) {
                        return Err(format!(
                            "Cardinality cannot count {}/{}",
                            predicate,
                            terms.len()
                        ));
                    }
                    let mut leaking = collect_vars(terms);
                    leaking.retain(|v| pos_vars.contains(v));
                    leaking
                }
                BodyAtom::TopK {
                    predicate,
                    terms,
                    partition,
                    order,
                    ..
                } => {
                    validate_static_int_term(order, "TopK ordering")?;
                    if !predicate_allows_topk(predicate, terms.len()) {
                        return Err(format!("TopK cannot select {}/{}", predicate, terms.len()));
                    }
                    validate_topk_partition(partition, terms)?;
                    if order.term().is_anonymous() || !terms.iter().any(|term| term == order.term())
                    {
                        return Err("TopK ordering term must occur in the candidate atom and must not be anonymous".into());
                    }
                    collect_vars(terms)
                }
                _ => continue,
            };
            for v in &constrained {
                if !pos_vars.contains(v) {
                    return Err(format!(
                        "unsafe variable {:?} in body of rule for {}",
                        v.0, rule.head_predicate
                    ));
                }
            }
        }
    }

    validate_predicate_names_and_shadowing(rules)?;
    check_variable_kinds(rules)?;

    Ok(())
}

fn validate_static_int_term(expression: &OrderingExpression, context: &str) -> Result<(), String> {
    let OrderingExpression::Int(Term::Const(value)) = expression else {
        return Ok(());
    };
    if is_canonical_unsigned_decimal(value) {
        Ok(())
    } else {
        Err(format!(
            "{context} Int constant must be canonical unsigned decimal text"
        ))
    }
}

fn validate_predicate_names_and_shadowing(rules: &[Rule]) -> Result<(), String> {
    let definitions = rules
        .iter()
        .map(|rule| (rule.head_predicate.as_str(), rule.head_terms.len()))
        .collect::<HashSet<_>>();
    for rule in rules {
        validate_head_predicate(rule.head_predicate.as_str(), rule.head_terms.len())?;
        validate_body_predicate_names(&rule.body)?;
        for (predicate, terms, _) in rule.body.iter().filter_map(BodyAtom::relation) {
            if predicate_name_class(predicate) == PredicateNameClass::Lowercase
                && !definitions.contains(&(predicate.as_str(), terms.len()))
            {
                return Err(format!("undefined helper {predicate}/{}", terms.len()));
            }
        }
    }
    Ok(())
}

fn validate_body_predicate_names(body: &[BodyAtom]) -> Result<(), String> {
    for atom in body {
        match atom {
            BodyAtom::Positive { predicate, terms } | BodyAtom::Negated { predicate, terms } => {
                validate_body_predicate(predicate, terms.len())?;
            }
            BodyAtom::Cardinality {
                predicate, terms, ..
            } => {
                validate_body_predicate(predicate, terms.len())?;
                if !predicate_allows_cardinality(predicate, terms.len()) {
                    return Err(format!(
                        "Cardinality cannot count {}/{}",
                        predicate,
                        terms.len()
                    ));
                }
            }
            BodyAtom::TopK {
                predicate, terms, ..
            } => {
                validate_body_predicate(predicate, terms.len())?;
                if !predicate_allows_topk(predicate, terms.len()) {
                    return Err(format!("TopK cannot select {}/{}", predicate, terms.len()));
                }
            }
            BodyAtom::Compare { .. }
            | BodyAtom::TextShape { .. }
            | BodyAtom::TextMatch { .. }
            | BodyAtom::NotEqual(_, _)
            | BodyAtom::True => {}
        }
    }
    Ok(())
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum KindPosition {
    Record,
    RecordReference,
    TextOnly,
    GenericText,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum HelperArgumentKind {
    Text,
    RecordReference,
    GenericText,
}

type InferredHelperArgumentKinds = HashMap<(String, usize), Vec<Option<HelperArgumentKind>>>;

fn check_variable_kinds(rules: &[Rule]) -> Result<(), String> {
    infer_helper_kinds(rules).map(|_| ())
}

fn infer_helper_kinds(rules: &[Rule]) -> Result<InferredHelperArgumentKinds, String> {
    let mut helper_kinds = InferredHelperArgumentKinds::new();
    let defined_helpers = rules
        .iter()
        .filter(|rule| profile::lookup(rule.head_predicate.as_str()).is_none())
        .map(|rule| (rule.head_predicate.to_string(), rule.head_terms.len()))
        .collect::<HashSet<_>>();
    for rule in rules {
        if profile::lookup(rule.head_predicate.as_str()).is_none() {
            helper_kinds
                .entry((rule.head_predicate.to_string(), rule.head_terms.len()))
                .or_insert_with(|| vec![None; rule.head_terms.len()]);
        }
        for (name, terms, _) in rule.body.iter().filter_map(BodyAtom::relation) {
            if profile::lookup(name).is_none() {
                helper_kinds
                    .entry((name.to_string(), terms.len()))
                    .or_insert_with(|| vec![None; terms.len()]);
            }
        }
    }

    let mut changed = true;
    while changed {
        changed = false;
        for rule in rules {
            // Provisional signatures only widen, so a conflict against one may
            // vanish once a later rule widens it: a head's kinds come from its
            // body alone, a body that does not yet fit contributes nothing
            // this pass, and a conflicting contribution is skipped. The final
            // pass checks every rule against the settled signatures.
            let mut vars = HashMap::<Variable, HelperArgumentKind>::new();
            if apply_atom_kind_constraints(&rule.body, &helper_kinds, &mut vars).is_err() {
                continue;
            }
            let head_key = (rule.head_predicate.to_string(), rule.head_terms.len());
            if let Some(sig) = helper_kinds.get_mut(&head_key) {
                for (idx, term) in rule.head_terms.iter().enumerate() {
                    if let Some(kind) = term_inferred_kind(term, &vars)
                        && merge_kind(&mut sig[idx], kind)
                    {
                        changed = true;
                    }
                }
            }
        }
    }

    for (key, sig) in helper_kinds.iter_mut() {
        if !defined_helpers.contains(key) {
            continue;
        }
        for kind in sig.iter_mut() {
            if kind.is_none() {
                *kind = Some(HelperArgumentKind::Text);
            }
        }
    }

    for rule in rules {
        let mut vars = HashMap::<Variable, HelperArgumentKind>::new();
        constrain_predicate_terms(
            rule.head_predicate.as_str(),
            &rule.head_terms,
            &helper_kinds,
            &mut vars,
        )?;
        apply_atom_kind_constraints(&rule.body, &helper_kinds, &mut vars)?;
    }
    Ok(helper_kinds)
}

/// Variable kinds that preserve complete candidate tuples in generated rules.
/// Generic positions use text variables so record-shaped values are not required.
pub fn predicate_variable_kinds(
    rules: &[Rule],
) -> Result<HashMap<(String, usize), Vec<VariableKind>>, String> {
    let mut kinds = infer_helper_kinds(rules)?
        .into_iter()
        .map(|(key, kinds)| {
            let kinds = kinds
                .into_iter()
                .map(|kind| match kind {
                    Some(HelperArgumentKind::RecordReference) => VariableKind::Record,
                    _ => VariableKind::Text,
                })
                .collect();
            (key, kinds)
        })
        .collect::<HashMap<_, _>>();
    for rule in rules {
        for (name, terms, _) in rule.body.iter().filter_map(BodyAtom::relation) {
            // Use unconstrained positions: a call-site field-name constant must
            // not narrow the kind of a complete Field candidate relation.
            let anonymous = vec![Term::Anonymous; terms.len()];
            if let Some(mark) = standard_predicate_mark(name, terms.len(), &anonymous) {
                kinds.insert(
                    (name.to_string(), terms.len()),
                    mark.into_iter()
                        .map(|kind| match kind {
                            KindPosition::Record | KindPosition::RecordReference => {
                                VariableKind::Record
                            }
                            _ => VariableKind::Text,
                        })
                        .collect(),
                );
            }
        }
    }
    Ok(kinds)
}

fn apply_atom_kind_constraints(
    body: &[BodyAtom],
    helper_kinds: &InferredHelperArgumentKinds,
    vars: &mut HashMap<Variable, HelperArgumentKind>,
) -> Result<(), String> {
    for atom in body {
        if let Some((predicate, terms, _)) = atom.relation() {
            constrain_predicate_terms(predicate, terms, helper_kinds, vars)?;
        }
        match atom {
            BodyAtom::TopK { order, .. } => {
                let position = match order {
                    OrderingExpression::Lex(_) => KindPosition::GenericText,
                    OrderingExpression::Int(_) => KindPosition::TextOnly,
                    OrderingExpression::AdmitOrder(_) => KindPosition::Record,
                };
                constrain_term(order.term(), position, vars)?;
            }
            BodyAtom::Compare { left, right, .. } => {
                let position = match left {
                    OrderingExpression::Lex(_) => KindPosition::GenericText,
                    OrderingExpression::Int(_) => KindPosition::TextOnly,
                    OrderingExpression::AdmitOrder(_) => KindPosition::Record,
                };
                constrain_term(left.term(), position, vars)?;
                constrain_term(right.term(), position, vars)?;
            }
            BodyAtom::TextMatch { text, .. } => {
                constrain_term(text, KindPosition::GenericText, vars)?;
            }
            BodyAtom::TextShape {
                text, start, end, ..
            } => {
                constrain_term(text, KindPosition::GenericText, vars)?;
                constrain_term(start, KindPosition::GenericText, vars)?;
                constrain_term(end, KindPosition::GenericText, vars)?;
            }
            BodyAtom::NotEqual(a, b) => {
                if let (Term::Var(left), Term::Var(right)) = (a, b)
                    && left.kind() != right.kind()
                {
                    return Err("!= compares incompatible variable kinds".into());
                }
            }
            BodyAtom::Positive { .. }
            | BodyAtom::Negated { .. }
            | BodyAtom::Cardinality { .. }
            | BodyAtom::True => {}
        }
    }
    Ok(())
}

fn constrain_predicate_terms(
    name: &str,
    terms: &[Term],
    helper_kinds: &InferredHelperArgumentKinds,
    vars: &mut HashMap<Variable, HelperArgumentKind>,
) -> Result<(), String> {
    if let Some(mark) = standard_predicate_mark(name, terms.len(), terms) {
        for (term, position) in terms.iter().zip(mark) {
            constrain_term(term, position, vars)?;
        }
        return Ok(());
    }
    if let Some(mark) = helper_kinds.get(&(name.to_string(), terms.len())) {
        for (term, kind) in terms.iter().zip(mark) {
            if let Some(kind) = kind {
                constrain_term(term, helper_argument_position(*kind), vars)?;
            }
        }
    }
    Ok(())
}

fn helper_argument_position(kind: HelperArgumentKind) -> KindPosition {
    match kind {
        HelperArgumentKind::Text => KindPosition::TextOnly,
        HelperArgumentKind::RecordReference => KindPosition::RecordReference,
        HelperArgumentKind::GenericText => KindPosition::GenericText,
    }
}

fn constrain_term(
    term: &Term,
    position: KindPosition,
    vars: &mut HashMap<Variable, HelperArgumentKind>,
) -> Result<(), String> {
    match term {
        Term::Anonymous => Ok(()),
        Term::Const(value) => {
            match position {
                KindPosition::Record => {
                    crate::records::RecordHash::parse(value).map_err(|_| {
                        format!("record-position constant is not a supported record hash: {value}")
                    })?;
                }
                KindPosition::RecordReference => {
                    if !is_record_reference_text(value) {
                        return Err(format!(
                            "record-reference constant is not record-reference text: {value}"
                        ));
                    }
                }
                KindPosition::TextOnly | KindPosition::GenericText => {}
            }
            Ok(())
        }
        Term::Var(variable) => {
            let required = match position {
                KindPosition::Record | KindPosition::RecordReference => {
                    if variable.kind() != VariableKind::Record {
                        return Err(format!(
                            "{} is text; this position expects a record-reference variable such as `Record`",
                            variable.name()
                        ));
                    }
                    HelperArgumentKind::RecordReference
                }
                KindPosition::TextOnly => {
                    if variable.kind() != VariableKind::Text {
                        return Err(format!(
                            "{} is a record-reference variable; this position expects text",
                            render_term(term)
                        ));
                    }
                    HelperArgumentKind::Text
                }
                KindPosition::GenericText => HelperArgumentKind::GenericText,
            };
            merge_variable_kind(vars, variable.clone(), required)
        }
    }
}

fn term_inferred_kind(
    term: &Term,
    vars: &HashMap<Variable, HelperArgumentKind>,
) -> Option<HelperArgumentKind> {
    match term {
        Term::Var(variable) => vars.get(variable).copied().or(Some(match variable.kind() {
            VariableKind::Record => HelperArgumentKind::RecordReference,
            VariableKind::Text => HelperArgumentKind::Text,
        })),
        Term::Const(value) => {
            is_record_reference_text(value).then_some(HelperArgumentKind::GenericText)
        }
        Term::Anonymous => None,
    }
}

fn merge_variable_kind(
    vars: &mut HashMap<Variable, HelperArgumentKind>,
    variable: Variable,
    kind: HelperArgumentKind,
) -> Result<(), String> {
    match vars.get_mut(&variable) {
        Some(existing) => merge_variable_kind_value(existing, kind).map(|_| ()),
        None => {
            vars.insert(variable, kind);
            Ok(())
        }
    }
}

fn merge_variable_kind_value(
    existing: &mut HelperArgumentKind,
    incoming: HelperArgumentKind,
) -> Result<bool, String> {
    match (*existing, incoming) {
        (left, right) if left == right => Ok(false),
        (HelperArgumentKind::GenericText, specific) => {
            *existing = specific;
            Ok(true)
        }
        (_, HelperArgumentKind::GenericText) => Ok(false),
        _ => Err("variable used with incompatible kinds".into()),
    }
}

/// Widens a provisional helper slot by one rule's contribution, reporting
/// whether it changed. A conflicting contribution leaves the slot as it is.
fn merge_kind(slot: &mut Option<HelperArgumentKind>, kind: HelperArgumentKind) -> bool {
    match (*slot, kind) {
        (None, _) => {
            *slot = Some(kind);
            true
        }
        (Some(HelperArgumentKind::GenericText), _) => false,
        (Some(existing), _) if existing == kind => false,
        (Some(_), HelperArgumentKind::GenericText) => {
            *slot = Some(HelperArgumentKind::GenericText);
            true
        }
        (Some(_), _) => false,
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PredicateNameClass {
    Uppercase,
    Lowercase,
    Other,
}

fn predicate_name_class(name: &str) -> PredicateNameClass {
    if is_record_reference_variable_name(name) {
        PredicateNameClass::Uppercase
    } else if is_text_variable_name(name) {
        PredicateNameClass::Lowercase
    } else {
        PredicateNameClass::Other
    }
}

fn validate_head_predicate(name: &str, arity: usize) -> Result<(), String> {
    let Some(predicate) = profile_predicate(name, arity)? else {
        return Ok(());
    };
    if predicate.role != Role::Facet {
        return Err(format!(
            "derived rule shadows profile predicate {name}/{arity}"
        ));
    }
    Ok(())
}

/// A body may reference profile base facts and helpers, never a facet.
fn validate_body_predicate(name: &str, arity: usize) -> Result<(), String> {
    match profile_predicate(name, arity)? {
        Some(predicate) if predicate.role == Role::Facet => Err(format!(
            "rule bodies must not reference output facet {name}/{arity}"
        )),
        _ => Ok(()),
    }
}

/// The declared profile predicate for a capitalized name, or `None` for a
/// lowercase helper. Capitalized names are reserved for the profile.
fn profile_predicate(
    name: &str,
    arity: usize,
) -> Result<Option<&'static ProfilePredicate>, String> {
    if is_reserved_logic_word(name) {
        return Err(format!(
            "`{name}` is reserved logic syntax, not a predicate"
        ));
    }
    match predicate_name_class(name) {
        PredicateNameClass::Uppercase => {
            let Some(predicate) = profile::lookup(name) else {
                return Err(format!(
                    "unknown capitalized profile predicate {name}/{arity}; capitalized names are reserved for the profile"
                ));
            };
            if predicate.arity() != arity {
                return Err(format!(
                    "{name} expects arity {}, got {arity}",
                    predicate.arity()
                ));
            }
            Ok(Some(predicate))
        }
        PredicateNameClass::Lowercase => Ok(None),
        PredicateNameClass::Other => Err(format!("invalid predicate name {name}/{arity}")),
    }
}

fn predicate_allows_cardinality(name: &str, arity: usize) -> bool {
    match profile::lookup(name) {
        Some(predicate) => {
            predicate.arity() == arity && matches!(predicate.role, Role::Base { .. })
        }
        None => predicate_name_class(name) == PredicateNameClass::Lowercase,
    }
}

fn predicate_allows_topk(name: &str, arity: usize) -> bool {
    predicate_allows_cardinality(name, arity)
}

fn standard_predicate_mark(name: &str, arity: usize, terms: &[Term]) -> Option<Vec<KindPosition>> {
    let predicate = profile::lookup(name).filter(|predicate| predicate.arity() == arity)?;
    Some(
        predicate
            .kinds
            .iter()
            .map(|kind| match kind {
                ArgKind::Record => KindPosition::Record,
                ArgKind::RecordReference => KindPosition::RecordReference,
                ArgKind::Text => KindPosition::TextOnly,
                ArgKind::GenericText => KindPosition::GenericText,
                ArgKind::FieldValue => match terms.get(1).and_then(Term::as_const) {
                    Some(":Type" | ":Length" | "Data-Length" | "By" | "Mark") => {
                        KindPosition::TextOnly
                    }
                    _ => KindPosition::GenericText,
                },
            })
            .collect(),
    )
}

fn validate_topk_partition(partition: &[Variable], terms: &[Term]) -> Result<(), String> {
    let mut seen = HashSet::new();
    for variable in partition {
        if !seen.insert(variable) {
            return Err("TopK PartitionBy variables must be distinct".into());
        }
        if !terms
            .iter()
            .any(|term| term == &Term::Var(variable.clone()))
        {
            return Err("TopK PartitionBy variable must occur in the candidate atom".into());
        }
    }
    Ok(())
}

fn collect_positive_nonbuiltin_vars(rule: &Rule) -> HashSet<Variable> {
    let mut vars = HashSet::new();
    for atom in &rule.body {
        match atom {
            BodyAtom::Positive { terms, .. } | BodyAtom::TopK { terms, .. } => {
                vars.extend(collect_vars(terms));
            }
            _ => {}
        }
    }
    vars
}

fn collect_head_vars(rule: &Rule) -> HashSet<Variable> {
    collect_vars(&rule.head_terms)
}

fn collect_vars(terms: &[Term]) -> HashSet<Variable> {
    let mut vars = HashSet::new();
    for t in terms {
        if let Term::Var(v) = t {
            vars.insert(v.clone());
        }
    }
    vars
}

fn collect_terms_vars(terms: &[Term]) -> HashSet<Variable> {
    collect_vars(terms)
}