vidi/util/

solver.rs

/* Copyright (C) 2025 DorotaC
 * SPDX-License-Identifier: MPL-2.0 OR LGPL-2.1-or-later OR MPL-2.0
 */
 
/*! Common routines for the Prolog solver. */

use logru_arithmetic::logru;
use logru::Symbols;
use logru::ast::{Term, VarScope};
use logru::search::Solution;
use logru::textual::UniverseQuery;
use std::array;
use std::collections::HashMap;

/** The standard library of Prolog predicates.
 *
 * The current Prolog syntax doesn't accept "=", ",", and other goodies in queries, so this makes them possible:
 * 
 * - `A = config(_, _).` becomes `A, eq(A, config(_, _))`.
 */
pub const STDLIB: &'static str = "
% ===== Standard library

eq(A, A).

% All the following are susceptible to going out of bounds on i64 values.

isLess(fraction(Fn, Fd), fraction(Gn, Gd)) :- isLess(mul(Fn, Gd), mul(Gn, Fd)).
is(fraction(Fn, Fd), fraction(Gn, Gd)) :- is(mul(Fn, Gd), mul(Gn, Fd)).
isGreater(fraction(Fn, Fd), fraction(Gn, Gd)) :- isGreater(mul(Fn, Gd), mul(Gn, Fd)).
isLessEq(fraction(Fn, Fd), fraction(Gn, Gd)) :- isLessEq(mul(Fn, Gd), mul(Gn, Fd)).
isGreaterEq(fraction(Fn, Fd), fraction(Gn, Gd)) :- isGreaterEq(mul(Fn, Gd), mul(Gn, Fd)).

% list. Empty list is `nil`
% l(Elem, Next).

% Returns the last element of the list (the one that is followed by `nil`)
last(l(Elem, nil), Elem).
last(l(_, Elems), Elem) :- last(Elems, Elem).

% append a single element to list
append(Es, Elem, l(Elem, Es)).
";

pub type SolutionHash = HashMap<String, Option<Term>>;

pub fn solution_to_hash(names: &VarScope, solution: Solution) -> SolutionHash {
    solution.iter_vars()
        .filter_map(|(var, term)| names
            .get_name(var)
            .and_then(|name| if name.ends_with("_") {
                None
            } else {
                Some(name)
            })
            .map(|n| (n.into(), term.cloned()))
        )
        .collect::<SolutionHash>()
}

/// Returns a native integer
pub fn get_int(term: &Term) -> Option<i64> {
    match term {
        Term::Int(i) => Some(*i),
        _ => None
    }
}

pub fn get_u32(term: &logru::ast::Term) -> Option<u32> {
    match term {
        logru::ast::Term::Int(val) => u32::try_from(*val).ok(),
        _ => None,
    }
}

pub fn get_u16(term: &logru::ast::Term) -> Option<u16> {
    match term {
        logru::ast::Term::Int(val) => u16::try_from(*val).ok(),
        _ => None,
    }
}

/// Extracts an atom (an Identifier term with no arguments).
pub fn get_atom<'a>(query: &'a UniverseQuery<'a>, term: &logru::ast::Term) -> Option<&'a str> {
    match term {
        logru::ast::Term::App(term) => {
            if term.args.len() == 0 {
                query.symbols().get_symbol_name(term.functor)
            } else {
                None
            }
        },
        _ => None,
    }
}

/// Returns the name of the `term` and its arguments, or None if the name doesn't exist in this query.
pub fn get_app_term<'a, 'b>(query: &'a UniverseQuery<'a>, term: &'b Term)
    -> Option<(&'a str, &'b [Term])>
{
    match term {
        Term::App(term) => query.symbols()
            .get_symbol_name(term.functor)
            .map(|n| (n, term.args.as_slice())),
        _ => None,
    }
}

/// Returns the exact number of arguments of term `term` if its name matches `name`.
pub fn get_term_args<'a, 'b, const N: usize>(
    query: &'a UniverseQuery<'a>,
    term: &'b Term,
    name: &str,
) -> Option<[&'b Term; N]> {
    let (n, args) = get_app_term(query, term)?;
    if n == name && args.len() == N{
        let mut args = args.iter();
        Some(array::from_fn(|_| args.next().unwrap()))
    } else {
        None
    }
}

/// Converts a list built with `l/2` to a `Vec`.
/// Fails if any of the elements are not a list.
pub fn list_to_vec<'a, 'b>(
    query: &'a UniverseQuery<'a>,
    mut list_term: &'b Term,
) -> Option<Vec<&'b Term>> {
    let mut out = Vec::new();
    loop {
        match get_app_term(query, list_term)? {
            ("nil", args) if args.len() == 0 => {
                return Some(out);
            },
            ("l", args) if args.len() == 2 => {
                out.push(&args[0]);
                list_term = &args[1];
            },
            _ => {
                return None;
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use logru::query_dfs;
    use logru::resolve::ResolverExt;
    use logru::textual::TextualUniverse;
    use logru_arithmetic::ArithmeticResolver;

    #[test]
    fn check_list() {
        let mut universe = TextualUniverse::new();
        universe.load_str(STDLIB).unwrap();
        universe.load_str("l(a, l(b, l(c, nil))).").unwrap();

        let mut query = universe.prepare_query(
            "eq(A, l(_, _)), A."
        ).unwrap();
        let resolver = ArithmeticResolver::new(query.symbols_mut()).or_else(universe.resolver());
        let names = query.query().scope.as_ref().unwrap();
        let solutions = query_dfs(resolver, &query.query());
        let mut solutions = solutions.map(|solution| solution_to_hash(names, solution));
        let solution = solutions.next().unwrap();
        
        let terms = list_to_vec(
            &query,
            solution.get("A").unwrap().as_ref().unwrap()
        ).unwrap();
        let names = terms.into_iter()
            .map(|term| get_app_term(&query, term).unwrap())
            .collect::<Vec<_>>();
        assert_eq!(
            names,
            vec![("a", [].as_slice()), ("b", &[]), ("c", &[])],
        );
        
        assert_eq!(solutions.next(), None);
    }
}