conjure_cp_core/ast/

eval.rs

#![allow(dead_code)]
use crate::ast::{
    AbstractLiteral, Atom, DeclarationKind, Expression as Expr, Literal as Lit, Metadata,
    comprehension::{Comprehension, ComprehensionQualifier},
    matrix,
};
use crate::into_matrix;
use itertools::{Itertools as _, izip};
use std::cmp::Ordering as CmpOrdering;
use std::collections::HashSet;

/// Simplify an expression to a constant if possible
/// Returns:
/// `None` if the expression cannot be simplified to a constant (e.g. if it contains a variable)
/// `Some(Const)` if the expression can be simplified to a constant
pub fn eval_constant(expr: &Expr) -> Option<Lit> {
    match expr {
        Expr::Supset(_, a, b) => match (a.as_ref(), b.as_ref()) {
            (
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
            ) => {
                let a_set: HashSet<Lit> = a.iter().cloned().collect();
                let b_set: HashSet<Lit> = b.iter().cloned().collect();

                if a_set.difference(&b_set).count() > 0 {
                    Some(Lit::Bool(a_set.is_superset(&b_set)))
                } else {
                    Some(Lit::Bool(false))
                }
            }
            _ => None,
        },
        Expr::SupsetEq(_, a, b) => match (a.as_ref(), b.as_ref()) {
            (
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
            ) => Some(Lit::Bool(
                a.iter()
                    .cloned()
                    .collect::<HashSet<Lit>>()
                    .is_superset(&b.iter().cloned().collect::<HashSet<Lit>>()),
            )),
            _ => None,
        },
        Expr::Subset(_, a, b) => match (a.as_ref(), b.as_ref()) {
            (
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
            ) => {
                let a_set: HashSet<Lit> = a.iter().cloned().collect();
                let b_set: HashSet<Lit> = b.iter().cloned().collect();

                if b_set.difference(&a_set).count() > 0 {
                    Some(Lit::Bool(a_set.is_subset(&b_set)))
                } else {
                    Some(Lit::Bool(false))
                }
            }
            _ => None,
        },
        Expr::SubsetEq(_, a, b) => match (a.as_ref(), b.as_ref()) {
            (
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
            ) => Some(Lit::Bool(
                a.iter()
                    .cloned()
                    .collect::<HashSet<Lit>>()
                    .is_subset(&b.iter().cloned().collect::<HashSet<Lit>>()),
            )),
            _ => None,
        },
        Expr::Intersect(_, a, b) => match (a.as_ref(), b.as_ref()) {
            (
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
            ) => {
                let mut res: Vec<Lit> = Vec::new();
                for lit in a.iter() {
                    if b.contains(lit) && !res.contains(lit) {
                        res.push(lit.clone());
                    }
                }
                Some(Lit::AbstractLiteral(AbstractLiteral::Set(res)))
            }
            _ => None,
        },
        Expr::Union(_, a, b) => match (a.as_ref(), b.as_ref()) {
            (
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(a)))),
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(b)))),
            ) => {
                let mut res: Vec<Lit> = Vec::new();
                for lit in a.iter() {
                    res.push(lit.clone());
                }
                for lit in b.iter() {
                    if !res.contains(lit) {
                        res.push(lit.clone());
                    }
                }
                Some(Lit::AbstractLiteral(AbstractLiteral::Set(res)))
            }
            _ => None,
        },
        Expr::In(_, a, b) => {
            if let (
                Expr::Atomic(_, Atom::Literal(Lit::Int(c))),
                Expr::Atomic(_, Atom::Literal(Lit::AbstractLiteral(AbstractLiteral::Set(d)))),
            ) = (a.as_ref(), b.as_ref())
            {
                for lit in d.iter() {
                    if let Lit::Int(x) = lit
                        && c == x
                    {
                        return Some(Lit::Bool(true));
                    }
                }
                Some(Lit::Bool(false))
            } else {
                None
            }
        }
        Expr::FromSolution(_, _) => None,
        Expr::DominanceRelation(_, _) => None,
        Expr::InDomain(_, e, domain) => {
            let Expr::Atomic(_, Atom::Literal(lit)) = e.as_ref() else {
                return None;
            };

            domain.contains(lit).ok().map(Into::into)
        }
        Expr::Atomic(_, Atom::Literal(c)) => Some(c.clone()),
        Expr::Atomic(_, Atom::Reference(reference)) => reference.resolve_constant(),
        Expr::AbstractLiteral(_, a) => Some(Lit::AbstractLiteral(a.clone().into_literals()?)),
        Expr::Comprehension(_, comprehension) => {
            eval_constant_comprehension(comprehension.as_ref())
        }
        Expr::AbstractComprehension(_, _) => None,
        Expr::UnsafeIndex(_, subject, indices) | Expr::SafeIndex(_, subject, indices) => {
            let subject: Lit = eval_constant(subject.as_ref())?;
            let indices: Vec<Lit> = indices
                .iter()
                .map(eval_constant)
                .collect::<Option<Vec<Lit>>>()?;

            match subject {
                Lit::AbstractLiteral(subject @ AbstractLiteral::Matrix(_, _)) => {
                    matrix::flatten_enumerate(subject)
                        .find(|(i, _)| i == &indices)
                        .map(|(_, x)| x)
                }
                Lit::AbstractLiteral(subject @ AbstractLiteral::Tuple(_)) => {
                    let AbstractLiteral::Tuple(elems) = subject else {
                        return None;
                    };

                    assert!(indices.len() == 1, "nested tuples not supported yet");

                    let Lit::Int(index) = indices[0].clone() else {
                        return None;
                    };

                    if elems.len() < index as usize || index < 1 {
                        return None;
                    }

                    // -1 for 0-indexing vs 1-indexing
                    let item = elems[index as usize - 1].clone();

                    Some(item)
                }
                Lit::AbstractLiteral(subject @ AbstractLiteral::Record(_)) => {
                    let AbstractLiteral::Record(elems) = subject else {
                        return None;
                    };

                    assert!(indices.len() == 1, "nested record not supported yet");

                    let Lit::Int(index) = indices[0].clone() else {
                        return None;
                    };

                    if elems.len() < index as usize || index < 1 {
                        return None;
                    }

                    // -1 for 0-indexing vs 1-indexing
                    let item = elems[index as usize - 1].clone();
                    Some(item.value)
                }
                _ => None,
            }
        }
        Expr::UnsafeSlice(_, subject, indices) | Expr::SafeSlice(_, subject, indices) => {
            let subject: Lit = eval_constant(subject.as_ref())?;
            let Lit::AbstractLiteral(subject @ AbstractLiteral::Matrix(_, _)) = subject else {
                return None;
            };

            let hole_dim = indices
                .iter()
                .cloned()
                .position(|x| x.is_none())
                .expect("slice expression should have a hole dimension");

            let missing_domain = matrix::index_domains(subject.clone())[hole_dim].clone();

            let indices: Vec<Option<Lit>> = indices
                .iter()
                .cloned()
                .map(|x| {
                    // the outer option represents success of this iterator, the inner the index
                    // slice.
                    match x {
                        Some(x) => eval_constant(&x).map(Some),
                        None => Some(None),
                    }
                })
                .collect::<Option<Vec<Option<Lit>>>>()?;

            let indices_in_slice: Vec<Vec<Lit>> = missing_domain
                .values()
                .ok()?
                .map(|i| {
                    let mut indices = indices.clone();
                    indices[hole_dim] = Some(i);
                    // These unwraps will only fail if we have multiple holes.
                    // As this is invalid, panicking is fine.
                    indices.into_iter().map(|x| x.unwrap()).collect_vec()
                })
                .collect_vec();

            // Note: indices_in_slice is not necessarily sorted, so this is the best way.
            let elems = matrix::flatten_enumerate(subject)
                .filter(|(i, _)| indices_in_slice.contains(i))
                .map(|(_, elem)| elem)
                .collect();

            Some(Lit::AbstractLiteral(into_matrix![elems]))
        }
        Expr::Abs(_, e) => un_op::<i32, i32>(|a| a.abs(), e).map(Lit::Int),
        Expr::Eq(_, a, b) => bin_op::<i32, bool>(|a, b| a == b, a, b)
            .or_else(|| bin_op::<bool, bool>(|a, b| a == b, a, b))
            .map(Lit::Bool),
        Expr::Neq(_, a, b) => bin_op::<i32, bool>(|a, b| a != b, a, b).map(Lit::Bool),
        Expr::Lt(_, a, b) => bin_op::<i32, bool>(|a, b| a < b, a, b).map(Lit::Bool),
        Expr::Gt(_, a, b) => bin_op::<i32, bool>(|a, b| a > b, a, b).map(Lit::Bool),
        Expr::Leq(_, a, b) => bin_op::<i32, bool>(|a, b| a <= b, a, b).map(Lit::Bool),
        Expr::Geq(_, a, b) => bin_op::<i32, bool>(|a, b| a >= b, a, b).map(Lit::Bool),
        Expr::Not(_, expr) => un_op::<bool, bool>(|e| !e, expr).map(Lit::Bool),
        Expr::And(_, e) => {
            vec_lit_op::<bool, bool>(|e| e.iter().all(|&e| e), e.as_ref()).map(Lit::Bool)
        }
        Expr::Table(_, _, _) => None,
        Expr::NegativeTable(_, _, _) => None,
        Expr::Root(_, _) => None,
        Expr::Or(_, es) => {
            // possibly cheating; definitely should be in partial eval instead
            for e in (**es).clone().unwrap_list()? {
                if let Expr::Atomic(_, Atom::Literal(Lit::Bool(true))) = e {
                    return Some(Lit::Bool(true));
                };
            }

            vec_lit_op::<bool, bool>(|e| e.iter().any(|&e| e), es.as_ref()).map(Lit::Bool)
        }
        Expr::Imply(_, box1, box2) => {
            let a: &Atom = (&**box1).try_into().ok()?;
            let b: &Atom = (&**box2).try_into().ok()?;

            let a: bool = a.try_into().ok()?;
            let b: bool = b.try_into().ok()?;

            if a {
                // true -> b ~> b
                Some(Lit::Bool(b))
            } else {
                // false -> b ~> true
                Some(Lit::Bool(true))
            }
        }
        Expr::Iff(_, box1, box2) => {
            let a: &Atom = (&**box1).try_into().ok()?;
            let b: &Atom = (&**box2).try_into().ok()?;

            let a: bool = a.try_into().ok()?;
            let b: bool = b.try_into().ok()?;

            Some(Lit::Bool(a == b))
        }
        Expr::Sum(_, exprs) => vec_lit_op::<i32, i32>(|e| e.iter().sum(), exprs).map(Lit::Int),
        Expr::Product(_, exprs) => {
            vec_lit_op::<i32, i32>(|e| e.iter().product(), exprs).map(Lit::Int)
        }
        Expr::FlatIneq(_, a, b, c) => {
            let a: i32 = a.try_into().ok()?;
            let b: i32 = b.try_into().ok()?;
            let c: i32 = c.try_into().ok()?;

            Some(Lit::Bool(a <= b + c))
        }
        Expr::FlatSumGeq(_, exprs, a) => {
            let sum = exprs.iter().try_fold(0, |acc, atom: &Atom| {
                let n: i32 = atom.try_into().ok()?;
                let acc = acc + n;
                Some(acc)
            })?;

            Some(Lit::Bool(sum >= a.try_into().ok()?))
        }
        Expr::FlatSumLeq(_, exprs, a) => {
            let sum = exprs.iter().try_fold(0, |acc, atom: &Atom| {
                let n: i32 = atom.try_into().ok()?;
                let acc = acc + n;
                Some(acc)
            })?;

            Some(Lit::Bool(sum >= a.try_into().ok()?))
        }
        Expr::Min(_, e) => {
            opt_vec_lit_op::<i32, i32>(|e| e.iter().min().copied(), e.as_ref()).map(Lit::Int)
        }
        Expr::Max(_, e) => {
            opt_vec_lit_op::<i32, i32>(|e| e.iter().max().copied(), e.as_ref()).map(Lit::Int)
        }
        Expr::UnsafeDiv(_, a, b) | Expr::SafeDiv(_, a, b) => {
            if unwrap_expr::<i32>(b)? == 0 {
                return None;
            }
            bin_op::<i32, i32>(|a, b| ((a as f32) / (b as f32)).floor() as i32, a, b).map(Lit::Int)
        }
        Expr::UnsafeMod(_, a, b) | Expr::SafeMod(_, a, b) => {
            if unwrap_expr::<i32>(b)? == 0 {
                return None;
            }
            bin_op::<i32, i32>(|a, b| a - b * (a as f32 / b as f32).floor() as i32, a, b)
                .map(Lit::Int)
        }
        Expr::MinionDivEqUndefZero(_, a, b, c) => {
            // div always rounds down
            let a: i32 = a.try_into().ok()?;
            let b: i32 = b.try_into().ok()?;
            let c: i32 = c.try_into().ok()?;

            if b == 0 {
                return None;
            }

            let a = a as f32;
            let b = b as f32;
            let div: i32 = (a / b).floor() as i32;
            Some(Lit::Bool(div == c))
        }
        Expr::Bubble(_, a, b) => bin_op::<bool, bool>(|a, b| a && b, a, b).map(Lit::Bool),
        Expr::MinionReify(_, a, b) => {
            let result = eval_constant(a)?;

            let result: bool = result.try_into().ok()?;
            let b: bool = b.try_into().ok()?;

            Some(Lit::Bool(b == result))
        }
        Expr::MinionReifyImply(_, a, b) => {
            let result = eval_constant(a)?;

            let result: bool = result.try_into().ok()?;
            let b: bool = b.try_into().ok()?;

            if b {
                Some(Lit::Bool(result))
            } else {
                Some(Lit::Bool(true))
            }
        }
        Expr::MinionModuloEqUndefZero(_, a, b, c) => {
            // From Savile Row. Same semantics as division.
            //
            //   a - (b * floor(a/b))
            //
            // We don't use % as it has the same semantics as /. We don't use / as we want to round
            // down instead, not towards zero.

            let a: i32 = a.try_into().ok()?;
            let b: i32 = b.try_into().ok()?;
            let c: i32 = c.try_into().ok()?;

            if b == 0 {
                return None;
            }

            let modulo = a - b * (a as f32 / b as f32).floor() as i32;
            Some(Lit::Bool(modulo == c))
        }
        Expr::MinionPow(_, a, b, c) => {
            // only available for positive a b c

            let a: i32 = a.try_into().ok()?;
            let b: i32 = b.try_into().ok()?;
            let c: i32 = c.try_into().ok()?;

            if a <= 0 {
                return None;
            }

            if b <= 0 {
                return None;
            }

            if c <= 0 {
                return None;
            }

            Some(Lit::Bool(a ^ b == c))
        }
        Expr::MinionWInSet(_, _, _) => None,
        Expr::MinionWInIntervalSet(_, x, intervals) => {
            let x_lit: &Lit = x.try_into().ok()?;

            let x_lit = match x_lit.clone() {
                Lit::Int(i) => Some(i),
                Lit::Bool(true) => Some(1),
                Lit::Bool(false) => Some(0),
                _ => None,
            }?;

            let mut intervals = intervals.iter();
            loop {
                let Some(lower) = intervals.next() else {
                    break;
                };

                let Some(upper) = intervals.next() else {
                    break;
                };
                if &x_lit >= lower && &x_lit <= upper {
                    return Some(Lit::Bool(true));
                }
            }

            Some(Lit::Bool(false))
        }
        Expr::Flatten(_, _, _) => {
            // TODO
            None
        }
        Expr::AllDiff(_, e) => {
            let es = (**e).clone().unwrap_list()?;
            let mut lits: HashSet<Lit> = HashSet::new();
            for expr in es {
                let Expr::Atomic(_, Atom::Literal(x)) = expr else {
                    return None;
                };
                match x {
                    Lit::Int(_) | Lit::Bool(_) => {
                        if lits.contains(&x) {
                            return Some(Lit::Bool(false));
                        } else {
                            lits.insert(x.clone());
                        }
                    }
                    Lit::AbstractLiteral(_) => return None, // Reject AbstractLiteral cases
                }
            }
            Some(Lit::Bool(true))
        }
        Expr::FlatAllDiff(_, es) => {
            let mut lits: HashSet<Lit> = HashSet::new();
            for atom in es {
                let Atom::Literal(x) = atom else {
                    return None;
                };

                match x {
                    Lit::Int(_) | Lit::Bool(_) => {
                        if lits.contains(x) {
                            return Some(Lit::Bool(false));
                        } else {
                            lits.insert(x.clone());
                        }
                    }
                    Lit::AbstractLiteral(_) => return None, // Reject AbstractLiteral cases
                }
            }
            Some(Lit::Bool(true))
        }
        Expr::FlatWatchedLiteral(_, _, _) => None,
        Expr::AuxDeclaration(_, _, _) => None,
        Expr::Neg(_, a) => match eval_constant(a.as_ref())? {
            Lit::Int(a) => Some(Lit::Int(-a)),
            _ => None,
        },
        Expr::Factorial(_, _) => None,
        Expr::Minus(_, a, b) => bin_op::<i32, i32>(|a, b| a - b, a, b).map(Lit::Int),
        Expr::FlatMinusEq(_, a, b) => {
            let a: i32 = a.try_into().ok()?;
            let b: i32 = b.try_into().ok()?;
            Some(Lit::Bool(a == -b))
        }
        Expr::FlatProductEq(_, a, b, c) => {
            let a: i32 = a.try_into().ok()?;
            let b: i32 = b.try_into().ok()?;
            let c: i32 = c.try_into().ok()?;
            Some(Lit::Bool(a * b == c))
        }
        Expr::FlatWeightedSumLeq(_, cs, vs, total) => {
            let cs: Vec<i32> = cs
                .iter()
                .map(|x| TryInto::<i32>::try_into(x).ok())
                .collect::<Option<Vec<i32>>>()?;
            let vs: Vec<i32> = vs
                .iter()
                .map(|x| TryInto::<i32>::try_into(x).ok())
                .collect::<Option<Vec<i32>>>()?;
            let total: i32 = total.try_into().ok()?;

            let sum: i32 = izip!(cs, vs).fold(0, |acc, (c, v)| acc + (c * v));

            Some(Lit::Bool(sum <= total))
        }
        Expr::FlatWeightedSumGeq(_, cs, vs, total) => {
            let cs: Vec<i32> = cs
                .iter()
                .map(|x| TryInto::<i32>::try_into(x).ok())
                .collect::<Option<Vec<i32>>>()?;
            let vs: Vec<i32> = vs
                .iter()
                .map(|x| TryInto::<i32>::try_into(x).ok())
                .collect::<Option<Vec<i32>>>()?;
            let total: i32 = total.try_into().ok()?;

            let sum: i32 = izip!(cs, vs).fold(0, |acc, (c, v)| acc + (c * v));

            Some(Lit::Bool(sum >= total))
        }
        Expr::FlatAbsEq(_, x, y) => {
            let x: i32 = x.try_into().ok()?;
            let y: i32 = y.try_into().ok()?;

            Some(Lit::Bool(x == y.abs()))
        }
        Expr::UnsafePow(_, a, b) | Expr::SafePow(_, a, b) => {
            let a: &Atom = a.try_into().ok()?;
            let a: i32 = a.try_into().ok()?;

            let b: &Atom = b.try_into().ok()?;
            let b: i32 = b.try_into().ok()?;

            if (a != 0 || b != 0) && b >= 0 {
                Some(Lit::Int(a.pow(b as u32)))
            } else {
                None
            }
        }
        Expr::Metavar(_, _) => None,
        Expr::MinionElementOne(_, _, _, _) => None,
        Expr::ToInt(_, expression) => {
            let lit = eval_constant(expression.as_ref())?;
            match lit {
                Lit::Int(_) => Some(lit),
                Lit::Bool(true) => Some(Lit::Int(1)),
                Lit::Bool(false) => Some(Lit::Int(0)),
                _ => None,
            }
        }
        Expr::SATInt(_, _, _, _) => {
            // TODO: If this SATInt is composed of literals, we should evaluate it back to an
            // integer literal.
            //
            // This is important because `is_all_constant` currently returns true for SATInts
            // containing no references. If we don't evaluate them here, bubble rules will skip
            // them (thinking they'll be constant-folded later), but they'll actually reach
            // the solver adaptors as un-encoded unsafe operations, causing panics.
            None
        }
        Expr::PairwiseSum(_, a, b) => {
            match (eval_constant(a.as_ref())?, eval_constant(b.as_ref())?) {
                (Lit::Int(a_int), Lit::Int(b_int)) => Some(Lit::Int(a_int + b_int)),
                _ => None,
            }
        }
        Expr::PairwiseProduct(_, a, b) => {
            match (eval_constant(a.as_ref())?, eval_constant(b.as_ref())?) {
                (Lit::Int(a_int), Lit::Int(b_int)) => Some(Lit::Int(a_int * b_int)),
                _ => None,
            }
        }
        Expr::Defined(_, _) => todo!(),
        Expr::Range(_, _) => todo!(),
        Expr::Image(_, _, _) => todo!(),
        Expr::ImageSet(_, _, _) => todo!(),
        Expr::PreImage(_, _, _) => todo!(),
        Expr::Inverse(_, _, _) => todo!(),
        Expr::Restrict(_, _, _) => todo!(),
        Expr::LexLt(_, a, b) => {
            let lt = vec_expr_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
                pairs
                    .iter()
                    .find_map(|(a, b)| match a.cmp(b) {
                        CmpOrdering::Less => Some(true),     // First difference is <
                        CmpOrdering::Greater => Some(false), // First difference is >
                        CmpOrdering::Equal => None,          // No difference
                    })
                    .unwrap_or(a_len < b_len) // [1,1] <lex [1,1,x]
            })?;
            Some(lt.into())
        }
        Expr::LexLeq(_, a, b) => {
            let lt = vec_expr_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
                pairs
                    .iter()
                    .find_map(|(a, b)| match a.cmp(b) {
                        CmpOrdering::Less => Some(true),
                        CmpOrdering::Greater => Some(false),
                        CmpOrdering::Equal => None,
                    })
                    .unwrap_or(a_len <= b_len) // [1,1] <=lex [1,1,x]
            })?;
            Some(lt.into())
        }
        Expr::LexGt(_, a, b) => eval_constant(&Expr::LexLt(Metadata::new(), b.clone(), a.clone())),
        Expr::LexGeq(_, a, b) => {
            eval_constant(&Expr::LexLeq(Metadata::new(), b.clone(), a.clone()))
        }
        Expr::FlatLexLt(_, a, b) => {
            let lt = atoms_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
                pairs
                    .iter()
                    .find_map(|(a, b)| match a.cmp(b) {
                        CmpOrdering::Less => Some(true),
                        CmpOrdering::Greater => Some(false),
                        CmpOrdering::Equal => None,
                    })
                    .unwrap_or(a_len < b_len)
            })?;
            Some(lt.into())
        }
        Expr::FlatLexLeq(_, a, b) => {
            let lt = atoms_pairs_op::<i32, _>(a, b, |pairs, (a_len, b_len)| {
                pairs
                    .iter()
                    .find_map(|(a, b)| match a.cmp(b) {
                        CmpOrdering::Less => Some(true),
                        CmpOrdering::Greater => Some(false),
                        CmpOrdering::Equal => None,
                    })
                    .unwrap_or(a_len <= b_len)
            })?;
            Some(lt.into())
        }
    }
}

pub fn un_op<T, A>(f: fn(T) -> A, a: &Expr) -> Option<A>
where
    T: TryFrom<Lit>,
{
    let a = unwrap_expr::<T>(a)?;
    Some(f(a))
}

pub fn bin_op<T, A>(f: fn(T, T) -> A, a: &Expr, b: &Expr) -> Option<A>
where
    T: TryFrom<Lit>,
{
    let a = unwrap_expr::<T>(a)?;
    let b = unwrap_expr::<T>(b)?;
    Some(f(a, b))
}

#[allow(dead_code)]
pub fn tern_op<T, A>(f: fn(T, T, T) -> A, a: &Expr, b: &Expr, c: &Expr) -> Option<A>
where
    T: TryFrom<Lit>,
{
    let a = unwrap_expr::<T>(a)?;
    let b = unwrap_expr::<T>(b)?;
    let c = unwrap_expr::<T>(c)?;
    Some(f(a, b, c))
}

pub fn vec_op<T, A>(f: fn(Vec<T>) -> A, a: &[Expr]) -> Option<A>
where
    T: TryFrom<Lit>,
{
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
    Some(f(a))
}

pub fn vec_lit_op<T, A>(f: fn(Vec<T>) -> A, a: &Expr) -> Option<A>
where
    T: TryFrom<Lit>,
{
    Some(f(eval_list_items(a)?))
}

type PairsCallback<T, A> = fn(Vec<(T, T)>, (usize, usize)) -> A;

/// Calls the given function on each consecutive pair of elements in the list expressions.
/// Also passes the length of the two lists.
fn vec_expr_pairs_op<T, A>(a: &Expr, b: &Expr, f: PairsCallback<T, A>) -> Option<A>
where
    T: TryFrom<Lit>,
{
    let a_exprs = a.clone().unwrap_matrix_unchecked()?.0;
    let b_exprs = b.clone().unwrap_matrix_unchecked()?.0;
    let lens = (a_exprs.len(), b_exprs.len());

    let lit_pairs = std::iter::zip(a_exprs, b_exprs)
        .map(|(a, b)| Some((unwrap_expr(&a)?, unwrap_expr(&b)?)))
        .collect::<Option<Vec<(T, T)>>>()?;
    Some(f(lit_pairs, lens))
}

/// Same as [`vec_expr_pairs_op`], but over slices of atoms.
fn atoms_pairs_op<T, A>(a: &[Atom], b: &[Atom], f: PairsCallback<T, A>) -> Option<A>
where
    T: TryFrom<Atom>,
{
    let lit_pairs = Iterator::zip(a.iter(), b.iter())
        .map(|(a, b)| Some((a.clone().try_into().ok()?, b.clone().try_into().ok()?)))
        .collect::<Option<Vec<(T, T)>>>()?;
    Some(f(lit_pairs, (a.len(), b.len())))
}

pub fn opt_vec_op<T, A>(f: fn(Vec<T>) -> Option<A>, a: &[Expr]) -> Option<A>
where
    T: TryFrom<Lit>,
{
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
    f(a)
}

pub fn opt_vec_lit_op<T, A>(f: fn(Vec<T>) -> Option<A>, a: &Expr) -> Option<A>
where
    T: TryFrom<Lit>,
{
    f(eval_list_items(a)?)
}

#[allow(dead_code)]
pub fn flat_op<T, A>(f: fn(Vec<T>, T) -> A, a: &[Expr], b: &Expr) -> Option<A>
where
    T: TryFrom<Lit>,
{
    let a = a.iter().map(unwrap_expr).collect::<Option<Vec<T>>>()?;
    let b = unwrap_expr::<T>(b)?;
    Some(f(a, b))
}

pub fn unwrap_expr<T: TryFrom<Lit>>(expr: &Expr) -> Option<T> {
    let c = eval_constant(expr)?;
    TryInto::<T>::try_into(c).ok()
}

fn eval_list_items<T>(expr: &Expr) -> Option<Vec<T>>
where
    T: TryFrom<Lit>,
{
    if let Some(items) = expr
        .clone()
        .unwrap_matrix_unchecked()
        .map(|(items, _)| items)
    {
        return items.iter().map(unwrap_expr).collect();
    }

    let Lit::AbstractLiteral(list) = eval_constant(expr)? else {
        return None;
    };

    let items = list.unwrap_list()?;
    items
        .iter()
        .cloned()
        .map(TryInto::try_into)
        .collect::<Result<Vec<_>, _>>()
        .ok()
}

fn eval_constant_comprehension(comprehension: &Comprehension) -> Option<Lit> {
    let mut values = Vec::new();
    eval_comprehension_qualifiers(comprehension, 0, &mut values)?;
    Some(Lit::AbstractLiteral(
        AbstractLiteral::matrix_implied_indices(values),
    ))
}

fn eval_comprehension_qualifiers(
    comprehension: &Comprehension,
    qualifier_index: usize,
    values: &mut Vec<Lit>,
) -> Option<()> {
    if qualifier_index == comprehension.qualifiers.len() {
        values.push(eval_constant(&comprehension.return_expression)?);
        return Some(());
    }

    match &comprehension.qualifiers[qualifier_index] {
        ComprehensionQualifier::Generator { ptr } => {
            let domain = ptr.domain()?;
            let generator_values = domain.resolve()?.values().ok()?.collect_vec();

            for value in generator_values {
                with_temporary_quantified_binding(ptr, &value, || {
                    eval_comprehension_qualifiers(comprehension, qualifier_index + 1, values)
                })?;
            }
        }
        ComprehensionQualifier::ExpressionGenerator { ptr } => {
            // clone immediately so the read lock guard is dropped
            let expr = ptr.as_quantified_expr()?.clone();
            let generator_values = generator_values_from_expr(&expr)?;

            for value in generator_values {
                with_temporary_quantified_binding(ptr, &value, || {
                    eval_comprehension_qualifiers(comprehension, qualifier_index + 1, values)
                })?;
            }
        }
        ComprehensionQualifier::Condition(condition) => match eval_constant(condition)? {
            Lit::Bool(true) => {
                eval_comprehension_qualifiers(comprehension, qualifier_index + 1, values)?
            }
            Lit::Bool(false) => {}
            _ => return None,
        },
    }

    Some(())
}

fn generator_values_from_expr(expr: &Expr) -> Option<Vec<Lit>> {
    match eval_constant(expr)? {
        Lit::AbstractLiteral(AbstractLiteral::Set(values))
        | Lit::AbstractLiteral(AbstractLiteral::MSet(values))
        | Lit::AbstractLiteral(AbstractLiteral::Tuple(values)) => Some(values),
        Lit::AbstractLiteral(list) => list.unwrap_list().cloned(),
        _ => None,
    }
}

fn with_temporary_quantified_binding<T>(
    quantified: &crate::ast::DeclarationPtr,
    value: &Lit,
    f: impl FnOnce() -> Option<T>,
) -> Option<T> {
    let mut targets = vec![quantified.clone()];
    if let DeclarationKind::Quantified(inner) = &*quantified.kind()
        && let Some(generator) = inner.generator()
    {
        targets.push(generator.clone());
    }

    let mut originals = Vec::with_capacity(targets.len());
    for mut target in targets {
        let old_kind = target.replace_kind(DeclarationKind::TemporaryValueLetting(Expr::Atomic(
            Metadata::new(),
            Atom::Literal(value.clone()),
        )));
        originals.push((target, old_kind));
    }

    let result = f();

    for (mut target, old_kind) in originals.into_iter().rev() {
        let _ = target.replace_kind(old_kind);
    }

    result
}