xlog_logic/

parser.rs

//! Parser for XLOG programs using Pest
//!
//! This module parses XLOG source text into AST structures.
//! It uses the Pest parser generator with a grammar defined in `grammar.pest`.
#![expect(
    missing_docs,
    reason = "`pest_derive` emits a public grammar Rule enum and helpers without documentation hooks"
)]

use pest::iterators::Pair;
use pest::Parser;
use pest_derive::Parser;
use xlog_core::{symbol, Result, ScalarType, XlogError};

use crate::ast::{
    AggExpr, AggOp, AnnotatedDisjunction, ArithExpr, Atom, BodyLiteral, CompOp, Comparison,
    CondExpr, Constraint, DomainDecl, EpistemicLiteral, EpistemicMode, EpistemicOp, Evidence,
    FuncBody, FuncDef, FuncParam, IsExpr, LearnableRule, MagicSetsMode, NeuralLabel,
    NeuralPredDecl, PredColumn, PredDecl, ProbCache, ProbEngine, ProbFact, ProbMethod, ProbQuery,
    Program, Query, Rule as AstRule, Term, TypeRef, Univ, UseDecl,
};

/// Pest-based parser for XLOG Datalog syntax.
#[derive(Parser)]
#[grammar = "grammar.pest"]
pub struct XlogParser;

/// Parse result containing the parsed pairs (for low-level access)
pub type ParseResult<'a> = pest::iterators::Pairs<'a, Rule>;

/// Parse an XLOG program string into an AST Program.
///
/// Simple ground facts (`p(1, "a", sym, X, _).`) take a hand-written fast
/// path (see [`crate::fact_fast_path`]); everything else goes through pest.
/// The result is identical to [`parse_program_reference`], including errors.
pub fn parse_program(input: &str) -> Result<Program> {
    parse_program_with_stats(input).map(|(program, _)| program)
}

/// Pure-pest reference parser: the whole source through the grammar, no fast
/// path. `parse_program` is specified to be observationally equal to this.
pub fn parse_program_reference(input: &str) -> Result<Program> {
    let pairs =
        XlogParser::parse(Rule::program, input).map_err(|e| XlogError::Parse(e.to_string()))?;

    build_program(pairs)
}

/// How much of a parse the fact fast path handled.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct FastParseStats {
    /// Facts built without pest.
    pub fast_facts: usize,
    /// Bytes of source pest still had to parse (everything but the facts).
    pub residual_bytes: usize,
    /// The residual failed to parse/build, so the whole source was re-parsed by
    /// the reference parser (its result is what was returned).
    pub fell_back: bool,
}

/// [`parse_program`] plus fast-path statistics (for probes and tests).
pub fn parse_program_with_stats(input: &str) -> Result<(Program, FastParseStats)> {
    let mut scan = crate::fact_fast_path::scan(input);
    if scan.facts.is_empty() {
        let stats = FastParseStats {
            fast_facts: 0,
            residual_bytes: input.len(),
            fell_back: false,
        };
        return parse_program_reference(input).map(|program| (program, stats));
    }
    let mut stats = FastParseStats {
        fast_facts: scan.facts.len(),
        residual_bytes: input.len() - scan.fact_bytes,
        fell_back: false,
    };
    let facts = std::mem::take(&mut scan.facts);
    let merged = XlogParser::parse(Rule::program, &scan.residual)
        .map_err(|e| XlogError::Parse(e.to_string()))
        .and_then(|pairs| build_program_merged(pairs, facts));
    match merged {
        Ok(program) => Ok((program, stats)),
        Err(_) => {
            // Any failure on the residual: report exactly what the reference
            // parser reports for the original source (positions included).
            stats.fell_back = true;
            parse_program_reference(input).map(|program| (program, stats))
        }
    }
}

/// Parse a single statement (low-level, returns pest pairs)
pub fn parse_statement(input: &str) -> Result<ParseResult<'_>> {
    XlogParser::parse(Rule::statement, input).map_err(|e| XlogError::Parse(e.to_string()))
}

/// Parse a single atom (low-level, returns pest pairs)
pub fn parse_atom(input: &str) -> Result<ParseResult<'_>> {
    XlogParser::parse(Rule::atom, input).map_err(|e| XlogError::Parse(e.to_string()))
}

// =============================================================================
// AST Building Functions
// =============================================================================

/// Build a Program AST from the pest pairs of the residual source plus the
/// fast-path facts, interleaved in original source order (statement order is
/// semantically relevant: it fixes rule order and therefore the plan). The
/// residual has the facts blanked out, so pest's span offsets are original
/// offsets. Facts intern their symbols here, in merge order, so interning
/// order is source order exactly as in the reference parser.
fn build_program_merged(
    pairs: ParseResult<'_>,
    facts: Vec<crate::fact_fast_path::FastFact<'_>>,
) -> Result<Program> {
    let mut program = Program::new();
    let mut facts = facts.into_iter().peekable();
    for pair in pairs {
        if pair.as_rule() != Rule::program {
            continue;
        }
        for inner in pair.into_inner() {
            if inner.as_rule() != Rule::statement {
                continue;
            }
            let stmt_offset = inner.as_span().start();
            while let Some(fact) = facts.next_if(|fact| fact.offset < stmt_offset) {
                program.rules.push(fact.into_rule());
            }
            build_statement(inner, &mut program)?;
        }
    }
    for fact in facts {
        program.rules.push(fact.into_rule());
    }
    Ok(program)
}

/// Build a Program AST from parsed pest pairs
fn build_program(pairs: ParseResult<'_>) -> Result<Program> {
    let mut program = Program::new();

    for pair in pairs {
        if pair.as_rule() == Rule::program {
            for inner in pair.into_inner() {
                match inner.as_rule() {
                    Rule::statement => {
                        build_statement(inner, &mut program)?;
                    }
                    Rule::EOI => {}
                    _ => {}
                }
            }
        }
    }

    Ok(program)
}

/// Build a statement and add it to the program
fn build_statement(pair: Pair<'_, Rule>, program: &mut Program) -> Result<()> {
    for inner in pair.into_inner() {
        match inner.as_rule() {
            Rule::use_stmt => {
                program.imports.push(build_use_stmt(inner));
            }
            Rule::domain_decl => {
                program.domains.push(build_domain_decl(inner)?);
            }
            Rule::pred_decl => {
                program.predicates.push(build_pred_decl(inner)?);
            }
            Rule::pragma => {
                apply_pragma(inner, program)?;
            }
            Rule::rule_def => {
                program.rules.push(build_rule(inner)?);
            }
            Rule::fact => {
                program.rules.push(build_fact(inner)?);
            }
            Rule::prob_fact => {
                program.prob_facts.push(build_prob_fact(inner)?);
            }
            Rule::annotated_disjunction => {
                program
                    .annotated_disjunctions
                    .push(build_annotated_disjunction(inner)?);
            }
            Rule::evidence_stmt => {
                program.evidence.push(build_evidence(inner)?);
            }
            Rule::prob_query => {
                program.prob_queries.push(build_prob_query(inner)?);
            }
            Rule::constraint => {
                program.constraints.push(build_constraint(inner)?);
            }
            Rule::query => {
                program.queries.push(build_query(inner)?);
            }
            Rule::func_def => {
                program.functions.push(parse_func_def(inner)?);
            }
            Rule::neural_pred_decl => {
                program
                    .neural_predicates
                    .push(build_neural_pred_decl(inner)?);
            }
            Rule::learnable_rule => {
                program.learnable_rules.push(build_learnable_rule(inner)?);
            }
            _ => {}
        }
    }
    Ok(())
}

fn apply_pragma(pair: Pair<'_, Rule>, program: &mut Program) -> Result<()> {
    let pragma = pair
        .into_inner()
        .next()
        .ok_or_else(|| XlogError::Parse("Empty pragma".to_string()))?;

    match pragma.as_rule() {
        Rule::pragma_prob_engine => {
            let value = pragma
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing prob_engine value".to_string()))?;
            let engine = match value.as_str() {
                "exact_ddnnf" => ProbEngine::ExactDdnnf,
                "mc" => ProbEngine::Mc,
                other => {
                    return Err(XlogError::Parse(format!(
                        "Unknown prob_engine value: {}",
                        other
                    )))
                }
            };
            program.directives.prob_engine = Some(engine);
        }
        Rule::pragma_prob_cache => {
            let value = pragma
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing prob_cache value".to_string()))?;
            let cache = match value.as_str() {
                "on" => ProbCache::On,
                "off" => ProbCache::Off,
                other => {
                    return Err(XlogError::Parse(format!(
                        "Unknown prob_cache value: {}",
                        other
                    )))
                }
            };
            program.directives.prob_cache = Some(cache);
        }
        Rule::pragma_epistemic_mode => {
            let value = pragma
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing epistemic_mode value".to_string()))?;
            let mode = match value.as_str() {
                "g91" => EpistemicMode::G91,
                "faeel" => EpistemicMode::Faeel,
                other => {
                    return Err(XlogError::Parse(format!(
                        "Unknown epistemic_mode value: {}",
                        other
                    )))
                }
            };
            program.directives.epistemic_mode = Some(mode);
        }
        Rule::pragma_prob_samples => {
            let value = pragma
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing prob_samples value".to_string()))?;
            let samples: usize = value.as_str().parse().map_err(|_| {
                XlogError::Parse(format!("Invalid prob_samples value: {}", value.as_str()))
            })?;
            if samples == 0 {
                return Err(XlogError::Parse(
                    "Invalid prob_samples value: expected > 0".to_string(),
                ));
            }
            program.directives.prob_samples = Some(samples);
        }
        Rule::pragma_prob_seed => {
            let value = pragma
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing prob_seed value".to_string()))?;
            let seed: u64 = value.as_str().parse().map_err(|_| {
                XlogError::Parse(format!("Invalid prob_seed value: {}", value.as_str()))
            })?;
            program.directives.prob_seed = Some(seed);
        }
        Rule::pragma_prob_confidence => {
            let value = pragma
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing prob_confidence value".to_string()))?;
            let confidence: f64 = value.as_str().parse().map_err(|_| {
                XlogError::Parse(format!("Invalid prob_confidence value: {}", value.as_str()))
            })?;
            if !(0.0 < confidence && confidence < 1.0) || confidence.is_nan() {
                return Err(XlogError::Parse(format!(
                    "Invalid prob_confidence value: {}; expected 0 < confidence < 1",
                    value.as_str()
                )));
            }
            program.directives.prob_confidence = Some(confidence);
        }
        Rule::pragma_prob_method => {
            let value = pragma
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing prob_method value".to_string()))?;
            let method = match value.as_str() {
                "rejection" => ProbMethod::Rejection,
                "evidence_clamping" => ProbMethod::EvidenceClamping,
                other => {
                    return Err(XlogError::Parse(format!(
                        "Unknown prob_method value: {}",
                        other
                    )))
                }
            };
            program.directives.prob_method = Some(method);
        }
        Rule::pragma_prob_max_nonmonotone_iterations => {
            let value = pragma.into_inner().next().ok_or_else(|| {
                XlogError::Parse("Missing prob_max_nonmonotone_iterations value".to_string())
            })?;
            let iterations: usize = value.as_str().parse().map_err(|_| {
                XlogError::Parse(format!(
                    "Invalid prob_max_nonmonotone_iterations value: {}",
                    value.as_str()
                ))
            })?;
            if iterations == 0 {
                return Err(XlogError::Parse(
                    "Invalid prob_max_nonmonotone_iterations value: expected > 0".to_string(),
                ));
            }
            program.directives.prob_max_nonmonotone_iterations = Some(iterations);
        }
        Rule::pragma_max_recursion => {
            let value = pragma
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing max_recursion_depth value".to_string()))?;
            let depth: u32 = value.as_str().parse().map_err(|_| {
                XlogError::Parse(format!(
                    "Invalid max_recursion_depth value: {}",
                    value.as_str()
                ))
            })?;
            program.directives.max_recursion_depth = Some(depth);
        }
        Rule::pragma_magic_sets => {
            let value = pragma
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing magic_sets value".to_string()))?;
            let mode = match value.as_str() {
                "auto" => MagicSetsMode::Auto,
                "on" => MagicSetsMode::On,
                "off" => MagicSetsMode::Off,
                other => {
                    return Err(XlogError::Parse(format!(
                        "Unknown magic_sets value: {}",
                        other
                    )))
                }
            };
            program.directives.magic_sets = Some(mode);
        }
        _ => {}
    }

    Ok(())
}

/// Build a domain declaration
fn build_domain_decl(pair: Pair<'_, Rule>) -> Result<DomainDecl> {
    let mut inner = pair.into_inner();
    let name = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing domain name".to_string()))?
        .as_str()
        .to_string();
    let type_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing domain type".to_string()))?;
    let typ = build_scalar_type_spec(type_pair, "domain alias")?;

    Ok(DomainDecl { name, typ })
}

/// Parse a module path (e.g., "utils/math" -> ["utils", "math"])
fn parse_module_path(pair: Pair<Rule>) -> Vec<String> {
    pair.as_str().split('/').map(|s| s.to_string()).collect()
}

/// Build a use statement
fn build_use_stmt(pair: Pair<Rule>) -> UseDecl {
    let mut inner = pair.into_inner();

    // Parse module path
    let path_pair = inner.next().unwrap();
    let module_path = parse_module_path(path_pair);

    // Parse optional import list
    let imports = inner.next().map(|import_list| {
        import_list
            .into_inner()
            .map(|p| p.as_str().to_string())
            .collect()
    });

    UseDecl {
        module_path,
        imports,
    }
}

/// Build a predicate declaration
fn build_pred_decl(pair: Pair<'_, Rule>) -> Result<PredDecl> {
    let mut inner = pair.into_inner();
    let mut is_private = false;

    // Check for private modifier
    let first = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing predicate name".to_string()))?;

    let name_pair = if first.as_rule() == Rule::private_mod {
        is_private = true;
        inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing predicate name after private".to_string()))?
    } else {
        first
    };

    let name = name_pair.as_str().to_string();

    let mut columns = Vec::new();
    for type_pair in inner {
        if type_pair.as_rule() == Rule::type_list {
            for col in type_pair.into_inner() {
                columns.push(build_pred_column(col)?);
            }
        }
    }
    let types = columns.iter().map(|c| c.typ.clone()).collect();

    Ok(PredDecl {
        name,
        types,
        columns,
        is_private,
    })
}

fn build_pred_column(pair: Pair<'_, Rule>) -> Result<PredColumn> {
    let mut inner = pair.into_inner();
    let first = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Empty predicate column".to_string()))?;

    if first.as_rule() == Rule::ident {
        let typ_pair = inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing named column type".to_string()))?;
        Ok(PredColumn {
            name: Some(first.as_str().to_string()),
            typ: build_type_ref(typ_pair)?,
        })
    } else {
        Ok(PredColumn {
            name: None,
            typ: build_type_ref(first)?,
        })
    }
}

/// Build a source type reference.
fn build_type_ref(pair: Pair<'_, Rule>) -> Result<TypeRef> {
    if pair.as_rule() == Rule::type_spec {
        let raw = pair.as_str().to_string();
        let mut inner = pair.into_inner();
        if let Some(child) = inner.next() {
            return build_type_ref(child);
        }
        return match raw.as_str() {
            "term" => Ok(TypeRef::Term),
            "compound" => Ok(TypeRef::Compound),
            "predref" => Ok(TypeRef::PredRef),
            other => Err(XlogError::Parse(format!("Unknown type: {}", other))),
        };
    }

    match pair.as_rule() {
        Rule::list_type => {
            let item = pair
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing list element type".to_string()))?;
            Ok(TypeRef::List(Box::new(build_type_ref(item)?)))
        }
        Rule::scalar_type => build_scalar_type_name(pair.as_str()).map(TypeRef::Scalar),
        Rule::ident => Ok(TypeRef::Domain(pair.as_str().to_string())),
        _ => match pair.as_str() {
            "term" => Ok(TypeRef::Term),
            "compound" => Ok(TypeRef::Compound),
            "predref" => Ok(TypeRef::PredRef),
            other => Err(XlogError::Parse(format!("Unknown type: {}", other))),
        },
    }
}

fn build_scalar_type_spec(pair: Pair<'_, Rule>, context: &str) -> Result<ScalarType> {
    match build_type_ref(pair)? {
        TypeRef::Scalar(ty) => Ok(ty),
        other => Err(XlogError::Parse(format!(
            "v0.8.5 {} must use a scalar type, got {:?}",
            context, other
        ))),
    }
}

fn build_scalar_type_name(type_str: &str) -> Result<ScalarType> {
    match type_str {
        "u32" => Ok(ScalarType::U32),
        "u64" => Ok(ScalarType::U64),
        "i32" => Ok(ScalarType::I32),
        "i64" => Ok(ScalarType::I64),
        "f32" => Ok(ScalarType::F32),
        "f64" => Ok(ScalarType::F64),
        "bool" => Ok(ScalarType::Bool),
        "symbol" => Ok(ScalarType::Symbol),
        _ => Err(XlogError::Parse(format!(
            "Unknown scalar type: {}",
            type_str
        ))),
    }
}

/// Parse a function parameter: X or X: f64
fn parse_func_param(pair: Pair<'_, Rule>) -> Result<FuncParam> {
    let mut inner = pair.into_inner();
    let name = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing parameter name".to_string()))?
        .as_str()
        .to_string();
    let typ = inner
        .next()
        .map(|ta| {
            // type_annotation contains type_spec
            build_scalar_type_spec(
                ta.into_inner()
                    .next()
                    .expect("type_annotation must contain type_spec"),
                "function parameter type annotation",
            )
        })
        .transpose()?;
    Ok(FuncParam { name, typ })
}

/// Parse a comparison operator
fn parse_cmp_op(pair: Pair<'_, Rule>) -> CompOp {
    match pair.as_str() {
        "==" | "=" => CompOp::Eq,
        "!=" => CompOp::Ne,
        "<" => CompOp::Lt,
        "<=" => CompOp::Le,
        ">" => CompOp::Gt,
        ">=" => CompOp::Ge,
        _ => unreachable!("unexpected comparison operator: {}", pair.as_str()),
    }
}

/// Parse a conditional expression: if X < 0 then 0 - X else X
fn parse_cond_expr(pair: Pair<'_, Rule>) -> Result<CondExpr> {
    let mut inner = pair.into_inner();

    // Parse condition test
    let cond_test = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing condition test".to_string()))?;
    let mut test_inner = cond_test.into_inner();
    let cond_left = build_arith_expr(
        test_inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing left side of condition".to_string()))?,
    )?;
    let cond_op = parse_cmp_op(
        test_inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing condition operator".to_string()))?,
    );
    let cond_right = build_arith_expr(
        test_inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing right side of condition".to_string()))?,
    )?;

    // Parse then branch
    let then_branch =
        Box::new(parse_func_body(inner.next().ok_or_else(|| {
            XlogError::Parse("Missing then branch".to_string())
        })?)?);

    // Parse else branch
    let else_branch =
        Box::new(parse_func_body(inner.next().ok_or_else(|| {
            XlogError::Parse("Missing else branch".to_string())
        })?)?);

    Ok(CondExpr {
        cond_left,
        cond_op,
        cond_right,
        then_branch,
        else_branch,
    })
}

/// Parse a function body: arithmetic, conditional, or predicate-based
fn parse_func_body(pair: Pair<'_, Rule>) -> Result<FuncBody> {
    let inner = pair
        .into_inner()
        .next()
        .ok_or_else(|| XlogError::Parse("Empty function body".to_string()))?;
    match inner.as_rule() {
        Rule::func_body_pred => {
            let mut parts = inner.into_inner();
            let result = parts
                .next()
                .ok_or_else(|| XlogError::Parse("Missing result variable".to_string()))?
                .as_str()
                .to_string();
            let body = build_body(
                parts
                    .next()
                    .ok_or_else(|| XlogError::Parse("Missing predicate body".to_string()))?,
            )?;
            Ok(FuncBody::Predicate { result, body })
        }
        Rule::func_body_arith => {
            let arith_inner = inner
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Empty arithmetic body".to_string()))?;
            match arith_inner.as_rule() {
                Rule::cond_expr => Ok(FuncBody::Conditional(parse_cond_expr(arith_inner)?)),
                _ => Ok(FuncBody::Arithmetic(build_arith_expr(arith_inner)?)),
            }
        }
        _ => Err(XlogError::Parse(format!(
            "Unexpected rule in func_body: {:?}",
            inner.as_rule()
        ))),
    }
}

/// Parse a function definition
fn parse_func_def(pair: Pair<'_, Rule>) -> Result<FuncDef> {
    let mut inner = pair.into_inner();
    let mut is_private = false;

    // Check for private modifier
    let first = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Empty function definition".to_string()))?;
    let name_pair = if first.as_rule() == Rule::private_mod {
        is_private = true;
        inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing function name after private".to_string()))?
    } else {
        first
    };

    let name = name_pair.as_str().to_string();

    let mut params = Vec::new();
    let mut return_type = None;
    let mut body = None;

    for p in inner {
        match p.as_rule() {
            Rule::func_params => {
                params = p
                    .into_inner()
                    .map(parse_func_param)
                    .collect::<Result<Vec<_>>>()?;
            }
            Rule::return_type => {
                return_type = Some(build_scalar_type_spec(
                    p.into_inner()
                        .next()
                        .ok_or_else(|| XlogError::Parse("Missing return type".to_string()))?,
                    "function return type annotation",
                )?);
            }
            Rule::func_body => {
                body = Some(parse_func_body(p)?);
            }
            _ => {}
        }
    }

    Ok(FuncDef {
        name,
        params,
        return_type,
        body: body.ok_or_else(|| XlogError::Parse("Function must have a body".to_string()))?,
        is_private,
    })
}

/// Build a rule (with body)
fn build_rule(pair: Pair<'_, Rule>) -> Result<AstRule> {
    let mut inner = pair.into_inner();
    let head_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing rule head".to_string()))?;
    let head = build_head(head_pair)?;

    let body_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing rule body".to_string()))?;
    let body = build_body(body_pair)?;

    Ok(AstRule { head, body })
}

/// Build a fact (rule with empty body)
fn build_fact(pair: Pair<'_, Rule>) -> Result<AstRule> {
    let mut inner = pair.into_inner();
    let atom_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing fact atom".to_string()))?;
    let head = build_atom(atom_pair)?;

    Ok(AstRule { head, body: vec![] })
}

/// Build a constraint
fn build_constraint(pair: Pair<'_, Rule>) -> Result<Constraint> {
    let mut inner = pair.into_inner();
    let body_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing constraint body".to_string()))?;
    let body = build_body(body_pair)?;

    Ok(Constraint {
        authored_index: None,
        body,
    })
}

/// Build a query
fn build_query(pair: Pair<'_, Rule>) -> Result<Query> {
    let mut inner = pair.into_inner();
    let atom_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing query atom".to_string()))?;
    let atom = build_atom(atom_pair)?;

    Ok(Query { atom })
}

fn build_prob_fact(pair: Pair<'_, Rule>) -> Result<ProbFact> {
    let choice = pair
        .into_inner()
        .next()
        .ok_or_else(|| XlogError::Parse("Missing probabilistic fact".to_string()))?;
    build_prob_choice(choice)
}

fn build_annotated_disjunction(pair: Pair<'_, Rule>) -> Result<AnnotatedDisjunction> {
    let mut choices = Vec::new();
    for inner in pair.into_inner() {
        if inner.as_rule() == Rule::prob_choice {
            choices.push(build_prob_choice(inner)?);
        }
    }
    if choices.is_empty() {
        return Err(XlogError::Parse(
            "Annotated disjunction must have at least one choice".to_string(),
        ));
    }
    Ok(AnnotatedDisjunction { choices })
}

fn build_prob_choice(pair: Pair<'_, Rule>) -> Result<ProbFact> {
    let mut inner = pair.into_inner();
    let prob_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing probability".to_string()))?;
    let prob: f64 = prob_pair
        .as_str()
        .parse()
        .map_err(|_| XlogError::Parse(format!("Invalid probability: {}", prob_pair.as_str())))?;

    let atom_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing probabilistic atom".to_string()))?;
    let atom = build_atom(atom_pair)?;

    Ok(ProbFact { prob, atom })
}

fn build_evidence(pair: Pair<'_, Rule>) -> Result<Evidence> {
    let mut inner = pair.into_inner();
    let atom_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing evidence atom".to_string()))?;
    let atom = build_atom(atom_pair)?;

    let value_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing evidence value".to_string()))?;
    let value = match value_pair.as_str() {
        "true" => true,
        "false" => false,
        other => {
            return Err(XlogError::Parse(format!(
                "Invalid evidence value (expected true/false): {}",
                other
            )))
        }
    };

    Ok(Evidence { atom, value })
}

fn build_prob_query(pair: Pair<'_, Rule>) -> Result<ProbQuery> {
    let atom_pair = pair
        .into_inner()
        .next()
        .ok_or_else(|| XlogError::Parse("Missing query atom".to_string()))?;
    let atom = build_atom(atom_pair)?;
    Ok(ProbQuery { atom })
}

/// Build a neural predicate declaration
/// Syntax: nn(network, [inputs], output, [labels]) :: pred(args).
fn build_neural_pred_decl(pair: Pair<'_, Rule>) -> Result<NeuralPredDecl> {
    let mut inner = pair.into_inner();

    // Parse network name (ident)
    let network = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing network name in neural predicate".to_string()))?
        .as_str()
        .to_string();

    // Parse input list
    let input_list = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing input list in neural predicate".to_string()))?;
    let inputs: Vec<String> = input_list
        .into_inner()
        .map(|p| p.as_str().to_string())
        .collect();

    // Parse output variable
    let output = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing output variable in neural predicate".to_string()))?
        .as_str()
        .to_string();

    // Check for optional label list or atom
    let next = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing predicate in neural predicate".to_string()))?;

    let (labels, predicate) = if next.as_rule() == Rule::neural_label_list {
        // Parse labels
        let label_vec: Vec<NeuralLabel> = next
            .into_inner()
            .map(|label_pair| {
                let label_inner = label_pair
                    .into_inner()
                    .next()
                    .expect("neural_label must have inner");
                match label_inner.as_rule() {
                    Rule::integer => {
                        let val: i64 = label_inner.as_str().parse().expect("valid integer");
                        NeuralLabel::Integer(val)
                    }
                    Rule::ident => NeuralLabel::Symbol(label_inner.as_str().to_string()),
                    _ => unreachable!("neural_label should be integer or ident"),
                }
            })
            .collect();

        // Parse atom
        let atom_pair = inner.next().ok_or_else(|| {
            XlogError::Parse("Missing predicate after labels in neural predicate".to_string())
        })?;
        let predicate = build_atom(atom_pair)?;

        (Some(label_vec), predicate)
    } else {
        // No labels, next is the atom (embedding mode)
        let predicate = build_atom(next)?;
        (None, predicate)
    };

    Ok(NeuralPredDecl {
        network,
        inputs,
        output,
        labels,
        predicate,
    })
}

/// Build a learnable rule from a parsed pair.
/// Grammar: learnable_rule = { "learnable" ~ "(" ~ ident ~ ")" ~ "::" ~ head ~ ":-" ~ body ~ "." }
/// Uses build_head instead of build_atom because the grammar produces a `head` pair.
fn build_learnable_rule(pair: Pair<'_, Rule>) -> Result<LearnableRule> {
    let mut inner = pair.into_inner();
    let mask_name = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing learnable mask name".into()))?
        .as_str()
        .to_string();
    let head = build_head(
        inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing learnable head".into()))?,
    )?;
    let body = build_body(
        inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing learnable body".into()))?,
    )?;
    Ok(LearnableRule {
        mask_name,
        head,
        body,
    })
}

/// Build a head (atom that may contain aggregates)
fn build_head(pair: Pair<'_, Rule>) -> Result<Atom> {
    let mut inner = pair.into_inner();
    let predicate = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing head predicate".to_string()))?
        .as_str()
        .to_string();

    let mut terms = Vec::new();
    for term_list in inner {
        if term_list.as_rule() == Rule::head_term_list {
            for head_term in term_list.into_inner() {
                terms.push(build_head_term(head_term)?);
            }
        }
    }

    Ok(Atom { predicate, terms })
}

/// Build a head term (can be aggregate or regular term)
fn build_head_term(pair: Pair<'_, Rule>) -> Result<Term> {
    let inner = pair
        .into_inner()
        .next()
        .ok_or_else(|| XlogError::Parse("Empty head term".to_string()))?;

    match inner.as_rule() {
        Rule::aggregate => build_aggregate(inner),
        Rule::agg_term => {
            // agg_term can contain aggregate or term
            let agg_inner = inner
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Empty agg_term".to_string()))?;
            match agg_inner.as_rule() {
                Rule::aggregate => build_aggregate(agg_inner),
                Rule::term => build_term(agg_inner),
                _ => build_term(agg_inner),
            }
        }
        Rule::term => build_term(inner),
        _ => build_term(inner),
    }
}

/// Build an aggregate expression
fn build_aggregate(pair: Pair<'_, Rule>) -> Result<Term> {
    let mut inner = pair.into_inner();
    let op_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing aggregate operator".to_string()))?;
    let op = match op_pair.as_str() {
        "count" => AggOp::Count,
        "sum" => AggOp::Sum,
        "min" => AggOp::Min,
        "max" => AggOp::Max,
        "logsumexp" => AggOp::LogSumExp,
        _ => {
            return Err(XlogError::Parse(format!(
                "Unknown aggregate: {}",
                op_pair.as_str()
            )))
        }
    };

    let var_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing aggregate variable".to_string()))?;
    let variable = var_pair.as_str().to_string();

    Ok(Term::Aggregate(AggExpr { op, variable }))
}

/// Build an atom
fn build_atom(pair: Pair<'_, Rule>) -> Result<Atom> {
    let mut inner = pair.into_inner();
    let predicate = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing atom predicate".to_string()))?
        .as_str()
        .to_string();

    let mut terms = Vec::new();
    for term_list in inner {
        if term_list.as_rule() == Rule::term_list {
            for term in term_list.into_inner() {
                terms.push(build_term(term)?);
            }
        }
    }

    Ok(Atom { predicate, terms })
}

/// Build a body (list of literals)
fn build_body(pair: Pair<'_, Rule>) -> Result<Vec<BodyLiteral>> {
    let mut literals = Vec::new();

    for lit in pair.into_inner() {
        literals.push(build_body_literal(lit)?);
    }

    Ok(literals)
}

/// Build a body literal
fn build_body_literal(pair: Pair<'_, Rule>) -> Result<BodyLiteral> {
    let inner = pair
        .into_inner()
        .next()
        .ok_or_else(|| XlogError::Parse("Empty body literal".to_string()))?;

    match inner.as_rule() {
        Rule::nested_modal_chain => build_nested_modal_chain(inner),
        Rule::negated_epistemic_atom => build_epistemic_literal(inner, true),
        Rule::epistemic_atom => build_epistemic_literal(inner, false),
        Rule::negated_atom => {
            let atom_pair = inner
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Missing negated atom".to_string()))?;
            Ok(BodyLiteral::Negated(build_atom(atom_pair)?))
        }
        Rule::atom => Ok(BodyLiteral::Positive(build_atom(inner)?)),
        Rule::comparison => Ok(BodyLiteral::Comparison(build_comparison(inner)?)),
        Rule::is_expr => Ok(BodyLiteral::IsExpr(build_is_expr(inner)?)),
        Rule::univ => Ok(BodyLiteral::Univ(build_univ(inner)?)),
        _ => Err(XlogError::Parse(format!(
            "Unknown body literal: {:?}",
            inner.as_rule()
        ))),
    }
}

fn build_epistemic_literal(pair: Pair<'_, Rule>, negated: bool) -> Result<BodyLiteral> {
    let epistemic_pair = if pair.as_rule() == Rule::negated_epistemic_atom {
        pair.into_inner()
            .next()
            .ok_or_else(|| XlogError::Parse("Missing negated epistemic literal".to_string()))?
    } else {
        pair
    };

    let mut inner = epistemic_pair.into_inner();
    let op_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing epistemic operator".to_string()))?;
    let op = match op_pair.as_str() {
        "know" => EpistemicOp::Know,
        "possible" => EpistemicOp::Possible,
        other => {
            return Err(XlogError::Parse(format!(
                "Unknown epistemic operator: {}",
                other
            )))
        }
    };
    let atom_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing epistemic atom".to_string()))?;

    Ok(BodyLiteral::Epistemic(EpistemicLiteral {
        op,
        negated,
        atom: build_atom(atom_pair)?,
    }))
}

/// Build a single epistemic literal from a NESTED modal chain by applying the
/// sound modal-logic (KD45/S5) collapse equivalence.
///
/// Under the autoepistemic modal axioms XLOG's `know`/`possible` operators assume
/// (positive introspection, axiom 4 `Kp → KKp`, and negative introspection,
/// axiom 5 `¬Kp → K¬Kp`, evaluated relative to the admissible world-view set),
/// a chain of modal operators over an atom COLLAPSES to the operator ADJACENT to
/// the atom (the innermost one):
///   `know possible p ≡ possible p`     (KM ≡ M)
///   `possible know  p ≡ know p`        (MK ≡ K)
///   `know know       p ≡ know p`       (KK ≡ K)
///   `possible possible p ≡ possible p` (MM ≡ M)
/// because each adjacent operator pair reduces by 4/5 and the inner operator wins.
/// A single LEADING negation distributes over the whole chain:
///   `not know possible p ≡ not possible p`.
/// This equivalence holds in BOTH epistemic modes (FAEEL and G91): introspection
/// holds WITHIN any single admissible world view in both modes; the modes differ
/// only in WHICH world views are admissible (founded least vs all stable), which
/// is exactly the single-level per-mode difference the collapsed literal inherits
/// by routing through the ordinary single-level epistemic path. The collapse adds
/// no new world-of-worlds evaluator.
///
/// A `not` before any operator negates the modal subformula to its right. A
/// `not` before the atom dualizes the atom-adjacent modal (`know not p` becomes
/// `not possible p`; `possible not p` becomes `not know p`). Under the same S5
/// collapse, all outer modal operators preserve the already-global inner modal
/// truth value, so the final single-level literal is determined by the
/// atom-adjacent operator, atom negation duality, and parity of `not` placements.
fn build_nested_modal_chain(pair: Pair<'_, Rule>) -> Result<BodyLiteral> {
    // Walk the chain left-to-right. Each `not_kw` negates whatever named token
    // follows it (the next operator, or the trailing atom). We record the
    // operators in order plus the negation that immediately precedes each, and
    // whether the trailing atom is preceded by a `not`.
    let mut ops: Vec<EpistemicOp> = Vec::new();
    let mut neg_before_op: Vec<bool> = Vec::new();
    let mut atom_pair: Option<Pair<'_, Rule>> = None;
    let mut pending_not = false;
    let mut neg_before_atom = false;

    for token in pair.into_inner() {
        match token.as_rule() {
            Rule::not_kw => {
                pending_not = true;
            }
            Rule::epistemic_op => {
                let op = match token.as_str() {
                    "know" => EpistemicOp::Know,
                    "possible" => EpistemicOp::Possible,
                    other => {
                        return Err(XlogError::Parse(format!(
                            "Unknown epistemic operator: {}",
                            other
                        )))
                    }
                };
                ops.push(op);
                neg_before_op.push(pending_not);
                pending_not = false;
            }
            Rule::atom => {
                neg_before_atom = pending_not;
                pending_not = false;
                atom_pair = Some(token);
            }
            other => {
                return Err(XlogError::Parse(format!(
                    "Unexpected token in nested modal chain: {:?}",
                    other
                )));
            }
        }
    }

    let atom_pair = atom_pair
        .ok_or_else(|| XlogError::Parse("Missing atom in nested modal chain".to_string()))?;
    if ops.len() < 2 {
        return Err(XlogError::Parse(
            "Nested modal chain must contain at least two epistemic operators".to_string(),
        ));
    }

    let innermost = *ops
        .last()
        .ok_or_else(|| XlogError::Parse("Empty modal chain".to_string()))?;
    let op = if neg_before_atom {
        match innermost {
            EpistemicOp::Know => EpistemicOp::Possible,
            EpistemicOp::Possible => EpistemicOp::Know,
        }
    } else {
        innermost
    };
    let negated = neg_before_op
        .iter()
        .copied()
        .fold(neg_before_atom, |acc, neg| acc ^ neg);

    Ok(BodyLiteral::Epistemic(EpistemicLiteral {
        op,
        negated,
        atom: build_atom(atom_pair)?,
    }))
}

/// Build a finite univ literal.
fn build_univ(pair: Pair<'_, Rule>) -> Result<Univ> {
    let mut inner = pair.into_inner();
    let term = build_term(
        inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing univ term".to_string()))?,
    )?;
    let parts = build_term(
        inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing univ parts".to_string()))?,
    )?;
    Ok(Univ { term, parts })
}

/// Build a comparison
fn build_comparison(pair: Pair<'_, Rule>) -> Result<Comparison> {
    let mut inner = pair.into_inner();

    let left_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing comparison left operand".to_string()))?;
    let left = build_term(left_pair)?;

    let op_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing comparison operator".to_string()))?;
    let op = match op_pair.as_str() {
        "==" | "=" => CompOp::Eq,
        "!=" => CompOp::Ne,
        "<" => CompOp::Lt,
        "<=" => CompOp::Le,
        ">" => CompOp::Gt,
        ">=" => CompOp::Ge,
        _ => {
            return Err(XlogError::Parse(format!(
                "Unknown comparison operator: {}",
                op_pair.as_str()
            )))
        }
    };

    let right_pair = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing comparison right operand".to_string()))?;
    let right = build_term(right_pair)?;

    Ok(Comparison { left, op, right })
}

/// Build a term
fn build_term(pair: Pair<'_, Rule>) -> Result<Term> {
    // term can directly contain variable, integer, etc. or be wrapped
    let inner = if pair.as_rule() == Rule::term {
        pair.into_inner()
            .next()
            .ok_or_else(|| XlogError::Parse("Empty term".to_string()))?
    } else {
        pair
    };

    match inner.as_rule() {
        Rule::var_or_anon => {
            // Unwrap var_or_anon to get either anonymous or variable
            let var_inner = inner
                .into_inner()
                .next()
                .ok_or_else(|| XlogError::Parse("Empty var_or_anon".to_string()))?;
            match var_inner.as_rule() {
                Rule::anonymous => Ok(Term::Anonymous),
                Rule::variable => Ok(Term::Variable(var_inner.as_str().to_string())),
                _ => Err(XlogError::Parse(format!(
                    "Expected variable or anonymous, got: {:?}",
                    var_inner.as_rule()
                ))),
            }
        }
        Rule::variable => Ok(Term::Variable(inner.as_str().to_string())),
        Rule::anonymous => Ok(Term::Anonymous),
        Rule::integer => {
            let val: i64 = inner
                .as_str()
                .parse()
                .map_err(|_| XlogError::Parse(format!("Invalid integer: {}", inner.as_str())))?;
            Ok(Term::Integer(val))
        }
        Rule::float_num => {
            let val: f64 = inner
                .as_str()
                .parse()
                .map_err(|_| XlogError::Parse(format!("Invalid float: {}", inner.as_str())))?;
            Ok(Term::Float(val))
        }
        Rule::string_lit => {
            let s = inner.as_str();
            // Remove quotes
            let unquoted = &s[1..s.len() - 1];
            Ok(Term::String(unquoted.to_string()))
        }
        Rule::list_literal => {
            let items = inner
                .into_inner()
                .map(build_term)
                .collect::<Result<Vec<_>>>()?;
            Ok(Term::List(items))
        }
        Rule::cons_pattern => {
            let mut parts = inner.into_inner();
            let head = parts
                .next()
                .ok_or_else(|| XlogError::Parse("Missing cons head".to_string()))?;
            let tail = parts
                .next()
                .ok_or_else(|| XlogError::Parse("Missing cons tail".to_string()))?;
            Ok(Term::Cons {
                head: Box::new(build_term(head)?),
                tail: Box::new(build_term(tail)?),
            })
        }
        Rule::compound_term => {
            let mut parts = inner.into_inner();
            let functor = parts
                .next()
                .ok_or_else(|| XlogError::Parse("Missing compound functor".to_string()))?
                .as_str()
                .to_string();
            let mut args = Vec::new();
            if let Some(term_list) = parts.next() {
                for term in term_list.into_inner() {
                    args.push(build_term(term)?);
                }
            }
            Ok(Term::Compound { functor, args })
        }
        Rule::ident => Ok(Term::Symbol(symbol::intern(inner.as_str()))),
        _ => Err(XlogError::Parse(format!(
            "Unknown term type: {:?}",
            inner.as_rule()
        ))),
    }
}

/// Build an arithmetic expression (handles additive operations + -)
/// Grammar: arith_expr = { arith_term ~ (arith_op_add ~ arith_term)* }
fn build_arith_expr(pair: Pair<'_, Rule>) -> Result<ArithExpr> {
    let mut inner = pair.into_inner();

    // First operand is always an arith_term
    let first = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Empty arithmetic expression".to_string()))?;
    let mut result = build_arith_term(first)?;

    // Process remaining (operator, operand) pairs
    while let Some(op_pair) = inner.next() {
        let op_str = op_pair.as_str();
        let right_pair = inner.next().ok_or_else(|| {
            XlogError::Parse("Missing right operand in arith expression".to_string())
        })?;
        let right = build_arith_term(right_pair)?;

        result = match op_str {
            "+" => ArithExpr::Add(Box::new(result), Box::new(right)),
            "-" => ArithExpr::Sub(Box::new(result), Box::new(right)),
            _ => {
                return Err(XlogError::Parse(format!(
                    "Unknown additive operator: {}",
                    op_str
                )))
            }
        };
    }

    Ok(result)
}

/// Build an arithmetic term (handles multiplicative operations * / %)
/// Grammar: arith_term = { arith_primary ~ (arith_op_mul ~ arith_primary)* }
fn build_arith_term(pair: Pair<'_, Rule>) -> Result<ArithExpr> {
    let mut inner = pair.into_inner();

    // First operand is always an arith_primary
    let first = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Empty arithmetic term".to_string()))?;
    let mut result = build_arith_primary(first)?;

    // Process remaining (operator, operand) pairs
    while let Some(op_pair) = inner.next() {
        let op_str = op_pair.as_str();
        let right_pair = inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing right operand in arith term".to_string()))?;
        let right = build_arith_primary(right_pair)?;

        result = match op_str {
            "*" => ArithExpr::Mul(Box::new(result), Box::new(right)),
            "/" => ArithExpr::Div(Box::new(result), Box::new(right)),
            "%" => ArithExpr::Mod(Box::new(result), Box::new(right)),
            _ => {
                return Err(XlogError::Parse(format!(
                    "Unknown multiplicative operator: {}",
                    op_str
                )))
            }
        };
    }

    Ok(result)
}

/// Build an arithmetic primary (leaf nodes, parentheses, and builtin functions)
/// Grammar: arith_primary = {
///     builtin_fn ~ "(" ~ arith_expr ~ ("," ~ (arith_expr | type_spec))* ~ ")" |
///     "(" ~ arith_expr ~ ")" |
///     variable |
///     integer |
///     float_num
/// }
fn build_arith_primary(pair: Pair<'_, Rule>) -> Result<ArithExpr> {
    let mut inner = pair.into_inner();
    let first = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Empty arithmetic primary".to_string()))?;

    match first.as_rule() {
        Rule::builtin_fn => {
            let fn_name = first.as_str();
            // Collect all arguments
            let args: Vec<Pair<'_, Rule>> = inner.collect();

            match fn_name {
                "abs" => {
                    if args.len() != 1 {
                        return Err(XlogError::Parse(
                            "abs() takes exactly 1 argument".to_string(),
                        ));
                    }
                    let arg = build_arith_expr(args.into_iter().next().unwrap())?;
                    Ok(ArithExpr::Abs(Box::new(arg)))
                }
                "min" => {
                    if args.len() != 2 {
                        return Err(XlogError::Parse(
                            "min() takes exactly 2 arguments".to_string(),
                        ));
                    }
                    let mut args_iter = args.into_iter();
                    let arg1 = build_arith_expr(args_iter.next().unwrap())?;
                    let arg2 = build_arith_expr(args_iter.next().unwrap())?;
                    Ok(ArithExpr::Min(Box::new(arg1), Box::new(arg2)))
                }
                "max" => {
                    if args.len() != 2 {
                        return Err(XlogError::Parse(
                            "max() takes exactly 2 arguments".to_string(),
                        ));
                    }
                    let mut args_iter = args.into_iter();
                    let arg1 = build_arith_expr(args_iter.next().unwrap())?;
                    let arg2 = build_arith_expr(args_iter.next().unwrap())?;
                    Ok(ArithExpr::Max(Box::new(arg1), Box::new(arg2)))
                }
                "pow" => {
                    if args.len() != 2 {
                        return Err(XlogError::Parse(
                            "pow() takes exactly 2 arguments".to_string(),
                        ));
                    }
                    let mut args_iter = args.into_iter();
                    let arg1 = build_arith_expr(args_iter.next().unwrap())?;
                    let arg2 = build_arith_expr(args_iter.next().unwrap())?;
                    Ok(ArithExpr::Pow(Box::new(arg1), Box::new(arg2)))
                }
                "cast" => {
                    if args.len() != 2 {
                        return Err(XlogError::Parse(
                            "cast() takes exactly 2 arguments".to_string(),
                        ));
                    }
                    let mut args_iter = args.into_iter();
                    let arg1 = build_arith_expr(args_iter.next().unwrap())?;
                    let type_pair = args_iter.next().unwrap();
                    let target_type = build_scalar_type_spec(type_pair, "cast target")?;
                    Ok(ArithExpr::Cast(Box::new(arg1), target_type))
                }
                _ => Err(XlogError::Parse(format!(
                    "Unknown builtin function: {}",
                    fn_name
                ))),
            }
        }
        Rule::arith_expr => {
            // Parenthesized expression
            build_arith_expr(first)
        }
        Rule::variable => Ok(ArithExpr::Variable(first.as_str().to_string())),
        Rule::integer => {
            let val: i64 = first
                .as_str()
                .parse()
                .map_err(|_| XlogError::Parse(format!("Invalid integer: {}", first.as_str())))?;
            Ok(ArithExpr::Integer(val))
        }
        Rule::float_num => {
            let val: f64 = first
                .as_str()
                .parse()
                .map_err(|_| XlogError::Parse(format!("Invalid float: {}", first.as_str())))?;
            Ok(ArithExpr::Float(val))
        }
        Rule::func_call => {
            let mut call_inner = first.into_inner();
            let name = call_inner
                .next()
                .ok_or_else(|| XlogError::Parse("Missing function name".to_string()))?
                .as_str()
                .to_string();
            let args: Vec<ArithExpr> = call_inner
                .map(build_arith_expr)
                .collect::<Result<Vec<_>>>()?;
            Ok(ArithExpr::FuncCall { name, args })
        }
        _ => Err(XlogError::Parse(format!(
            "Unexpected token in arith_primary: {:?}",
            first.as_rule()
        ))),
    }
}

/// Build an is-expression: Z is X + Y
fn build_is_expr(pair: Pair<'_, Rule>) -> Result<IsExpr> {
    let mut inner = pair.into_inner();
    let target = inner
        .next()
        .ok_or_else(|| XlogError::Parse("Missing target variable in is expression".to_string()))?
        .as_str()
        .to_string();
    let expr = build_arith_expr(
        inner
            .next()
            .ok_or_else(|| XlogError::Parse("Missing expression in is expression".to_string()))?,
    )?;
    Ok(IsExpr { target, expr })
}

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

    #[test]
    fn test_parse_fact() {
        let input = "edge(1, 2).";
        let result = parse_program(input);
        assert!(result.is_ok(), "Failed to parse fact: {:?}", result.err());

        let program = result.unwrap();
        assert_eq!(program.rules.len(), 1);
        assert!(program.rules[0].is_fact());
        assert_eq!(program.rules[0].head.predicate, "edge");
        assert_eq!(program.rules[0].head.terms.len(), 2);
        assert_eq!(program.rules[0].head.terms[0], Term::Integer(1));
        assert_eq!(program.rules[0].head.terms[1], Term::Integer(2));
    }

    #[test]
    fn test_parse_rule() {
        let input = "reach(X, Y) :- edge(X, Y).";
        let result = parse_program(input);
        assert!(result.is_ok(), "Failed to parse rule: {:?}", result.err());

        let program = result.unwrap();
        assert_eq!(program.rules.len(), 1);
        assert!(!program.rules[0].is_fact());
        assert_eq!(program.rules[0].head.predicate, "reach");
        assert_eq!(program.rules[0].body.len(), 1);
    }

    #[test]
    fn test_parse_recursive_rule() {
        let input = "reach(X, Z) :- reach(X, Y), edge(Y, Z).";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse recursive rule: {:?}",
            result.err()
        );

        let program = result.unwrap();
        assert_eq!(program.rules.len(), 1);
        assert_eq!(program.rules[0].body.len(), 2);
    }

    #[test]
    fn test_parse_negation() {
        let input = "isolated(X) :- node(X), not edge(X, Y).";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse negation: {:?}",
            result.err()
        );

        let program = result.unwrap();
        assert!(program.rules[0].has_negation());
        assert!(matches!(&program.rules[0].body[1], BodyLiteral::Negated(_)));
    }

    #[test]
    fn test_parse_epistemic_literal_syntax() {
        let input = "believed(X) :- node(X), know edge(X).";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse epistemic literal: {:?}",
            result.err()
        );

        let program = result.unwrap();
        assert_eq!(program.rules.len(), 1);
        assert_eq!(program.rules[0].body.len(), 2);
    }

    #[test]
    fn test_parse_predicate_with_epistemic_keyword_prefix_as_atom() {
        let input = "friend_of_friend(A, C) :- knows(A, B), knows(B, C), A != C.";
        let program = parse_program(input).expect("parse ordinary predicate named knows");

        assert_eq!(program.rules.len(), 1);
        assert_eq!(program.rules[0].body.len(), 3);
        assert!(
            matches!(&program.rules[0].body[0], BodyLiteral::Positive(atom) if atom.predicate == "knows")
        );
        assert!(
            matches!(&program.rules[0].body[1], BodyLiteral::Positive(atom) if atom.predicate == "knows")
        );
    }

    #[test]
    fn test_parse_aggregate() {
        let input = "out_degree(X, count(Y)) :- edge(X, Y).";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse aggregate: {:?}",
            result.err()
        );

        let program = result.unwrap();
        assert!(program.rules[0].has_aggregation());
        if let Term::Aggregate(agg) = &program.rules[0].head.terms[1] {
            assert_eq!(agg.op, AggOp::Count);
            assert_eq!(agg.variable, "Y");
        } else {
            panic!("Expected aggregate term");
        }
    }

    #[test]
    fn test_parse_logsumexp_aggregate() {
        let input = "score(X, logsumexp(Y)) :- obs(X, Y).";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse logsumexp aggregate: {:?}",
            result.err()
        );

        let program = result.unwrap();
        assert!(program.rules[0].has_aggregation());
        if let Term::Aggregate(agg) = &program.rules[0].head.terms[1] {
            assert_eq!(agg.op, AggOp::LogSumExp);
            assert_eq!(agg.variable, "Y");
        } else {
            panic!("Expected aggregate term");
        }
    }

    #[test]
    fn test_parse_constraint() {
        let input = ":- reach(X, X).";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse constraint: {:?}",
            result.err()
        );

        let program = result.unwrap();
        assert_eq!(program.constraints.len(), 1);
        assert_eq!(program.constraints[0].body.len(), 1);
    }

    #[test]
    fn test_parse_query() {
        let input = "?- reach(1, N).";
        let result = parse_program(input);
        assert!(result.is_ok(), "Failed to parse query: {:?}", result.err());

        let program = result.unwrap();
        assert_eq!(program.queries.len(), 1);
        assert_eq!(program.queries[0].atom.predicate, "reach");
    }

    #[test]
    fn test_parse_full_program() {
        let input = r#"
            edge(1, 2).
            edge(2, 3).
            edge(3, 4).
            reach(X, Y) :- edge(X, Y).
            reach(X, Z) :- reach(X, Y), edge(Y, Z).
            ?- reach(1, N).
        "#;
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse full program: {:?}",
            result.err()
        );

        let program = result.unwrap();
        assert_eq!(program.rules.len(), 5); // 3 facts + 2 rules
        assert_eq!(program.queries.len(), 1);
        assert_eq!(program.facts().count(), 3);
        assert_eq!(program.proper_rules().count(), 2);
    }

    #[test]
    fn test_parse_comparison() {
        let input = "small(X) :- value(X), X < 10.";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse comparison: {:?}",
            result.err()
        );

        let program = result.unwrap();
        assert_eq!(program.rules[0].body.len(), 2);
        if let BodyLiteral::Comparison(cmp) = &program.rules[0].body[1] {
            assert_eq!(cmp.op, CompOp::Lt);
            assert_eq!(cmp.left, Term::Variable("X".to_string()));
            assert_eq!(cmp.right, Term::Integer(10));
        } else {
            panic!("Expected comparison");
        }
    }

    #[test]
    fn test_parse_pred_decl() {
        let input = "pred edge(u32, u32).";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse pred decl: {:?}",
            result.err()
        );

        let program = result.unwrap();
        assert_eq!(program.predicates.len(), 1);
        assert_eq!(program.predicates[0].name, "edge");
        assert_eq!(program.predicates[0].types.len(), 2);
        assert_eq!(
            program.predicates[0].types[0],
            TypeRef::Scalar(ScalarType::U32)
        );
        assert_eq!(
            program.predicates[0].types[1],
            TypeRef::Scalar(ScalarType::U32)
        );
    }

    #[test]
    fn test_parse_anonymous_wildcard() {
        // Test anonymous wildcard in body
        let input = "has_child(X) :- parent(X, _).";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse anonymous wildcard: {:?}",
            result.err()
        );

        let program = result.unwrap();
        assert_eq!(program.rules.len(), 1);
        let rule = &program.rules[0];
        assert_eq!(rule.head.predicate, "has_child");

        // Check body atom has anonymous term
        if let BodyLiteral::Positive(atom) = &rule.body[0] {
            assert_eq!(atom.predicate, "parent");
            assert_eq!(atom.terms.len(), 2);
            assert_eq!(atom.terms[0], Term::Variable("X".to_string()));
            assert_eq!(atom.terms[1], Term::Anonymous);
        } else {
            panic!("Expected positive atom");
        }
    }

    #[test]
    fn test_parse_multiple_wildcards() {
        // Multiple wildcards in same rule - each is independent
        let input = "exists(X) :- rel(X, _, _).";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse multiple wildcards: {:?}",
            result.err()
        );

        let program = result.unwrap();
        if let BodyLiteral::Positive(atom) = &program.rules[0].body[0] {
            assert_eq!(atom.terms.len(), 3);
            assert_eq!(atom.terms[0], Term::Variable("X".to_string()));
            assert_eq!(atom.terms[1], Term::Anonymous);
            assert_eq!(atom.terms[2], Term::Anonymous);
        }
    }

    #[test]
    fn test_parse_is_expr() {
        // Test that grammar accepts 'is' expressions before AST lowering.
        let input = "result(X, Z) :- input(X, Y), Z is Y + 1.";
        let result = XlogParser::parse(Rule::program, input);
        assert!(
            result.is_ok(),
            "Failed to parse is expression: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_parse_arithmetic_precedence() {
        // Multiplication before addition
        let input = "r(X, Z) :- p(X, A, B), Z is A + B * 2.";
        let result = parse_program(input).unwrap();
        let rule = &result.rules[0];
        assert_eq!(rule.body.len(), 2);
        assert!(matches!(&rule.body[1], BodyLiteral::IsExpr(_)));
    }

    #[test]
    fn test_parse_arithmetic_parentheses() {
        let input = "r(X, Z) :- p(X, A, B), Z is (A + B) * 2.";
        assert!(parse_program(input).is_ok());
    }

    #[test]
    fn test_parse_probabilistic_fact_syntax() {
        let input = "0.7::rain().";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse probabilistic fact: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_parse_annotated_disjunction_syntax() {
        let input = "0.6::coin(heads); 0.4::coin(tails).";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse annotated disjunction: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_parse_evidence_is_not_a_fact() {
        let input = "evidence(rain(), true).";
        let program = parse_program(input).unwrap();
        assert_eq!(
            program.rules.len(),
            0,
            "evidence/2 should not be parsed as a regular fact"
        );
    }

    #[test]
    fn test_parse_query_directive_is_not_a_fact() {
        let input = "query(reach(1,3)).";
        let program = parse_program(input).unwrap();
        assert_eq!(
            program.rules.len(),
            0,
            "query/1 should not be parsed as a regular fact"
        );
    }

    #[test]
    fn test_parse_prob_engine_pragma_syntax() {
        let input = "#pragma prob_engine = mc";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse prob_engine pragma: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_parse_epistemic_mode_pragma_syntax() {
        let input = "#pragma epistemic_mode = faeel";
        let result = parse_program(input);
        assert!(
            result.is_ok(),
            "Failed to parse epistemic_mode pragma: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_parse_probabilistic_fact_ast() {
        let program = parse_program("0.7::rain().").unwrap();
        assert_eq!(program.prob_facts.len(), 1);
        assert!((program.prob_facts[0].prob - 0.7).abs() < 1e-9);
        assert_eq!(program.prob_facts[0].atom.predicate, "rain");
        assert!(program.prob_facts[0].atom.terms.is_empty());
    }

    #[test]
    fn test_parse_annotated_disjunction_ast() {
        let program = parse_program("0.6::coin(heads); 0.4::coin(tails).").unwrap();
        assert_eq!(program.annotated_disjunctions.len(), 1);
        let ad = &program.annotated_disjunctions[0];
        assert_eq!(ad.choices.len(), 2);
        assert!((ad.choices[0].prob - 0.6).abs() < 1e-9);
        assert_eq!(ad.choices[0].atom.predicate, "coin");
        assert_eq!(ad.choices[0].atom.terms.len(), 1);
        assert_eq!(
            ad.choices[0].atom.terms[0],
            Term::Symbol(symbol::intern("heads"))
        );
        assert!((ad.choices[1].prob - 0.4).abs() < 1e-9);
        assert_eq!(
            ad.choices[1].atom.terms[0],
            Term::Symbol(symbol::intern("tails"))
        );
    }

    #[test]
    fn test_parse_evidence_ast() {
        let program = parse_program("evidence(rain(), true).").unwrap();
        assert_eq!(program.evidence.len(), 1);
        assert_eq!(program.evidence[0].atom.predicate, "rain");
        assert!(program.evidence[0].value);
    }

    #[test]
    fn test_parse_prob_query_ast() {
        let program = parse_program("query(reach(1,3)).").unwrap();
        assert_eq!(program.prob_queries.len(), 1);
        assert_eq!(program.prob_queries[0].atom.predicate, "reach");
        assert_eq!(program.prob_queries[0].atom.terms.len(), 2);
        assert_eq!(program.prob_queries[0].atom.terms[0], Term::Integer(1));
        assert_eq!(program.prob_queries[0].atom.terms[1], Term::Integer(3));
    }

    #[test]
    fn test_parse_prob_engine_pragma_ast() {
        let program = parse_program("#pragma prob_engine = mc").unwrap();
        assert_eq!(
            program.directives.prob_engine,
            Some(crate::ast::ProbEngine::Mc)
        );
        assert_eq!(program.prob_engine(), crate::ast::ProbEngine::Mc);
    }

    #[test]
    fn test_parse_arithmetic_builtins() {
        let inputs = [
            "r(X, Z) :- p(X, Y), Z is abs(Y).",
            "r(X, Z) :- p(X, A, B), Z is min(A, B).",
            "r(X, Z) :- p(X, A, B), Z is max(A, B).",
            "r(X, Z) :- p(X, A, B), Z is pow(A, B).",
            "r(X, Z) :- p(X, Y), Z is cast(Y, f64).",
        ];
        for input in inputs {
            assert!(parse_program(input).is_ok(), "Failed to parse: {}", input);
        }
    }

    #[test]
    fn test_parse_arithmetic_nested() {
        let input = "r(X, Z) :- p(X, A, B, C), Z is abs(A - B) + min(B, C) * 2.";
        assert!(parse_program(input).is_ok());
    }

    /// Parse a one-rule program whose single body literal is an epistemic literal
    /// and return that literal.
    fn first_epistemic_literal(src: &str) -> EpistemicLiteral {
        let program = parse_program(src).unwrap_or_else(|e| panic!("parse failed: {e:?}"));
        let rule = program.rules.first().expect("one rule");
        match rule.body.first().expect("one body literal") {
            BodyLiteral::Epistemic(lit) => lit.clone(),
            other => panic!("expected epistemic literal, got {other:?}"),
        }
    }

    #[test]
    fn test_nested_modal_chain_collapses_to_innermost_operator() {
        // know possible p ≡ possible p  (KM ≡ M, inner operator wins)
        let lit = first_epistemic_literal("q() :- know possible p().");
        assert_eq!(lit.op, EpistemicOp::Possible);
        assert!(!lit.negated);
        assert_eq!(lit.atom.predicate, "p");

        // possible know p ≡ know p  (MK ≡ K)
        let lit = first_epistemic_literal("q() :- possible know p().");
        assert_eq!(lit.op, EpistemicOp::Know);
        assert!(!lit.negated);

        // know know p ≡ know p  (KK ≡ K)
        let lit = first_epistemic_literal("q() :- know know p().");
        assert_eq!(lit.op, EpistemicOp::Know);

        // possible possible p ≡ possible p  (MM ≡ M)
        let lit = first_epistemic_literal("q() :- possible possible p().");
        assert_eq!(lit.op, EpistemicOp::Possible);

        // 3-deep chain: innermost (atom-adjacent) operator still wins.
        let lit = first_epistemic_literal("q() :- know possible know p().");
        assert_eq!(lit.op, EpistemicOp::Know);
    }

    #[test]
    fn test_nested_modal_chain_leading_negation_distributes() {
        // not know possible p ≡ not possible p
        let lit = first_epistemic_literal("q() :- not know possible p().");
        assert_eq!(lit.op, EpistemicOp::Possible);
        assert!(lit.negated);

        // not possible know p ≡ not know p
        let lit = first_epistemic_literal("q() :- not possible know p().");
        assert_eq!(lit.op, EpistemicOp::Know);
        assert!(lit.negated);
    }

    #[test]
    fn test_nested_modal_chain_interior_negation_dualizes() {
        // know not possible p ≡ not possible p
        let lit = first_epistemic_literal("q() :- know not possible p().");
        assert_eq!(lit.op, EpistemicOp::Possible);
        assert!(lit.negated);

        // possible not know p ≡ not know p
        let lit = first_epistemic_literal("q() :- possible not know p().");
        assert_eq!(lit.op, EpistemicOp::Know);
        assert!(lit.negated);

        // not know not possible p ≡ possible p
        let lit = first_epistemic_literal("q() :- not know not possible p().");
        assert_eq!(lit.op, EpistemicOp::Possible);
        assert!(!lit.negated);
    }

    #[test]
    fn test_nested_modal_chain_atom_adjacent_negation_dualizes() {
        // know possible not p ≡ not know p
        let lit = first_epistemic_literal("q() :- know possible not p().");
        assert_eq!(lit.op, EpistemicOp::Know);
        assert!(lit.negated);

        // possible know not p ≡ not possible p
        let lit = first_epistemic_literal("q() :- possible know not p().");
        assert_eq!(lit.op, EpistemicOp::Possible);
        assert!(lit.negated);
    }
}