conjure_cp_core/ast/

expressions.rs

use std::collections::{HashSet, VecDeque};
use std::fmt::{Display, Formatter};
use std::hash::{DefaultHasher, Hash, Hasher};
use std::sync::atomic::{AtomicU64, Ordering};

static HASH_HITS: AtomicU64 = AtomicU64::new(0);
static HASH_MISSES: AtomicU64 = AtomicU64::new(0);

pub fn print_hash_stats() {
    println!(
        "Expression hash stats: hits={}, misses={}",
        HASH_HITS.load(Ordering::Relaxed),
        HASH_MISSES.load(Ordering::Relaxed)
    );
}
use tracing::trace;

use conjure_cp_enum_compatibility_macro::{document_compatibility, generate_discriminants};
use itertools::Itertools;
use serde::{Deserialize, Serialize};
use tree_morph::cache::CacheHashable;
use ustr::Ustr;

use polyquine::Quine;
use uniplate::{Biplate, Uniplate};

use crate::ast::FuncAttr;
use crate::ast::metadata::NO_HASH;
use crate::bug;

use super::abstract_comprehension::AbstractComprehension;
use super::ac_operators::ACOperatorKind;
use super::categories::{Category, CategoryOf};
use super::comprehension::Comprehension;
use super::declaration::DeclarationKind;
use super::domains::HasDomain as _;
use super::pretty::{pretty_expressions_as_top_level, pretty_vec};
use super::records::Field;
use super::sat_encoding::SATIntEncoding;
use super::{
    AbstractLiteral, Atom, DeclarationPtr, Domain, DomainPtr, GroundDomain, IntVal, JectivityAttr,
    Literal, MSetAttr, Metadata, Model, Moo, Name, PartialityAttr, Range, Reference, RelAttr,
    ReturnType, SetAttr, SymbolTable, SymbolTablePtr, Typeable, UnresolvedDomain, matrix,
};

// Ensure that this type doesn't get too big
//
// If you triggered this assertion, you either made a variant of this enum that is too big, or you
// made Name,Literal,AbstractLiteral,Atom bigger, which made this bigger! To fix this, put some
// stuff in boxes.
//
// Enums take the size of their largest variant, so an enum with mostly small variants and a few
// large ones wastes memory... A larger Expression type also slows down Oxide.
//
// For more information, and more details on type sizes and how to measure them, see the commit
// message for 6012de809 (perf: reduce size of AST types, 2025-06-18).
//
// You can also see type sizes in the rustdoc documentation, generated by ./tools/gen_docs.sh
//
// https://github.com/conjure-cp/conjure-oxide/commit/6012de8096ca491ded91ecec61352fdf4e994f2e

// TODO: box all usages of Metadata to bring this down a bit more - I have added variants to
// ReturnType, and Metadata contains ReturnType, so Metadata has got bigger. Metadata will get a
// lot bigger still when we start using it for memoisation, so it should really be
// boxed ~niklasdewally

// expect size of Expression to be 112 bytes
static_assertions::assert_eq_size!([u8; 112], Expression);

/// Represents different types of expressions used to define rules and constraints in the model.
///
/// The `Expression` enum includes operations, constants, and variable references
/// used to build rules and conditions for the model.
#[generate_discriminants]
#[document_compatibility]
#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize, Uniplate, Quine)]
#[biplate(to=AbstractComprehension)]
#[biplate(to=AbstractLiteral<Expression>)]
#[biplate(to=AbstractLiteral<Literal>)]
#[biplate(to=Atom)]
#[biplate(to=Comprehension)]
#[biplate(to=DeclarationPtr)]
#[biplate(to=DomainPtr)]
#[biplate(to=Literal)]
#[biplate(to=Metadata)]
#[biplate(to=Name)]
#[biplate(to=Option<Expression>)]
#[biplate(to=Field<Expression>)]
#[biplate(to=Field<Literal>)]
#[biplate(to=Reference)]
#[biplate(to=Model)]
#[biplate(to=SymbolTable)]
#[biplate(to=SymbolTablePtr)]
#[biplate(to=Vec<Expression>)]
#[path_prefix(conjure_cp::ast)]
pub enum Expression {
    AbstractLiteral(Metadata, AbstractLiteral<Expression>),
    /// The top of the model
    Root(Metadata, Vec<Expression>),

    /// An expression representing "A is valid as long as B is true"
    /// Turns into a conjunction when it reaches a boolean context
    Bubble(Metadata, Moo<Expression>, Moo<Expression>),

    /// A comprehension.
    ///
    /// The inside of the comprehension opens a new scope.
    // todo (gskorokhod): Comprehension contains a symbol table which contains a bunch of pointers.
    // This makes implementing Quine tricky (it doesnt support Rc, by design). Skip it for now.
    #[polyquine_skip]
    Comprehension(Metadata, Moo<Comprehension>),

    /// Higher-level abstract comprehension
    #[polyquine_skip] // no idea what this is lol but it stops rustc screaming at me
    AbstractComprehension(Metadata, Moo<AbstractComprehension>),

    /// Defines dominance ("Solution A is preferred over Solution B")
    DominanceRelation(Metadata, Moo<Expression>),
    /// `fromSolution(name)` - Used in dominance relation definitions
    FromSolution(Metadata, Moo<Atom>),

    #[polyquine_with(arm = (_, name) => {
        let ident = proc_macro2::Ident::new(name.as_str(), proc_macro2::Span::call_site());
        quote::quote! { #ident.clone().into() }
    })]
    Metavar(Metadata, Ustr),

    Atomic(Metadata, Atom),

    /// Asserts that the given variant of a variant expression is in use.
    /// See also: [GroundDomain::Variant]
    #[compatible(JsonInput)]
    Active(Metadata, Moo<Expression>, Name),

    /// Indexing into a record expression, e.g `{foo = 1, bar = true}[foo]`
    /// See also: [GroundDomain::Record]
    #[compatible(JsonInput)]
    RecordField(Metadata, Moo<Expression>, Name),

    /// A matrix index.
    ///
    /// Defined iff the indices are within their respective index domains.
    #[compatible(JsonInput)]
    UnsafeIndex(Metadata, Moo<Expression>, Vec<Expression>),

    /// A safe matrix index.
    ///
    /// See [`Expression::UnsafeIndex`]
    #[compatible(SMT)]
    SafeIndex(Metadata, Moo<Expression>, Vec<Expression>),

    /// A matrix slice: `a[indices]`.
    ///
    /// One of the indicies may be `None`, representing the dimension of the matrix we want to take
    /// a slice of. For example, for some 3d matrix a, `a[1,..,2]` has the indices
    /// `Some(1),None,Some(2)`.
    ///
    /// It is assumed that the slice only has one "wild-card" dimension and thus is 1 dimensional.
    ///
    /// Defined iff the defined indices are within their respective index domains.
    #[compatible(JsonInput)]
    UnsafeSlice(Metadata, Moo<Expression>, Vec<Option<Expression>>),

    /// A safe matrix slice: `a[indices]`.
    ///
    /// See [`Expression::UnsafeSlice`].
    SafeSlice(Metadata, Moo<Expression>, Vec<Option<Expression>>),

    /// `inDomain(x,domain)` iff `x` is in the domain `domain`.
    ///
    /// This cannot be constructed from Essence input, nor passed to a solver: this expression is
    /// mainly used during the conversion of `UnsafeIndex` and `UnsafeSlice` to `SafeIndex` and
    /// `SafeSlice` respectively.
    InDomain(Metadata, Moo<Expression>, DomainPtr),

    /// `toInt(b)` casts boolean expression b to an integer.
    ///
    /// - If b is false, then `toInt(b) == 0`
    ///
    /// - If b is true, then `toInt(b) == 1`
    #[compatible(SMT)]
    ToInt(Metadata, Moo<Expression>),

    /// `|x|` - absolute value of `x`
    #[compatible(JsonInput, SMT)]
    Abs(Metadata, Moo<Expression>),

    /// `sum(<vec_expr>)`
    #[compatible(JsonInput, SMT)]
    Sum(Metadata, Moo<Expression>),

    /// `a * b * c * ...`
    #[compatible(JsonInput, SMT)]
    Product(Metadata, Moo<Expression>),

    /// `min(<vec_expr>)`
    #[compatible(JsonInput, SMT)]
    Min(Metadata, Moo<Expression>),

    /// `max(<vec_expr>)`
    #[compatible(JsonInput, SMT)]
    Max(Metadata, Moo<Expression>),

    /// `not(a)`
    #[compatible(JsonInput, SAT, SMT)]
    Not(Metadata, Moo<Expression>),

    /// `or(<vec_expr>)`
    #[compatible(JsonInput, SAT, SMT)]
    Or(Metadata, Moo<Expression>),

    /// `and(<vec_expr>)`
    #[compatible(JsonInput, SAT, SMT)]
    And(Metadata, Moo<Expression>),

    /// Ensures that `a->b` (material implication).
    #[compatible(JsonInput, SMT)]
    Imply(Metadata, Moo<Expression>, Moo<Expression>),

    /// `iff(a, b)` a <-> b
    #[compatible(JsonInput, SMT)]
    Iff(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    Union(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    In(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    Intersect(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    Supset(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    SupsetEq(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    Subset(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    SubsetEq(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput, SMT)]
    Eq(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput, SMT)]
    Neq(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput, SMT)]
    Geq(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput, SMT)]
    Leq(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput, SMT)]
    Gt(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput, SMT)]
    Lt(Metadata, Moo<Expression>, Moo<Expression>),

    /// `s subsequence t` tests whether the list of values taken by s occurs in the same order
    /// in the list of values taken by t
    #[compatible(JsonInput)]
    Subsequence(Metadata, Moo<Expression>, Moo<Expression>),

    /// `s substring t` tests whether the list of values taken by s occurs in the same order
    /// and contiguously in the list of values taken by t
    #[compatible(JsonInput)]
    Substring(Metadata, Moo<Expression>, Moo<Expression>),

    /// Division after preventing division by zero, usually with a bubble
    #[compatible(SMT)]
    SafeDiv(Metadata, Moo<Expression>, Moo<Expression>),

    /// Division with a possibly undefined value (division by 0)
    #[compatible(JsonInput)]
    UnsafeDiv(Metadata, Moo<Expression>, Moo<Expression>),

    /// Modulo after preventing mod 0, usually with a bubble
    #[compatible(SMT)]
    SafeMod(Metadata, Moo<Expression>, Moo<Expression>),

    /// Modulo with a possibly undefined value (mod 0)
    #[compatible(JsonInput)]
    UnsafeMod(Metadata, Moo<Expression>, Moo<Expression>),

    /// Negation: `-x`
    #[compatible(JsonInput, SMT)]
    Neg(Metadata, Moo<Expression>),

    /// Factorial: `x!` or 'factorial(x)`
    #[compatible(JsonInput)]
    Factorial(Metadata, Moo<Expression>),

    /// Set of domain values function is defined for
    #[compatible(JsonInput)]
    Defined(Metadata, Moo<Expression>),

    /// Set of codomain values function is defined for
    #[compatible(JsonInput)]
    Range(Metadata, Moo<Expression>),

    #[compatible(JsonInput)]
    ToSet(Metadata, Moo<Expression>),

    #[compatible(JsonInput)]
    ToMSet(Metadata, Moo<Expression>),

    #[compatible(JsonInput)]
    ToRelation(Metadata, Moo<Expression>),

    /// Unsafe power`x**y` (possibly undefined)
    ///
    /// Defined when (X!=0 \\/ Y!=0) /\ Y>=0
    #[compatible(JsonInput)]
    UnsafePow(Metadata, Moo<Expression>, Moo<Expression>),

    /// `UnsafePow` after preventing undefinedness
    SafePow(Metadata, Moo<Expression>, Moo<Expression>),

    /// Flatten matrix operator
    /// `flatten(M)` or `flatten(n, M)`
    /// where M is a matrix and n is an optional integer argument indicating depth of flattening
    Flatten(Metadata, Option<Moo<Expression>>, Moo<Expression>),

    /// `allDiff(<vec_expr>)`
    #[compatible(JsonInput)]
    AllDiff(Metadata, Moo<Expression>),

    /// `table([x1, x2, ...], [[r11, r12, ...], [r21, r22, ...], ...])`
    ///
    /// Represents a positive table constraint: the tuple `[x1, x2, ...]` must match one of the
    /// allowed rows.
    #[compatible(JsonInput)]
    Table(Metadata, Moo<Expression>, Moo<Expression>),

    /// `negativeTable([x1, x2, ...], [[r11, r12, ...], [r21, r22, ...], ...])`
    ///
    /// Represents a negative table constraint: the tuple `[x1, x2, ...]` must NOT match any of the
    /// forbidden rows.
    #[compatible(JsonInput)]
    NegativeTable(Metadata, Moo<Expression>, Moo<Expression>),
    /// Binary subtraction operator
    ///
    /// This is a parser-level construct, and is immediately normalised to `Sum([a,-b])`.
    /// TODO: make this compatible with Set Difference calculations - need to change return type and domain for this expression and write a set comprehension rule.
    /// have already edited minus_to_sum to prevent this from applying to sets
    #[compatible(JsonInput)]
    Minus(Metadata, Moo<Expression>, Moo<Expression>),

    /// Partition Operator: test if a list of elements are not all contained in one part of the partition
    /// First Expr Arg is a list of elements
    /// Second Expr Arg is the partition
    #[compatible(JsonInput)]
    Apart(Metadata, Moo<Expression>, Moo<Expression>),

    /// Partition Operator: union of all parts of a partition
    /// Expr Arg is a partition
    #[compatible(JsonInput)]
    Participants(Metadata, Moo<Expression>),

    /// Partition Operator: part of partition that contains specified element
    /// First Expr Arg is an element that should be contained
    /// Second Expr Arg is the partition that should contain that element
    #[compatible(JsonInput)]
    Party(Metadata, Moo<Expression>, Moo<Expression>),

    /// Partition Operator: partition to its set of parts
    /// Expr Arg is the partition from which the parts come from
    #[compatible(JsonInput)]
    Parts(Metadata, Moo<Expression>),

    /// Partition Operator: test if a list of elements are all in the same part of the partition
    /// First Expr Arg is the list of elements to test with
    /// Second Expr Arg is the partition to test on
    #[compatible(JsonInput)]
    Together(Metadata, Moo<Expression>, Moo<Expression>),

    /// Ensures that x=|y| i.e. x is the absolute value of y.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#abs)
    #[compatible(Minion)]
    FlatAbsEq(Metadata, Moo<Atom>, Moo<Atom>),

    /// Ensures that `alldiff([a,b,...])`.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#alldiff)
    #[compatible(Minion)]
    FlatAllDiff(Metadata, Vec<Atom>),

    /// Ensures that sum(vec) >= x.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#sumgeq)
    #[compatible(Minion)]
    FlatSumGeq(Metadata, Vec<Atom>, Atom),

    /// Ensures that sum(vec) <= x.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#sumleq)
    #[compatible(Minion)]
    FlatSumLeq(Metadata, Vec<Atom>, Atom),

    /// `ineq(x,y,k)` ensures that x <= y + k.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#ineq)
    #[compatible(Minion)]
    FlatIneq(Metadata, Moo<Atom>, Moo<Atom>, Box<Literal>),

    /// `w-literal(x,k)` ensures that x == k, where x is a variable and k a constant.
    ///
    /// Low-level Minion constraint.
    ///
    /// This is a low-level Minion constraint and you should probably use Eq instead. The main use
    /// of w-literal is to convert boolean variables to constraints so that they can be used inside
    /// watched-and and watched-or.
    ///
    /// # See also
    ///
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#minuseq)
    /// + `rules::minion::boolean_literal_to_wliteral`.
    #[compatible(Minion)]
    #[polyquine_skip]
    FlatWatchedLiteral(Metadata, Reference, Literal),

    /// `weightedsumleq(cs,xs,total)` ensures that cs.xs <= total, where cs.xs is the scalar dot
    /// product of cs and xs.
    ///
    /// Low-level Minion constraint.
    ///
    /// Represents a weighted sum of the form `ax + by + cz + ...`
    ///
    /// # See also
    ///
    /// + [Minion
    /// documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#weightedsumleq)
    FlatWeightedSumLeq(Metadata, Vec<Literal>, Vec<Atom>, Moo<Atom>),

    /// `weightedsumgeq(cs,xs,total)` ensures that cs.xs >= total, where cs.xs is the scalar dot
    /// product of cs and xs.
    ///
    /// Low-level Minion constraint.
    ///
    /// Represents a weighted sum of the form `ax + by + cz + ...`
    ///
    /// # See also
    ///
    /// + [Minion
    /// documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#weightedsumleq)
    FlatWeightedSumGeq(Metadata, Vec<Literal>, Vec<Atom>, Moo<Atom>),

    /// Ensures that x =-y, where x and y are atoms.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#minuseq)
    #[compatible(Minion)]
    FlatMinusEq(Metadata, Moo<Atom>, Moo<Atom>),

    /// Ensures that x*y=z.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#product)
    #[compatible(Minion)]
    FlatProductEq(Metadata, Moo<Atom>, Moo<Atom>, Moo<Atom>),

    /// Ensures that floor(x/y)=z. Always true when y=0.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#div_undefzero)
    #[compatible(Minion)]
    MinionDivEqUndefZero(Metadata, Moo<Atom>, Moo<Atom>, Moo<Atom>),

    /// Ensures that x%y=z. Always true when y=0.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#mod_undefzero)
    #[compatible(Minion)]
    MinionModuloEqUndefZero(Metadata, Moo<Atom>, Moo<Atom>, Moo<Atom>),

    /// Ensures that `x**y = z`.
    ///
    /// Low-level Minion constraint.
    ///
    /// This constraint is false when `y<0` except for `1**y=1` and `(-1)**y=z` (where z is 1 if y
    /// is odd and z is -1 if y is even).
    ///
    /// # See also
    ///
    /// + [Github comment about `pow` semantics](https://github.com/minion/minion/issues/40#issuecomment-2595914891)
    /// + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#pow)
    MinionPow(Metadata, Moo<Atom>, Moo<Atom>, Moo<Atom>),

    /// `reify(constraint,r)` ensures that r=1 iff `constraint` is satisfied, where r is a 0/1
    /// variable.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    ///  + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#reify)
    #[compatible(Minion)]
    MinionReify(Metadata, Moo<Expression>, Atom),

    /// `reifyimply(constraint,r)` ensures that `r->constraint`, where r is a 0/1 variable.
    /// variable.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    ///  + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#reifyimply)
    #[compatible(Minion)]
    MinionReifyImply(Metadata, Moo<Expression>, Atom),

    /// `w-inintervalset(x, [a1,a2, b1,b2, … ])` ensures that the value of x belongs to one of the
    /// intervals {a1,…,a2}, {b1,…,b2} etc.
    ///
    /// The list of intervals must be given in numerical order.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///>
    ///  + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#w-inintervalset)
    #[compatible(Minion)]
    MinionWInIntervalSet(Metadata, Atom, Vec<i32>),

    /// `w-inset(x, [v1, v2, … ])` ensures that the value of `x` is one of the explicitly given values `v1`, `v2`, etc.
    ///
    /// This constraint enforces membership in a specific set of discrete values rather than intervals.
    ///
    /// The list of values must be given in numerical order.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    ///  + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#w-inset)
    #[compatible(Minion)]
    MinionWInSet(Metadata, Atom, Vec<i32>),

    /// `element_one(vec, i, e)` specifies that `vec[i] = e`. This implies that i is
    /// in the range `[1..len(vec)]`.
    ///
    /// Low-level Minion constraint.
    ///
    /// # See also
    ///
    ///  + [Minion documentation](https://minion-solver.readthedocs.io/en/stable/usage/constraints.html#element_one)
    #[compatible(Minion)]
    MinionElementOne(Metadata, Vec<Atom>, Moo<Atom>, Moo<Atom>),

    /// Declaration of an auxiliary variable.
    ///
    /// As with Savile Row, we semantically distinguish this from `Eq`.
    #[compatible(Minion)]
    #[polyquine_skip]
    AuxDeclaration(Metadata, Reference, Moo<Expression>),

    /// This expression is for encoding ints for the SAT solver, it stores the encoding type, the vector of booleans and the min/max for the int.
    #[compatible(SAT)]
    SATInt(Metadata, SATIntEncoding, Moo<Expression>, (i32, i32)),

    /// Addition over a pair of expressions (i.e. a + b) rather than a vec-expr like Expression::Sum.
    /// This is for compatibility with backends that do not support addition over vectors.
    #[compatible(SMT)]
    PairwiseSum(Metadata, Moo<Expression>, Moo<Expression>),

    /// Multiplication over a pair of expressions (i.e. a * b) rather than a vec-expr like Expression::Product.
    /// This is for compatibility with backends that do not support multiplication over vectors.
    #[compatible(SMT)]
    PairwiseProduct(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    Image(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    ImageSet(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    PreImage(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    Inverse(Metadata, Moo<Expression>, Moo<Expression>),

    #[compatible(JsonInput)]
    Restrict(Metadata, Moo<Expression>, Moo<Expression>),

    /// Lexicographical < between two matrices.
    ///
    /// A <lex B iff: A[i] < B[i] for some i /\ (A[j] > B[j] for some j -> i < j)
    /// I.e. A must be less than B at some index i, and if it is greater than B at another index j,
    /// then j comes after i.
    /// I.e. A must be greater than B at the first index where they differ.
    ///
    /// E.g. [1, 1] <lex [2, 1] and [1, 1] <lex [1, 2]
    LexLt(Metadata, Moo<Expression>, Moo<Expression>),

    /// Lexicographical <= between two matrices
    LexLeq(Metadata, Moo<Expression>, Moo<Expression>),

    /// Lexicographical > between two matrices
    /// This is a parser-level construct, and is immediately normalised to LexLt(b, a)
    LexGt(Metadata, Moo<Expression>, Moo<Expression>),

    /// Lexicographical >= between two matrices
    /// This is a parser-level construct, and is immediately normalised to LexLeq(b, a)
    LexGeq(Metadata, Moo<Expression>, Moo<Expression>),

    /// Low-level minion constraint. See Expression::LexLt
    FlatLexLt(Metadata, Vec<Atom>, Vec<Atom>),

    /// Low-level minion constraint. See Expression::LexLeq
    FlatLexLeq(Metadata, Vec<Atom>, Vec<Atom>),

    /// Alters the shape of relations by projection
    #[compatible(JsonInput)]
    RelationProj(Metadata, Moo<Expression>, Vec<Option<Expression>>),

    /// Cardinality of a collection type
    #[compatible(JsonInput)]
    Card(Metadata, Moo<Expression>),
}

// for the given matrix literal, return a bounded domain from the min to max of applying op to each
// child expression.
//
// Op must be monotonic.
//
// Returns none if unbounded
fn bounded_i32_domain_for_matrix_literal_monotonic(
    e: &Expression,
    op: fn(i32, i32) -> Option<i32>,
) -> Option<DomainPtr> {
    // only care about the elements, not the indices
    let (mut exprs, _) = e.clone().unwrap_matrix_unchecked()?;

    // fold each element's domain into one using op.
    //
    // here, I assume that op is monotone. This means that the bounds of op([a1,a2],[b1,b2])  for
    // the ranges [a1,a2], [b1,b2] will be
    // [min(op(a1,b1),op(a2,b1),op(a1,b2),op(a2,b2)),max(op(a1,b1),op(a2,b1),op(a1,b2),op(a2,b2))].
    //
    // We used to not assume this, and work out the bounds by applying op on the Cartesian product
    // of A and B; however, this caused a combinatorial explosion and my computer to run out of
    // memory (on the hakank_eprime_xkcd test)...
    //Int
    // For example, to find the bounds of the intervals [1,4], [1,5] combined using op, we used to do
    //  [min(op(1,1), op(1,2),op(1,3),op(1,4),op(1,5),op(2,1)..
    //
    // +,-,/,* are all monotone, so this assumption should be fine for now...

    let expr = exprs.pop()?;
    let dom = expr.domain_of()?;
    let resolved = dom.resolve().ok()?;
    let GroundDomain::Int(ranges) = resolved.as_ref() else {
        return None;
    };

    let (mut current_min, mut current_max) = range_vec_bounds_i32(ranges)?;

    for expr in exprs {
        let dom = expr.domain_of()?;
        let resolved = dom.resolve().ok()?;
        let GroundDomain::Int(ranges) = resolved.as_ref() else {
            return None;
        };

        let (min, max) = range_vec_bounds_i32(ranges)?;

        // all the possible new values for current_min / current_max
        let minmax = op(min, current_max)?;
        let minmin = op(min, current_min)?;
        let maxmin = op(max, current_min)?;
        let maxmax = op(max, current_max)?;
        let vals = [minmax, minmin, maxmin, maxmax];

        current_min = *vals
            .iter()
            .min()
            .expect("vals iterator should not be empty, and should have a minimum.");
        current_max = *vals
            .iter()
            .max()
            .expect("vals iterator should not be empty, and should have a maximum.");
    }

    if current_min == current_max {
        Some(Domain::int(vec![Range::Single(current_min)]))
    } else {
        Some(Domain::int(vec![Range::Bounded(current_min, current_max)]))
    }
}

fn matrix_element_domain(e: &Expression) -> Option<DomainPtr> {
    let (elem_domain, _) = e.domain_of()?.as_matrix()?;
    elem_domain.as_ref().as_int()?;
    Some(elem_domain)
}

// Returns none if unbounded
fn range_vec_bounds_i32(ranges: &Vec<Range<i32>>) -> Option<(i32, i32)> {
    let mut min = i32::MAX;
    let mut max = i32::MIN;
    for r in ranges {
        match r {
            Range::Single(i) => {
                if *i < min {
                    min = *i;
                }
                if *i > max {
                    max = *i;
                }
            }
            Range::Bounded(i, j) => {
                if *i < min {
                    min = *i;
                }
                if *j > max {
                    max = *j;
                }
            }
            Range::UnboundedR(_) | Range::UnboundedL(_) | Range::Unbounded => return None,
        }
    }
    Some((min, max))
}

impl Expression {
    /// Returns the possible values of the expression, recursing to leaf expressions
    pub fn domain_of(&self) -> Option<DomainPtr> {
        match self {
            Expression::Union(_, a, b) => Some(Domain::set(
                SetAttr::<IntVal>::default(),
                a.domain_of()?.union(&b.domain_of()?).ok()?,
            )),
            Expression::Intersect(_, a, b) => Some(Domain::set(
                SetAttr::<IntVal>::default(),
                a.domain_of()?.intersect(&b.domain_of()?).ok()?,
            )),
            Expression::In(_, _, _) => Some(Domain::bool()),
            Expression::Supset(_, _, _) => Some(Domain::bool()),
            Expression::SupsetEq(_, _, _) => Some(Domain::bool()),
            Expression::Subset(_, _, _) => Some(Domain::bool()),
            Expression::SubsetEq(_, _, _) => Some(Domain::bool()),
            Expression::AbstractLiteral(_, abslit) => abslit.domain_of(),
            Expression::DominanceRelation(_, _) => Some(Domain::bool()),
            Expression::FromSolution(_, expr) => Some(expr.domain_of()),
            Expression::Metavar(_, _) => None,
            Expression::Comprehension(_, comprehension) => comprehension.domain_of(),
            Expression::AbstractComprehension(_, comprehension) => comprehension.domain_of(),
            Expression::RecordField(_, rec, field_name) => {
                let rec_ents = rec.domain_of()?.as_record()?;
                for ent in rec_ents {
                    if ent.name.eq(field_name) {
                        return Some(ent.value);
                    }
                }
                None
            }
            Expression::UnsafeIndex(_, matrix, index) | Expression::SafeIndex(_, matrix, index) => {
                let dom = matrix.domain_of()?;
                if let Some((elem_domain, _)) = dom.as_matrix() {
                    return Some(elem_domain);
                }

                // may actually use the value in the future
                #[allow(clippy::redundant_pattern_matching)]
                if let Some(_) = dom.as_tuple() {
                    // TODO: We can implement proper indexing for tuples
                    return None;
                }

                if let Some(doms) = dom.as_variant().or(dom.as_record()) {
                    let index_expr = index.first()?;
                    return match index_expr {
                        Expression::Atomic(_, atom) => {
                            let decl = atom.clone().into_declaration();
                            for inner_dom in doms {
                                if *decl.name() == inner_dom.name {
                                    return Some(inner_dom.value);
                                }
                            }
                            None
                        }
                        _ => None,
                    };
                }

                bug!(
                    "subject of an index operation should support indexing, but got {matrix}: {dom}"
                )
            }
            Expression::UnsafeSlice(_, matrix, indices)
            | Expression::SafeSlice(_, matrix, indices) => {
                let sliced_dimension = indices.iter().position(Option::is_none);

                let dom = matrix.domain_of()?;
                let Some((elem_domain, index_domains)) = dom.as_matrix() else {
                    bug!("subject of an index operation should be a matrix");
                };

                match sliced_dimension {
                    Some(dimension) => Some(Domain::matrix(
                        elem_domain,
                        vec![index_domains[dimension].clone()],
                    )),

                    // same as index
                    None => Some(elem_domain),
                }
            }
            Expression::InDomain(_, _, _) => Some(Domain::bool()),
            Expression::Atomic(_, atom) => Some(atom.domain_of()),
            Expression::Sum(_, e) => {
                bounded_i32_domain_for_matrix_literal_monotonic(e, |x, y| Some(x + y))
            }
            Expression::Product(_, e) => {
                bounded_i32_domain_for_matrix_literal_monotonic(e, |x, y| Some(x * y))
            }
            Expression::Min(_, e) => bounded_i32_domain_for_matrix_literal_monotonic(e, |x, y| {
                Some(if x < y { x } else { y })
            })
            .or_else(|| matrix_element_domain(e)),
            Expression::Max(_, e) => bounded_i32_domain_for_matrix_literal_monotonic(e, |x, y| {
                Some(if x > y { x } else { y })
            })
            .or_else(|| matrix_element_domain(e)),
            Expression::UnsafeDiv(_, a, b) => a
                .domain_of()?
                .resolve()
                .ok()?
                .apply_i32(
                    // rust integer division is truncating; however, we want to always round down,
                    // including for negative numbers.
                    |x, y| {
                        if y != 0 {
                            Some((x as f32 / y as f32).floor() as i32)
                        } else {
                            None
                        }
                    },
                    b.domain_of()?.resolve().ok()?.as_ref(),
                )
                .map(DomainPtr::from)
                .ok(),
            Expression::SafeDiv(_, a, b) => {
                // rust integer division is truncating; however, we want to always round down
                // including for negative numbers.
                let domain = a
                    .domain_of()?
                    .resolve()
                    .ok()?
                    .apply_i32(
                        |x, y| {
                            if y != 0 {
                                Some((x as f32 / y as f32).floor() as i32)
                            } else {
                                None
                            }
                        },
                        b.domain_of()?.resolve().ok()?.as_ref(),
                    )
                    .unwrap_or_else(|err| bug!("Got {err} when computing domain of {self}"));

                if let GroundDomain::Int(ranges) = domain {
                    let mut ranges = ranges;
                    ranges.push(Range::Single(0));
                    Some(Domain::int(ranges))
                } else {
                    bug!("Domain of {self} was not integer")
                }
            }
            Expression::UnsafeMod(_, a, b) => a
                .domain_of()?
                .resolve()
                .ok()?
                .apply_i32(
                    |x, y| if y != 0 { Some(x % y) } else { None },
                    b.domain_of()?.resolve().ok()?.as_ref(),
                )
                .map(DomainPtr::from)
                .ok(),
            Expression::SafeMod(_, a, b) => {
                let domain = a
                    .domain_of()?
                    .resolve()
                    .ok()?
                    .apply_i32(
                        |x, y| if y != 0 { Some(x % y) } else { None },
                        b.domain_of()?.resolve().ok()?.as_ref(),
                    )
                    .unwrap_or_else(|err| bug!("Got {err} when computing domain of {self}"));

                if let GroundDomain::Int(ranges) = domain {
                    let mut ranges = ranges;
                    ranges.push(Range::Single(0));
                    Some(Domain::int(ranges))
                } else {
                    bug!("Domain of {self} was not integer")
                }
            }
            Expression::SafePow(_, a, b) | Expression::UnsafePow(_, a, b) => a
                .domain_of()?
                .resolve()
                .ok()?
                .apply_i32(
                    |x, y| {
                        if (x != 0 || y != 0) && y >= 0 {
                            Some(x.pow(y as u32))
                        } else {
                            None
                        }
                    },
                    b.domain_of()?.resolve().ok()?.as_ref(),
                )
                .map(DomainPtr::from)
                .ok(),
            Expression::Root(_, _) => None,
            Expression::Bubble(_, inner, _) => inner.domain_of(),
            Expression::AuxDeclaration(_, _, _) => Some(Domain::bool()),
            Expression::And(_, _) => Some(Domain::bool()),
            Expression::Not(_, _) => Some(Domain::bool()),
            Expression::Or(_, _) => Some(Domain::bool()),
            Expression::Imply(_, _, _) => Some(Domain::bool()),
            Expression::Iff(_, _, _) => Some(Domain::bool()),
            Expression::Eq(_, _, _) => Some(Domain::bool()),
            Expression::Neq(_, _, _) => Some(Domain::bool()),
            Expression::Geq(_, _, _) => Some(Domain::bool()),
            Expression::Leq(_, _, _) => Some(Domain::bool()),
            Expression::Gt(_, _, _) => Some(Domain::bool()),
            Expression::Lt(_, _, _) => Some(Domain::bool()),
            Expression::Factorial(_, _) => None, // not implemented
            Expression::FlatAbsEq(_, _, _) => Some(Domain::bool()),
            Expression::FlatSumGeq(_, _, _) => Some(Domain::bool()),
            Expression::FlatSumLeq(_, _, _) => Some(Domain::bool()),
            Expression::MinionDivEqUndefZero(_, _, _, _) => Some(Domain::bool()),
            Expression::MinionModuloEqUndefZero(_, _, _, _) => Some(Domain::bool()),
            Expression::FlatIneq(_, _, _, _) => Some(Domain::bool()),
            Expression::Flatten(_, n, m) => {
                if let Some(expr) = n {
                    if expr.return_type() == ReturnType::Int {
                        // TODO: handle flatten with depth argument
                        return None;
                    }
                } else {
                    // TODO: currently only works for matrices
                    let dom = m.domain_of()?.resolve().ok()?;
                    let (val_dom, idx_doms) = match dom.as_ref() {
                        GroundDomain::Matrix(val, idx) => (val, idx),
                        _ => return None,
                    };
                    let num_elems = matrix::num_elements(idx_doms).ok()? as i32;

                    let new_index_domain = Domain::int(vec![Range::Bounded(1, num_elems)]);
                    return Some(Domain::matrix(
                        val_dom.clone().into(),
                        vec![new_index_domain],
                    ));
                }
                None
            }
            Expression::AllDiff(_, _) => Some(Domain::bool()),
            Expression::Table(_, _, _) => Some(Domain::bool()),
            Expression::NegativeTable(_, _, _) => Some(Domain::bool()),
            Expression::FlatWatchedLiteral(_, _, _) => Some(Domain::bool()),
            Expression::MinionReify(_, _, _) => Some(Domain::bool()),
            Expression::MinionReifyImply(_, _, _) => Some(Domain::bool()),
            Expression::MinionWInIntervalSet(_, _, _) => Some(Domain::bool()),
            Expression::MinionWInSet(_, _, _) => Some(Domain::bool()),
            Expression::MinionElementOne(_, _, _, _) => Some(Domain::bool()),
            Expression::Neg(_, x) => {
                let dom = x.domain_of()?;
                let mut ranges = dom.as_int()?;

                ranges = ranges
                    .into_iter()
                    .map(|r| match r {
                        Range::Single(x) => Range::Single(-x),
                        Range::Bounded(x, y) => Range::Bounded(-y, -x),
                        Range::UnboundedR(i) => Range::UnboundedL(-i),
                        Range::UnboundedL(i) => Range::UnboundedR(-i),
                        Range::Unbounded => Range::Unbounded,
                    })
                    .collect();

                Some(Domain::int(ranges))
            }
            Expression::Minus(_, a, b) => {
                let a_resolved = a.domain_of()?.resolve().ok()?;
                let b_resolved = b.domain_of()?.resolve().ok()?;

                if matches!(a_resolved.as_ref(), GroundDomain::Int(_))
                    && matches!(b_resolved.as_ref(), GroundDomain::Int(_))
                {
                    a_resolved
                        .apply_i32(|x, y| Some(x - y), b_resolved.as_ref())
                        .map(DomainPtr::from)
                        .ok()
                } else if matches!(a_resolved.as_ref(), GroundDomain::Set(_, _))
                    && matches!(b_resolved.as_ref(), GroundDomain::Set(_, _))
                {
                    Some(DomainPtr::from(a_resolved))
                } else {
                    None
                }
            }
            Expression::FlatAllDiff(_, _) => Some(Domain::bool()),
            Expression::FlatMinusEq(_, _, _) => Some(Domain::bool()),
            Expression::FlatProductEq(_, _, _, _) => Some(Domain::bool()),
            Expression::FlatWeightedSumLeq(_, _, _, _) => Some(Domain::bool()),
            Expression::FlatWeightedSumGeq(_, _, _, _) => Some(Domain::bool()),
            Expression::Abs(_, a) => a
                .domain_of()?
                .resolve()
                .ok()?
                .apply_i32(
                    |a, _| Some(a.abs()),
                    a.domain_of()?.resolve().ok()?.as_ref(),
                )
                .map(DomainPtr::from)
                .ok(),
            Expression::MinionPow(_, _, _, _) => Some(Domain::bool()),
            Expression::ToInt(_, _) => Some(Domain::int(vec![Range::Bounded(0, 1)])),
            Expression::SATInt(_, _, _, (low, high)) => {
                Some(Domain::int_ground(vec![Range::Bounded(*low, *high)]))
            }
            Expression::PairwiseSum(_, a, b) => a
                .domain_of()?
                .resolve()
                .ok()?
                .apply_i32(|a, b| Some(a + b), b.domain_of()?.resolve().ok()?.as_ref())
                .map(DomainPtr::from)
                .ok(),
            Expression::PairwiseProduct(_, a, b) => a
                .domain_of()?
                .resolve()
                .ok()?
                .apply_i32(|a, b| Some(a * b), b.domain_of()?.resolve().ok()?.as_ref())
                .map(DomainPtr::from)
                .ok(),
            Expression::Defined(_, function) => {
                let (attrs, domain, codomain) = function.domain_of()?.as_function()?;
                let size = Self::function_elements_size(attrs, &domain, &codomain);
                if let Some(size) = size {
                    Some(Domain::set(SetAttr::new(size), domain))
                } else {
                    Some(Domain::empty(ReturnType::Set(Box::new(
                        domain.return_type(),
                    ))))
                }
            }
            Expression::Range(_, function) => {
                let (attrs, domain, codomain) = function.domain_of()?.as_function()?;
                let jectivity = attrs.resolve().ok()?.jectivity;

                let size_size = attrs.resolve().ok()?.size;
                let size_size = match size_size {
                    Range::Unbounded => Range::UnboundedR(0),
                    // If lower bound we can guarantee one mapping (unless size = 0)
                    Range::Single(x) => match jectivity {
                        JectivityAttr::Injective | JectivityAttr::Surjective => Range::Single(x),
                        _ => Range::Bounded(Ord::min(1, x), x),
                    },
                    // Upper bound guarantees the same upper bound
                    Range::UnboundedL(x) => Range::Bounded(0, x),
                    // If not bounded by 0 can guarantee min 1
                    Range::UnboundedR(x) => match jectivity {
                        JectivityAttr::Injective | JectivityAttr::Surjective => {
                            Range::UnboundedR(x)
                        }
                        _ => Range::UnboundedR(Ord::min(1, x)),
                    },
                    Range::Bounded(x, y) => Range::Bounded(Ord::min(1, x), y),
                };

                // Gets the size imposed by the partiality and jectivity attributes
                let partiality = attrs.resolve().ok()?.partiality;
                let codomain_length = codomain.length_signed();
                let attr_size = match jectivity {
                    // Bijective and surjective functions must have every element in the codomain mapped to
                    JectivityAttr::Bijective | JectivityAttr::Surjective => match codomain_length {
                        Ok(co_len) => Some(Range::Single(co_len)),
                        Err(_) => None,
                    },
                    JectivityAttr::Injective => {
                        let domain_length = domain.length_signed();
                        match domain_length {
                            Ok(len) => match codomain_length {
                                Ok(co_len) => match partiality {
                                    // When its injective we can guarantee 1 to 1, so the maximum domain length is a single bound
                                    PartialityAttr::Total => {
                                        Some(Range::Single(Ord::min(len, co_len)))
                                    }
                                    PartialityAttr::Partial => {
                                        Some(Range::Bounded(0, Ord::min(len, co_len)))
                                    }
                                },
                                Err(_) => None,
                            },
                            Err(_) => None,
                        }
                    }
                    JectivityAttr::None => {
                        let domain_length = domain.length_signed();
                        match domain_length {
                            // This is the general case, where we know there cannot be more codomain elements mapped to that domain elements
                            Ok(len) => Some(Range::Bounded(0, len)),
                            Err(_) => None,
                        }
                    }
                };

                let size = match attr_size {
                    Some(attr_size) => {
                        let unsafe_range = Range::minimal(&[size_size, attr_size]);
                        match unsafe_range {
                            Ok(range) => range,
                            Err(_) => {
                                return Some(Domain::empty(ReturnType::Set(Box::new(
                                    domain.return_type(),
                                ))));
                            }
                        }
                    }
                    None => size_size,
                };
                Some(Domain::set(SetAttr::new(size), codomain))
            }
            Expression::Image(_, function, _) => get_function_codomain(function),
            Expression::ImageSet(_, function, _) => {
                let codomain = get_function_codomain(function);
                // An imageSet is the converted to a set, and can be empty
                codomain.map(|inner_dom| Domain::set(SetAttr::new(Range::Bounded(0, 1)), inner_dom))
            }
            Expression::PreImage(_, function, _) => {
                let (attrs, domain, codomain) = function.domain_of()?.as_function()?;

                let size_size = attrs.resolve().ok()?.size;
                let size_size = match size_size {
                    // Our only guarantee is an upper bound is the same
                    Range::Unbounded => Range::UnboundedR(0),
                    Range::Single(x) => Range::Bounded(0, x),
                    Range::UnboundedL(x) => Range::Bounded(0, x),
                    Range::UnboundedR(_) => Range::UnboundedR(0),
                    Range::Bounded(_, y) => Range::Bounded(0, y),
                };

                let jectivity = attrs.resolve().ok()?.jectivity;
                let codomain_length = codomain.length_signed();
                let attr_size = match jectivity {
                    // When there is 1-to-1 mapping we can guarantee no more than 1 occurrence
                    JectivityAttr::Bijective => Some(Range::Single(1)),
                    JectivityAttr::Injective => match size_size {
                        Range::Single(x) | Range::UnboundedL(x) | Range::Bounded(x, _) => {
                            match codomain_length {
                                Ok(co_len) => {
                                    if x >= co_len {
                                        Some(Range::Single(1))
                                    } else {
                                        Some(Range::Bounded(0, 1))
                                    }
                                }
                                Err(_) => Some(Range::Bounded(0, 1)),
                            }
                        }
                        _ => Some(Range::Bounded(0, 1)),
                    },
                    JectivityAttr::Surjective => {
                        let domain_length = domain.length_signed();
                        match domain_length {
                            Ok(len) => match codomain_length {
                                // We know the element is mapped but not how many times
                                // Every element must be mapped so it cannot be every element of domain
                                Ok(co_len) => match size_size {
                                    Range::Bounded(_, x)
                                    | Range::UnboundedL(x)
                                    | Range::Single(x) => Some(Range::Bounded(
                                        1,
                                        Ord::max(Ord::min(len, x) - co_len + 1, 0),
                                    )),
                                    _ => Some(Range::Bounded(1, Ord::max(len - co_len + 1, 0))),
                                },
                                Err(_) => Some(Range::UnboundedR(1)),
                            },
                            Err(_) => Some(Range::UnboundedR(1)),
                        }
                    }
                    JectivityAttr::None => {
                        let domain_length = domain.length_signed();
                        match domain_length {
                            Ok(len) => Some(Range::Bounded(0, len)),
                            Err(_) => Some(Range::UnboundedR(0)),
                        }
                    }
                };

                let size = match attr_size {
                    Some(attr_size) => {
                        let unsafe_range = Range::minimal(&[size_size, attr_size]);
                        match unsafe_range {
                            Ok(range) => range,
                            Err(_) => {
                                return Some(Domain::empty(ReturnType::Set(Box::new(
                                    domain.return_type(),
                                ))));
                            }
                        }
                    }
                    None => size_size,
                };
                Some(Domain::set(SetAttr::new(size), domain))
            }
            Expression::Restrict(_, function, new_domain) => {
                let mut domain = function.domain_of()?;
                let (attrs_mut, dom, codom_mut) = domain.as_function_mut()?;

                // Stops other references being mutable
                let attrs: &FuncAttr<IntVal> = attrs_mut;
                let codom: &Moo<Domain> = codom_mut;

                // Gets the minimal range between the old domain and new domain
                let mut new_dom = new_domain.domain_of()?;
                // If domains cannot be resolved we just stick to the restricted one
                if let Some(new_rng) = new_dom.as_int_ground_mut()
                    && let Some(old_rng) = dom.as_int_ground_mut()
                {
                    new_rng.append(old_rng);
                    if let Ok(rng) = Range::minimal(new_rng) {
                        let ranges = vec![rng];
                        new_dom = Domain::int(ranges);
                    }
                }
                let attr_size = attrs.resolve().ok()?.size;
                let new_size = match new_dom.length_signed() {
                    // Combines current size attributes with length of new domain
                    Ok(len) => match Range::minimal(&[attr_size, Range::Bounded(0, len)]) {
                        Ok(size) => size,
                        Err(_) => {
                            // Means the restriction is impossible
                            return Some(Domain::empty(ReturnType::Function(
                                Box::new(new_dom.return_type()),
                                Box::new(codom.return_type()),
                            )));
                        }
                    },
                    Err(_) => attr_size,
                };
                let jectivity = attrs.jectivity.clone();
                let partiality = attrs.partiality.clone();
                let new_attrs = FuncAttr {
                    size: new_size,
                    jectivity,
                    partiality,
                };
                Some(Domain::function(new_attrs, new_dom, codom.clone()))
            }
            Expression::Subsequence(_, _, _) => Some(Domain::bool()),
            Expression::Substring(_, _, _) => Some(Domain::bool()),
            Expression::Inverse(..) => Some(Domain::bool()),
            Expression::LexLt(..) => Some(Domain::bool()),
            Expression::LexLeq(..) => Some(Domain::bool()),
            Expression::LexGt(..) => Some(Domain::bool()),
            Expression::LexGeq(..) => Some(Domain::bool()),
            Expression::FlatLexLt(..) => Some(Domain::bool()),
            Expression::FlatLexLeq(..) => Some(Domain::bool()),
            Expression::Active(..) => Some(Domain::bool()),
            Expression::ToSet(_, other) => {
                if let Some((attrs, dom, codom)) = other.domain_of()?.as_function() {
                    let set_attrs = SetAttr { size: attrs.size };
                    Some(Domain::set(set_attrs, Domain::tuple(vec![dom, codom])))
                } else if let Some((attrs, doms)) = other.domain_of()?.as_relation() {
                    let set_attrs = SetAttr { size: attrs.size };
                    Some(Domain::set(set_attrs, Domain::tuple(doms)))
                } else if let Some((attrs, dom)) = other.domain_of()?.as_mset() {
                    let set_attrs = SetAttr { size: attrs.size };
                    Some(Domain::set(set_attrs, dom))
                } else if let Some((dom, dimensions)) = other.domain_of()?.as_matrix() {
                    // We combine all matrix domains into a tuple
                    let mut doms = vec![];
                    for _ in dimensions {
                        doms.push(dom.clone());
                    }
                    let doms_sizes: Result<Vec<i32>, _> =
                        doms.iter().map(|x| x.length_signed()).collect();
                    let attr = match doms_sizes {
                        Ok(vals) => {
                            if let Some(&size) = vals.iter().min() {
                                SetAttr::new(Range::Single(size))
                            } else {
                                SetAttr::<i32>::default()
                            }
                        }
                        // We do not know the ground dimensions yet so default is chosen
                        Err(_) => SetAttr::<i32>::default(),
                    };
                    Some(Domain::set(attr, Domain::tuple(doms)))
                } else {
                    bug!(
                        "Domain of {self} needed to be a function, relation, mset, or matrix for ToSet"
                    )
                }
            }
            Expression::ToMSet(_, other) => {
                if let Some((attrs, dom, codom)) = other.domain_of()?.as_function() {
                    let set_attrs = MSetAttr {
                        size: attrs.size,
                        occurrence: Range::Single(IntVal::Const(1)),
                    };
                    Some(Domain::mset(set_attrs, Domain::tuple(vec![dom, codom])))
                } else if let Some((attrs, doms)) = other.domain_of()?.as_relation() {
                    let set_attrs = MSetAttr {
                        size: attrs.size,
                        occurrence: Range::Single(IntVal::Const(1)),
                    };
                    Some(Domain::mset(set_attrs, Domain::tuple(doms)))
                } else if let Some((attrs, dom)) = other.domain_of()?.as_set() {
                    let set_attrs = MSetAttr {
                        size: attrs.size,
                        occurrence: Range::Single(IntVal::Const(1)),
                    };
                    Some(Domain::mset(set_attrs, dom))
                } else {
                    bug!("Domain of {self} needed to be a function, relation, or set for ToMSet")
                }
            }
            Expression::ToRelation(_, function) => {
                let (attrs, domain, codomain) = function.domain_of()?.as_function()?;
                // Function attributes apply to the relation
                let rel_attrs = RelAttr {
                    size: attrs.size,
                    binary: vec![],
                };
                Some(Domain::relation(rel_attrs, vec![domain, codomain]))
            }
            Expression::RelationProj(_, relation, projections) => {
                let (_, domains) = relation.domain_of()?.as_relation()?;
                let new_doms = domains
                    .iter()
                    .zip(projections.iter())
                    .filter_map(|(domain, included)| {
                        if included.is_none() {
                            // The domains corresponding to projections which are None remain in the relation
                            Some(domain.clone())
                        } else {
                            None
                        }
                    })
                    .collect();
                Some(Domain::relation(RelAttr::<IntVal>::default(), new_doms))
            }
            Expression::Apart(_, _, _) => Some(Domain::bool()),
            Expression::Together(_, _, _) => Some(Domain::bool()),
            Expression::Participants(_, p) => {
                // Every single element of the domain _must_ be in the set, so fixed size on that.
                let (attr, inner) = p.domain_of()?.as_partition()?;
                let len = inner.length_signed().ok()?;

                let p_parts = attr.resolve().ok()?.num_parts;
                let p_card = attr.resolve().ok()?.part_len;

                // if
                match (p_parts.low(), p_parts.high(), p_card.low(), p_card.high()) {
                    (Some(p), Some(q), Some(r), Some(s)) => {
                        let lo = p * r;
                        let hi = q * s;
                        if len < lo || len > hi {
                            return Some(Domain::empty(ReturnType::Set(Box::new(
                                inner.return_type(),
                            ))));
                        }
                    }

                    (None, Some(q), None, Some(s)) => {
                        let hi = q * s;
                        if len > hi {
                            return Some(Domain::empty(ReturnType::Set(Box::new(
                                inner.return_type(),
                            ))));
                        }
                    }

                    (Some(p), None, Some(r), None) => {
                        let lo = p * r;
                        if len < lo {
                            return Some(Domain::empty(ReturnType::Set(Box::new(
                                inner.return_type(),
                            ))));
                        }
                    }

                    _ => {}
                }

                Some(Domain::set(
                    SetAttr::new_size(len),
                    Domain::int(inner.as_int()?),
                ))
            }
            Expression::Party(_, _, p) => {
                // Will pick a part, so set will share same attrs
                let (attr, inner) = p.domain_of()?.as_partition()?;

                Some(Domain::set(SetAttr::new(attr.part_len), inner))
            }
            Expression::Parts(_, p) => {
                let (attr, inner) = p.domain_of()?.as_partition()?;

                Some(Domain::set(
                    SetAttr::new(attr.num_parts.clone()),
                    Domain::set(SetAttr::new(attr.part_len), inner),
                ))
            }
            Expression::Card(_, collection) => {
                let domain = collection.domain_of()?;
                if let Some((_, dimensions)) = domain.as_matrix() {
                    let doms_ground: Result<Vec<i32>, _> =
                        dimensions.iter().map(|x| x.length_signed()).collect();
                    if let Ok(doms_ground) = doms_ground {
                        let size: Range<i32> = Range::Single(doms_ground.iter().product());
                        Some(Domain::int(vec![size]))
                    } else {
                        Some(Domain::int(vec![Range::<i32>::Unbounded]))
                    }
                } else if let Some((attr, dom)) = domain.as_set() {
                    let attr_size = attr.resolve().ok()?.size;
                    if let Ok(length) = dom.length_signed() {
                        let unsafe_range = Range::minimal(&[attr_size, Range::Bounded(0, length)]);
                        return match unsafe_range {
                            Ok(range) => Some(Domain::int(vec![range])),
                            Err(_) => None,
                        };
                    }
                    // If the domain is not known we just need to go off of attributes
                    Some(Domain::int(vec![attr_size]))
                } else if let Some((attrs, dom)) = domain.as_mset() {
                    let attrs_gd = attrs.resolve().ok()?;
                    // Gets maximum value of the occurrence
                    let attr_occ = match attrs_gd.occurrence {
                        Range::Single(x) => Some(x),
                        Range::Unbounded | Range::UnboundedR(_) => None,
                        Range::Bounded(_, x) => Some(x),
                        Range::UnboundedL(x) => Some(x),
                    };
                    if let Some(occ) = attr_occ {
                        if let Ok(length) = dom.length_signed() {
                            let unsafe_range =
                                Range::minimal(&[attrs_gd.size, Range::Bounded(0, length * occ)]);
                            match unsafe_range {
                                Ok(range) => Some(Domain::int(vec![range])),
                                Err(_) => None,
                            }
                        } else {
                            // If the domain is not known we just need to go off of attributes
                            Some(Domain::int(vec![attrs_gd.size]))
                        }
                    } else {
                        // If no occurrence is provided then it must have bounded size
                        Some(Domain::int(vec![attrs_gd.size]))
                    }
                } else if let Some((attrs, doms)) = domain.as_relation() {
                    // TODO: Further inference may be possible using the binary attributes

                    let attrs_gd = attrs.resolve().ok()?;
                    // See if all domains are ground
                    let doms_sizes: Result<Vec<i32>, _> =
                        doms.iter().map(|x| x.length_signed()).collect();
                    if let Ok(doms_sizes) = doms_sizes {
                        let length = Range::Bounded(0, doms_sizes.iter().product());
                        // Combine the attributes and the domain possibilities
                        let unsafe_range = Range::minimal(&[attrs_gd.size, length]);
                        return match unsafe_range {
                            Ok(range) => Some(Domain::int(vec![range])),
                            Err(_) => None,
                        };
                    }
                    // If the domain is not known we just need to go off of attributes
                    Some(Domain::int(vec![attrs_gd.size]))
                } else if let Some((attrs, dom, codom)) = domain.as_function() {
                    let size = Self::function_elements_size(attrs, &dom, &codom);
                    size.map(|size| Domain::int(vec![size]))
                } else {
                    bug!(
                        "Domain of {self} needed to be a matrix, set, mset, relation, or function for cardinality"
                    )
                }
            }
        }
    }

    // Gets the number of domain elements mapped in a function. This is the cardinality and also the defined
    fn function_elements_size(
        attrs: FuncAttr<IntVal>,
        domain: &DomainPtr,
        codomain: &DomainPtr,
    ) -> Option<Range> {
        let attrs_gd = attrs.resolve().ok()?;
        let domain_length = domain.length_signed();
        // We can only infer if the domain is ground and the length is known
        let attr_size = match domain_length {
            Ok(len) => match attrs_gd.partiality {
                PartialityAttr::Total => Some(Range::Single(len)),
                PartialityAttr::Partial => {
                    // When partial we also need the codomain to be ground and known
                    let codomain_length = codomain.length_signed();
                    match codomain_length {
                        Ok(co_len) => match attrs_gd.jectivity {
                            JectivityAttr::Bijective => Some(Range::Single(co_len)),
                            JectivityAttr::Surjective => Some(Range::Bounded(co_len, len)),
                            JectivityAttr::Injective => {
                                Some(Range::Bounded(0, Ord::min(len, co_len)))
                            }
                            JectivityAttr::None => Some(Range::Bounded(0, len)),
                        },
                        Err(_) => None,
                    }
                }
            },
            Err(_) => None,
        };
        // We combine the sizes:
        // attrs_gd.size relates to size constraints imposed by the size attributes of the function
        // attr_size relates to size constraints imposed by the jectivity and partiality attributes.
        //       This uses inference from the domain and codomain lengths.
        // If the attributes clash the function is unsolveable, and an empty domain is returned
        match attr_size {
            Some(attr_size) => {
                let unsafe_range = Range::minimal(&[attrs_gd.size, attr_size]);
                unsafe_range.ok()
            }
            None => Some(attrs_gd.size),
        }
    }

    /// Returns a reference to this expression's metadata without cloning.
    pub fn meta_ref(&self) -> &Metadata {
        macro_rules! match_meta_ref {
            ($($variant:ident),* $(,)?) => {
                match self {
                    $(Expression::$variant(meta, ..) => meta,)*
                }
            };
        }
        match_meta_ref!(
            AbstractLiteral,
            Root,
            Bubble,
            Comprehension,
            AbstractComprehension,
            DominanceRelation,
            FromSolution,
            Metavar,
            Atomic,
            RecordField,
            UnsafeIndex,
            SafeIndex,
            UnsafeSlice,
            SafeSlice,
            InDomain,
            ToInt,
            Abs,
            Sum,
            Product,
            Min,
            Max,
            Not,
            Or,
            And,
            Imply,
            Iff,
            Union,
            In,
            Intersect,
            Supset,
            SupsetEq,
            Subset,
            SubsetEq,
            Eq,
            Neq,
            Geq,
            Leq,
            Gt,
            Lt,
            SafeDiv,
            UnsafeDiv,
            SafeMod,
            UnsafeMod,
            Apart,
            Together,
            Participants,
            Party,
            Parts,
            Neg,
            Defined,
            Range,
            UnsafePow,
            SafePow,
            Flatten,
            AllDiff,
            Minus,
            Factorial,
            FlatAbsEq,
            FlatAllDiff,
            FlatSumGeq,
            FlatSumLeq,
            FlatIneq,
            FlatWatchedLiteral,
            FlatWeightedSumLeq,
            FlatWeightedSumGeq,
            FlatMinusEq,
            FlatProductEq,
            MinionDivEqUndefZero,
            MinionModuloEqUndefZero,
            MinionPow,
            MinionReify,
            MinionReifyImply,
            MinionWInIntervalSet,
            MinionWInSet,
            MinionElementOne,
            AuxDeclaration,
            SATInt,
            PairwiseSum,
            PairwiseProduct,
            Image,
            ImageSet,
            PreImage,
            Inverse,
            Restrict,
            LexLt,
            LexLeq,
            LexGt,
            LexGeq,
            FlatLexLt,
            FlatLexLeq,
            NegativeTable,
            Table,
            Active,
            ToSet,
            ToMSet,
            ToRelation,
            RelationProj,
            Card,
            Subsequence,
            Substring,
        )
    }

    pub fn get_meta(&self) -> Metadata {
        let metas: VecDeque<Metadata> = self.children_bi();
        metas[0].clone()
    }

    pub fn set_meta(&self, meta: Metadata) {
        self.transform_bi(&|_| meta.clone());
    }

    /// Checks whether this expression is safe.
    ///
    /// An expression is unsafe if can be undefined, or if any of its children can be undefined.
    ///
    /// Unsafe expressions are (typically) prefixed with Unsafe in our AST, and can be made
    /// safe through the use of bubble rules.
    pub fn is_safe(&self) -> bool {
        // TODO: memoise in Metadata
        for expr in self.universe() {
            match expr {
                Expression::UnsafeDiv(_, _, _)
                | Expression::UnsafeMod(_, _, _)
                | Expression::UnsafePow(_, _, _)
                | Expression::UnsafeIndex(_, _, _)
                | Expression::Bubble(_, _, _)
                | Expression::UnsafeSlice(_, _, _) => {
                    return false;
                }
                _ => {}
            }
        }
        true
    }

    /// True if the expression is an associative and commutative operator
    pub fn is_associative_commutative_operator(&self) -> bool {
        TryInto::<ACOperatorKind>::try_into(self).is_ok()
    }

    /// True if the expression is a matrix literal.
    ///
    /// This is true for both forms of matrix literals: those with elements of type [`Literal`] and
    /// [`Expression`].
    pub fn is_matrix_literal(&self) -> bool {
        matches!(
            self,
            Expression::AbstractLiteral(_, AbstractLiteral::Matrix(_, _))
                | Expression::Atomic(
                    _,
                    Atom::Literal(Literal::AbstractLiteral(AbstractLiteral::Matrix(_, _))),
                )
        )
    }

    /// True iff self and other are both atomic and identical.
    ///
    /// This method is useful to cheaply check equivalence. Assuming CSE is enabled, any unifiable
    /// expressions will be rewritten to a common variable. This is much cheaper than checking the
    /// entire subtrees of `self` and `other`.
    pub fn identical_atom_to(&self, other: &Expression) -> bool {
        let atom1: Result<&Atom, _> = self.try_into();
        let atom2: Result<&Atom, _> = other.try_into();

        if let (Ok(atom1), Ok(atom2)) = (atom1, atom2) {
            atom2 == atom1
        } else {
            false
        }
    }

    /// If the expression is a list, returns a *copied* vector of the inner expressions.
    ///
    /// A list is any a matrix with the domain `int(1..)`. This includes matrix literals without
    /// any explicitly specified domain.
    pub fn unwrap_list(&self) -> Option<Vec<Expression>> {
        match self {
            Expression::AbstractLiteral(_, matrix @ AbstractLiteral::Matrix(_, _)) => {
                matrix.unwrap_list().cloned()
            }
            Expression::Atomic(
                _,
                Atom::Literal(Literal::AbstractLiteral(matrix @ AbstractLiteral::Matrix(_, _))),
            ) => matrix.unwrap_list().map(|elems| {
                elems
                    .clone()
                    .into_iter()
                    .map(|x: Literal| Expression::Atomic(Metadata::new(), Atom::Literal(x)))
                    .collect_vec()
            }),
            _ => None,
        }
    }

    /// If the expression is a matrix, gets it elements and index domain.
    ///
    /// **Consider using the safer [`Expression::unwrap_list`] instead.**
    ///
    /// It is generally undefined to edit the length of a matrix unless it is a list (as defined by
    /// [`Expression::unwrap_list`]). Users of this function should ensure that, if the matrix is
    /// reconstructed, the index domain and the number of elements in the matrix remain the same.
    pub fn unwrap_matrix_unchecked(self) -> Option<(Vec<Expression>, DomainPtr)> {
        match self {
            Expression::AbstractLiteral(_, AbstractLiteral::Matrix(elems, domain)) => {
                Some((elems, domain))
            }
            Expression::Atomic(
                _,
                Atom::Literal(Literal::AbstractLiteral(AbstractLiteral::Matrix(elems, domain))),
            ) => Some((
                elems
                    .into_iter()
                    .map(|x: Literal| Expression::Atomic(Metadata::new(), Atom::Literal(x)))
                    .collect_vec(),
                domain.into(),
            )),

            _ => None,
        }
    }

    /// For a Root expression, extends the inner vec with the given vec.
    ///
    /// # Panics
    /// Panics if the expression is not Root.
    pub fn extend_root(self, exprs: Vec<Expression>) -> Expression {
        match self {
            Expression::Root(meta, mut children) => {
                children.extend(exprs);
                Expression::Root(meta, children)
            }
            _ => panic!("extend_root called on a non-Root expression"),
        }
    }

    /// Converts the expression to a literal, if possible.
    pub fn into_literal(self) -> Option<Literal> {
        match self {
            Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
            Expression::AbstractLiteral(_, abslit) => {
                Some(Literal::AbstractLiteral(abslit.into_literals()?))
            }
            Expression::Neg(_, e) => {
                let Literal::Int(i) = Moo::unwrap_or_clone(e).into_literal()? else {
                    bug!("negated literal should be an int");
                };

                Some(Literal::Int(-i))
            }

            _ => None,
        }
    }

    /// If this expression is an associative-commutative operator, return its [ACOperatorKind].
    pub fn to_ac_operator_kind(&self) -> Option<ACOperatorKind> {
        TryFrom::try_from(self).ok()
    }

    /// Returns the categories of all sub-expressions of self.
    pub fn universe_categories(&self) -> HashSet<Category> {
        self.universe()
            .into_iter()
            .map(|x| x.category_of())
            .collect()
    }
}

pub fn get_function_codomain(function: &Moo<Expression>) -> Option<DomainPtr> {
    let function_domain = function.domain_of()?;
    match function_domain.resolve().as_ref() {
        Ok(d) => {
            match d.as_ref() {
                GroundDomain::Function(_, _, codomain) => Some(codomain.clone().into()),
                // Not defined for anything other than a function
                _ => None,
            }
        }
        Err(_) => {
            match function_domain.as_unresolved()? {
                UnresolvedDomain::Function(_, _, codomain) => Some(codomain.clone()),
                // Not defined for anything other than a function
                _ => None,
            }
        }
    }
}

impl TryFrom<&Expression> for i32 {
    type Error = ();

    fn try_from(value: &Expression) -> Result<Self, Self::Error> {
        let Expression::Atomic(_, atom) = value else {
            return Err(());
        };

        let Atom::Literal(lit) = atom else {
            return Err(());
        };

        let Literal::Int(i) = lit else {
            return Err(());
        };

        Ok(*i)
    }
}

impl TryFrom<Expression> for i32 {
    type Error = ();

    fn try_from(value: Expression) -> Result<Self, Self::Error> {
        TryFrom::<&Expression>::try_from(&value)
    }
}
impl From<i32> for Expression {
    fn from(i: i32) -> Self {
        Expression::Atomic(Metadata::new(), Atom::Literal(Literal::Int(i)))
    }
}

impl From<bool> for Expression {
    fn from(b: bool) -> Self {
        Expression::Atomic(Metadata::new(), Atom::Literal(Literal::Bool(b)))
    }
}

impl From<Atom> for Expression {
    fn from(value: Atom) -> Self {
        Expression::Atomic(Metadata::new(), value)
    }
}

impl From<Literal> for Expression {
    fn from(value: Literal) -> Self {
        Expression::Atomic(Metadata::new(), value.into())
    }
}

impl From<AbstractLiteral<Expression>> for Expression {
    fn from(value: AbstractLiteral<Expression>) -> Self {
        Expression::AbstractLiteral(Metadata::new(), value)
    }
}

impl From<Moo<Expression>> for Expression {
    fn from(val: Moo<Expression>) -> Self {
        val.as_ref().clone()
    }
}

impl CategoryOf for Expression {
    fn category_of(&self) -> Category {
        // take highest category of all the expressions children
        let category = self.cata(&move |x,children| {

            if let Some(max_category) = children.iter().max() {
                // if this expression contains subexpressions, return the maximum category of the
                // subexpressions
                *max_category
            } else {
                // this expression has no children
                let mut max_category = Category::Bottom;

                // calculate the category by looking at all atoms, submodels, comprehensions, and
                // declarationptrs inside this expression

                // this should generically cover all leaf types we currently have in oxide.

                // if x contains submodels (including comprehensions)
                if !Biplate::<Model>::universe_bi(&x).is_empty() {
                    // assume that the category is decision
                    return Category::Decision;
                }

                // if x contains atoms
                if let Some(max_atom_category) = Biplate::<Atom>::universe_bi(&x).iter().map(|x| x.category_of()).max()
                // and those atoms have a higher category than we already know about
                && max_atom_category > max_category{
                    // update category
                    max_category = max_atom_category;
                }

                // if x contains declarationPtrs
                if let Some(max_declaration_category) = Biplate::<DeclarationPtr>::universe_bi(&x).iter().map(|x| x.category_of()).max()
                // and those pointers have a higher category than we already know about
                && max_declaration_category > max_category{
                    // update category
                    max_category = max_declaration_category;
                }
                max_category

            }
        });

        if cfg!(debug_assertions) {
            trace!(
                category= %category,
                expression= %self,
                "Called Expression::category_of()"
            );
        };
        category
    }
}

impl Display for Expression {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match &self {
            Expression::Union(_, box1, box2) => {
                write!(f, "({} union {})", box1.clone(), box2.clone())
            }
            Expression::In(_, e1, e2) => {
                write!(f, "{e1} in {e2}")
            }
            Expression::Intersect(_, box1, box2) => {
                write!(f, "({} intersect {})", box1.clone(), box2.clone())
            }
            Expression::Supset(_, box1, box2) => {
                write!(f, "({} supset {})", box1.clone(), box2.clone())
            }
            Expression::SupsetEq(_, box1, box2) => {
                write!(f, "({} supsetEq {})", box1.clone(), box2.clone())
            }
            Expression::Subset(_, box1, box2) => {
                write!(f, "({} subset {})", box1.clone(), box2.clone())
            }
            Expression::SubsetEq(_, box1, box2) => {
                write!(f, "({} subsetEq {})", box1.clone(), box2.clone())
            }

            Expression::AbstractLiteral(_, l) => l.fmt(f),
            Expression::Comprehension(_, c) => c.fmt(f),
            Expression::AbstractComprehension(_, c) => c.fmt(f),
            Expression::UnsafeIndex(_, e1, e2) => write!(f, "{e1}{}", pretty_vec(e2)),
            Expression::RecordField(_, r, fld) => {
                write!(f, "{r}[{fld}]")
            }
            Expression::SafeIndex(_, e1, e2) => write!(f, "SafeIndex({e1},{})", pretty_vec(e2)),
            Expression::UnsafeSlice(_, e1, es) => {
                let args = es
                    .iter()
                    .map(|x| match x {
                        Some(x) => format!("{x}"),
                        None => "..".into(),
                    })
                    .join(",");

                write!(f, "{e1}[{args}]")
            }
            Expression::SafeSlice(_, e1, es) => {
                let args = es
                    .iter()
                    .map(|x| match x {
                        Some(x) => format!("{x}"),
                        None => "..".into(),
                    })
                    .join(",");

                write!(f, "SafeSlice({e1},[{args}])")
            }
            Expression::InDomain(_, e, domain) => {
                write!(f, "__inDomain({e},{domain})")
            }
            Expression::Root(_, exprs) => {
                write!(f, "{}", pretty_expressions_as_top_level(exprs))
            }
            Expression::DominanceRelation(_, expr) => write!(f, "DominanceRelation({expr})"),
            Expression::FromSolution(_, expr) => write!(f, "FromSolution({expr})"),
            Expression::Metavar(_, name) => write!(f, "&{name}"),
            Expression::Atomic(_, atom) => atom.fmt(f),
            Expression::Abs(_, a) | Expression::Card(_, a) => write!(f, "|{a}|"),
            Expression::Sum(_, e) => {
                write!(f, "sum({e})")
            }
            Expression::Product(_, e) => {
                write!(f, "product({e})")
            }
            Expression::Min(_, e) => {
                write!(f, "min({e})")
            }
            Expression::Max(_, e) => {
                write!(f, "max({e})")
            }
            Expression::Not(_, expr_box) => {
                write!(f, "!({})", expr_box.clone())
            }
            Expression::Or(_, e) => {
                write!(f, "or({e})")
            }
            Expression::And(_, e) => {
                write!(f, "and({e})")
            }
            Expression::Imply(_, box1, box2) => {
                write!(f, "({box1}) -> ({box2})")
            }
            Expression::Iff(_, box1, box2) => {
                write!(f, "({box1}) <-> ({box2})")
            }
            Expression::Eq(_, box1, box2) => {
                write!(f, "({} = {})", box1.clone(), box2.clone())
            }
            Expression::Neq(_, box1, box2) => {
                write!(f, "({} != {})", box1.clone(), box2.clone())
            }
            Expression::Geq(_, box1, box2) => {
                write!(f, "({} >= {})", box1.clone(), box2.clone())
            }
            Expression::Leq(_, box1, box2) => {
                write!(f, "({} <= {})", box1.clone(), box2.clone())
            }
            Expression::Gt(_, box1, box2) => {
                write!(f, "({} > {})", box1.clone(), box2.clone())
            }
            Expression::Lt(_, box1, box2) => {
                write!(f, "({} < {})", box1.clone(), box2.clone())
            }
            Expression::Apart(_, list, partition) => {
                write!(f, "apart({list}, {partition})")
            }
            Expression::Together(_, list, partition) => {
                write!(f, "together({list}, {partition})")
            }
            Expression::Participants(_, partition) => {
                write!(f, "participants({partition})")
            }
            Expression::Party(_, element, partition) => {
                write!(f, "party({element}, {partition})")
            }
            Expression::Parts(_, partition) => {
                write!(f, "parts({partition})")
            }
            Expression::FlatSumGeq(_, box1, box2) => {
                write!(f, "SumGeq({}, {})", pretty_vec(box1), box2.clone())
            }
            Expression::FlatSumLeq(_, box1, box2) => {
                write!(f, "SumLeq({}, {})", pretty_vec(box1), box2.clone())
            }
            Expression::FlatIneq(_, box1, box2, box3) => write!(
                f,
                "Ineq({}, {}, {})",
                box1.clone(),
                box2.clone(),
                box3.clone()
            ),
            Expression::Flatten(_, n, m) => {
                if let Some(n) = n {
                    write!(f, "flatten({n}, {m})")
                } else {
                    write!(f, "flatten({m})")
                }
            }
            Expression::AllDiff(_, e) => {
                write!(f, "allDiff({e})")
            }
            Expression::Table(_, tuple_expr, rows_expr) => {
                write!(f, "table({tuple_expr}, {rows_expr})")
            }
            Expression::NegativeTable(_, tuple_expr, rows_expr) => {
                write!(f, "negativeTable({tuple_expr}, {rows_expr})")
            }
            Expression::Bubble(_, box1, box2) => {
                write!(f, "{{{} @ {}}}", box1.clone(), box2.clone())
            }
            Expression::SafeDiv(_, box1, box2) => {
                write!(f, "SafeDiv({}, {})", box1.clone(), box2.clone())
            }
            Expression::UnsafeDiv(_, box1, box2) => {
                write!(f, "({} / {})", box1.clone(), box2.clone())
            }
            Expression::UnsafePow(_, box1, box2) => {
                write!(f, "({} ** {})", box1.clone(), box2.clone())
            }
            Expression::SafePow(_, box1, box2) => {
                write!(f, "SafePow({}, {})", box1.clone(), box2.clone())
            }
            Expression::Subsequence(_, s, t) => {
                write!(f, "{} subsequence {}", s.clone(), t.clone())
            }
            Expression::Substring(_, s, t) => {
                write!(f, "{} substring {}", s.clone(), t.clone())
            }
            Expression::MinionDivEqUndefZero(_, box1, box2, box3) => {
                write!(
                    f,
                    "DivEq({}, {}, {})",
                    box1.clone(),
                    box2.clone(),
                    box3.clone()
                )
            }
            Expression::MinionModuloEqUndefZero(_, box1, box2, box3) => {
                write!(
                    f,
                    "ModEq({}, {}, {})",
                    box1.clone(),
                    box2.clone(),
                    box3.clone()
                )
            }
            Expression::FlatWatchedLiteral(_, x, l) => {
                write!(f, "WatchedLiteral({x},{l})")
            }
            Expression::MinionReify(_, box1, box2) => {
                write!(f, "Reify({}, {})", box1.clone(), box2.clone())
            }
            Expression::MinionReifyImply(_, box1, box2) => {
                write!(f, "ReifyImply({}, {})", box1.clone(), box2.clone())
            }
            Expression::MinionWInIntervalSet(_, atom, intervals) => {
                let intervals = intervals.iter().join(",");
                write!(f, "__minion_w_inintervalset({atom},[{intervals}])")
            }
            Expression::MinionWInSet(_, atom, values) => {
                let values = values.iter().join(",");
                write!(f, "__minion_w_inset({atom},[{values}])")
            }
            Expression::AuxDeclaration(_, reference, e) => {
                write!(f, "{} =aux {}", reference, e.clone())
            }
            Expression::UnsafeMod(_, a, b) => {
                write!(f, "{} % {}", a.clone(), b.clone())
            }
            Expression::SafeMod(_, a, b) => {
                write!(f, "SafeMod({},{})", a.clone(), b.clone())
            }
            Expression::Neg(_, a) => {
                write!(f, "-({})", a.clone())
            }
            Expression::Factorial(_, a) => {
                write!(f, "({})!", a.clone())
            }
            Expression::Minus(_, a, b) => {
                write!(f, "({} - {})", a.clone(), b.clone())
            }
            Expression::FlatAllDiff(_, es) => {
                write!(f, "__flat_alldiff({})", pretty_vec(es))
            }
            Expression::FlatAbsEq(_, a, b) => {
                write!(f, "AbsEq({},{})", a.clone(), b.clone())
            }
            Expression::FlatMinusEq(_, a, b) => {
                write!(f, "MinusEq({},{})", a.clone(), b.clone())
            }
            Expression::FlatProductEq(_, a, b, c) => {
                write!(
                    f,
                    "FlatProductEq({},{},{})",
                    a.clone(),
                    b.clone(),
                    c.clone()
                )
            }
            Expression::FlatWeightedSumLeq(_, cs, vs, total) => {
                write!(
                    f,
                    "FlatWeightedSumLeq({},{},{})",
                    pretty_vec(cs),
                    pretty_vec(vs),
                    total.clone()
                )
            }
            Expression::FlatWeightedSumGeq(_, cs, vs, total) => {
                write!(
                    f,
                    "FlatWeightedSumGeq({},{},{})",
                    pretty_vec(cs),
                    pretty_vec(vs),
                    total.clone()
                )
            }
            Expression::MinionPow(_, atom, atom1, atom2) => {
                write!(f, "MinionPow({atom},{atom1},{atom2})")
            }
            Expression::MinionElementOne(_, atoms, atom, atom1) => {
                let atoms = atoms.iter().join(",");
                write!(f, "__minion_element_one([{atoms}],{atom},{atom1})")
            }

            Expression::ToInt(_, expr) => {
                write!(f, "toInt({expr})")
            }

            Expression::SATInt(_, encoding, bits, (min, max)) => {
                write!(f, "SATInt({encoding:?}, {bits} [{min}, {max}])")
            }

            Expression::PairwiseSum(_, a, b) => write!(f, "PairwiseSum({a}, {b})"),
            Expression::PairwiseProduct(_, a, b) => write!(f, "PairwiseProduct({a}, {b})"),

            Expression::Defined(_, function) => write!(f, "defined({function})"),
            Expression::Range(_, function) => write!(f, "range({function})"),
            Expression::Image(_, function, elems) => write!(f, "image({function},{elems})"),
            Expression::ImageSet(_, function, elems) => write!(f, "imageSet({function},{elems})"),
            Expression::PreImage(_, function, elems) => write!(f, "preImage({function},{elems})"),
            Expression::Inverse(_, a, b) => write!(f, "inverse({a},{b})"),
            Expression::Restrict(_, function, domain) => write!(f, "restrict({function},{domain})"),

            Expression::LexLt(_, a, b) => write!(f, "({a} <lex {b})"),
            Expression::LexLeq(_, a, b) => write!(f, "({a} <=lex {b})"),
            Expression::LexGt(_, a, b) => write!(f, "({a} >lex {b})"),
            Expression::LexGeq(_, a, b) => write!(f, "({a} >=lex {b})"),
            Expression::FlatLexLt(_, a, b) => {
                write!(f, "FlatLexLt({}, {})", pretty_vec(a), pretty_vec(b))
            }
            Expression::FlatLexLeq(_, a, b) => {
                write!(f, "FlatLexLeq({}, {})", pretty_vec(a), pretty_vec(b))
            }
            Expression::Active(_, variant, field_name) => {
                write!(f, "active({variant}, {field_name})")
            }
            Expression::ToSet(_, other) => write!(f, "toSet({other})"),
            Expression::ToMSet(_, other) => write!(f, "toMSet({other})"),
            Expression::ToRelation(_, function) => write!(f, "toRelation({function})"),
            Expression::RelationProj(_, relation, projections) => {
                let projections_str = projections
                    .iter()
                    .map(|x| {
                        if let Some(x) = x {
                            x.to_string()
                        } else {
                            String::from("_")
                        }
                    })
                    .join(", ");
                write!(f, "{relation}({projections_str})")
            }
        }
    }
}

fn minus_operand_return_type(expr: &Expression) -> ReturnType {
    match expr {
        Expression::Atomic(_, Atom::Reference(reference)) => {
            let decl_kind = reference.ptr.kind().clone();
            match decl_kind {
                DeclarationKind::Find(var) => var.return_type(),
                DeclarationKind::Given(domain)
                | DeclarationKind::DomainLetting(domain) => domain.return_type(),
                DeclarationKind::Quantified(inner) => inner.domain().return_type(),
                DeclarationKind::QuantifiedExpr(inner)
                | DeclarationKind::TemporaryValueLetting(inner)
                // not sure if i should ever be looking at the domain ptr but seems to work
                | DeclarationKind::ValueLetting(inner, _) => inner.return_type(),
            }
        }
        _ => expr.return_type(),
    }
}

impl Typeable for Expression {
    fn return_type(&self) -> ReturnType {
        match self {
            Expression::Union(_, subject, _) => ReturnType::Set(Box::new(subject.return_type())),
            Expression::Intersect(_, subject, _) => {
                ReturnType::Set(Box::new(subject.return_type()))
            }
            Expression::In(_, _, _) => ReturnType::Bool,
            Expression::Supset(_, _, _) => ReturnType::Bool,
            Expression::SupsetEq(_, _, _) => ReturnType::Bool,
            Expression::Subset(_, _, _) => ReturnType::Bool,
            Expression::SubsetEq(_, _, _) => ReturnType::Bool,
            Expression::AbstractLiteral(_, lit) => lit.return_type(),
            Expression::RecordField(_, rec, field_name) => {
                if let ReturnType::Record(ents) = rec.return_type() {
                    for Field { name, value } in ents {
                        if name.eq(field_name) {
                            return value;
                        }
                    }
                }
                ReturnType::Unknown
            }
            Expression::UnsafeIndex(_, subject, idx) | Expression::SafeIndex(_, subject, idx) => {
                let subject_ty = subject.return_type();
                match subject_ty {
                    ReturnType::Matrix(_) => {
                        // For n-dimensional matrices, unwrap the element type until
                        // we either get to the innermost element type or the last index
                        let mut elem_typ = subject_ty;
                        let mut idx_len = idx.len();
                        while idx_len > 0
                            && let ReturnType::Matrix(new_elem_typ) = &elem_typ
                        {
                            elem_typ = *new_elem_typ.clone();
                            idx_len -= 1;
                        }
                        elem_typ
                    }
                    // TODO: We can implement indexing for these eventually
                    ReturnType::Record(_) | ReturnType::Tuple(_) | ReturnType::Variant(_) => {
                        ReturnType::Unknown
                    }
                    _ => bug!(
                        "Invalid indexing operation: expected the operand to be a collection, got {self}: {subject_ty}"
                    ),
                }
            }
            Expression::UnsafeSlice(_, subject, _) | Expression::SafeSlice(_, subject, _) => {
                ReturnType::Matrix(Box::new(subject.return_type()))
            }
            Expression::InDomain(_, _, _) => ReturnType::Bool,
            Expression::Comprehension(_, comp) => comp.return_type(),
            Expression::AbstractComprehension(_, comp) => comp.return_type(),
            Expression::Root(_, _) => ReturnType::Bool,
            Expression::DominanceRelation(_, _) => ReturnType::Bool,
            Expression::FromSolution(_, expr) => expr.return_type(),
            Expression::Metavar(_, _) => ReturnType::Unknown,
            Expression::Atomic(_, atom) => atom.return_type(),
            Expression::Abs(_, _) => ReturnType::Int,
            Expression::Sum(_, _) => ReturnType::Int,
            Expression::Product(_, _) => ReturnType::Int,
            Expression::Min(_, _) => ReturnType::Int,
            Expression::Max(_, _) => ReturnType::Int,
            Expression::Not(_, _) => ReturnType::Bool,
            Expression::Or(_, _) => ReturnType::Bool,
            Expression::Imply(_, _, _) => ReturnType::Bool,
            Expression::Iff(_, _, _) => ReturnType::Bool,
            Expression::And(_, _) => ReturnType::Bool,
            Expression::Eq(_, _, _) => ReturnType::Bool,
            Expression::Neq(_, _, _) => ReturnType::Bool,
            Expression::Geq(_, _, _) => ReturnType::Bool,
            Expression::Leq(_, _, _) => ReturnType::Bool,
            Expression::Gt(_, _, _) => ReturnType::Bool,
            Expression::Lt(_, _, _) => ReturnType::Bool,
            Expression::Apart(_, _, _) => ReturnType::Bool,
            Expression::Together(_, _, _) => ReturnType::Bool,
            Expression::Party(_, _, subject) => ReturnType::Set(Box::new(subject.return_type())),
            Expression::Participants(_, subject) => {
                ReturnType::Set(Box::new(subject.return_type()))
            }
            Expression::Parts(_, subject) => {
                ReturnType::Set(Box::new(ReturnType::Set(Box::new(subject.return_type()))))
            }
            Expression::SafeDiv(_, _, _) => ReturnType::Int,
            Expression::UnsafeDiv(_, _, _) => ReturnType::Int,
            Expression::FlatAllDiff(_, _) => ReturnType::Bool,
            Expression::FlatSumGeq(_, _, _) => ReturnType::Bool,
            Expression::FlatSumLeq(_, _, _) => ReturnType::Bool,
            Expression::MinionDivEqUndefZero(_, _, _, _) => ReturnType::Bool,
            Expression::FlatIneq(_, _, _, _) => ReturnType::Bool,
            Expression::Flatten(_, _, matrix) => {
                let matrix_type = matrix.return_type();
                match matrix_type {
                    ReturnType::Matrix(_) => {
                        // unwrap until we get to innermost element
                        let mut elem_type = matrix_type;
                        while let ReturnType::Matrix(new_elem_type) = &elem_type {
                            elem_type = *new_elem_type.clone();
                        }
                        ReturnType::Matrix(Box::new(elem_type))
                    }
                    _ => bug!(
                        "Invalid indexing operation: expected the operand to be a collection, got {self}: {matrix_type}"
                    ),
                }
            }
            Expression::AllDiff(_, _) => ReturnType::Bool,
            Expression::Table(_, _, _) => ReturnType::Bool,
            Expression::NegativeTable(_, _, _) => ReturnType::Bool,
            Expression::Bubble(_, inner, _) => inner.return_type(),
            Expression::FlatWatchedLiteral(_, _, _) => ReturnType::Bool,
            Expression::MinionReify(_, _, _) => ReturnType::Bool,
            Expression::MinionReifyImply(_, _, _) => ReturnType::Bool,
            Expression::MinionWInIntervalSet(_, _, _) => ReturnType::Bool,
            Expression::MinionWInSet(_, _, _) => ReturnType::Bool,
            Expression::MinionElementOne(_, _, _, _) => ReturnType::Bool,
            Expression::AuxDeclaration(_, _, _) => ReturnType::Bool,
            Expression::UnsafeMod(_, _, _) => ReturnType::Int,
            Expression::SafeMod(_, _, _) => ReturnType::Int,
            Expression::MinionModuloEqUndefZero(_, _, _, _) => ReturnType::Bool,
            Expression::Neg(_, _) => ReturnType::Int,
            Expression::Factorial(_, _) => ReturnType::Int,
            Expression::UnsafePow(_, _, _) => ReturnType::Int,
            Expression::SafePow(_, _, _) => ReturnType::Int,
            Expression::Minus(_, a, b) => {
                // rather than calling .return_type on a and b which sometimes errors on references that don't have domains
                // use custom function that extracts return type from atomic references based on each declaration variant
                let a_type = minus_operand_return_type(a);
                let b_type = minus_operand_return_type(b);

                if a_type == ReturnType::Int && b_type == ReturnType::Int {
                    ReturnType::Int
                } else if let ReturnType::Set(a_inner) = a_type
                    && let ReturnType::Set(b_inner) = b_type
                    && a_inner == b_inner
                {
                    ReturnType::Set(a_inner)
                } else {
                    bug!(
                        "Invalid minus operation: operands are of different or invalid types for this operation"
                    )
                }
            }
            Expression::FlatAbsEq(_, _, _) => ReturnType::Bool,
            Expression::FlatMinusEq(_, _, _) => ReturnType::Bool,
            Expression::FlatProductEq(_, _, _, _) => ReturnType::Bool,
            Expression::FlatWeightedSumLeq(_, _, _, _) => ReturnType::Bool,
            Expression::FlatWeightedSumGeq(_, _, _, _) => ReturnType::Bool,
            Expression::MinionPow(_, _, _, _) => ReturnType::Bool,
            Expression::ToInt(_, _) => ReturnType::Int,
            Expression::SATInt(..) => ReturnType::Int,
            Expression::PairwiseSum(_, _, _) => ReturnType::Int,
            Expression::PairwiseProduct(_, _, _) => ReturnType::Int,
            Expression::Defined(_, function) => {
                let subject = function.return_type();
                match subject {
                    ReturnType::Function(domain, _) => ReturnType::Set(Box::new(*domain)),
                    _ => bug!(
                        "Invalid defined operation: expected the operand to be a function, got {self}: {subject}"
                    ),
                }
            }
            Expression::Range(_, function) => {
                let subject = function.return_type();
                match subject {
                    ReturnType::Function(_, codomain) => ReturnType::Set(Box::new(*codomain)),
                    _ => bug!(
                        "Invalid range operation: expected the operand to be a function, got {self}: {subject}"
                    ),
                }
            }
            Expression::Image(_, function, _) => {
                let subject = function.return_type();
                match subject {
                    ReturnType::Function(_, codomain) => *codomain,
                    _ => bug!(
                        "Invalid image operation: expected the operand to be a function, got {self}: {subject}"
                    ),
                }
            }
            Expression::ImageSet(_, function, _) => {
                let subject = function.return_type();
                match subject {
                    ReturnType::Function(_, codomain) => ReturnType::Set(Box::new(*codomain)),
                    _ => bug!(
                        "Invalid imageSet operation: expected the operand to be a function, got {self}: {subject}"
                    ),
                }
            }
            Expression::PreImage(_, function, _) => {
                let subject = function.return_type();
                match subject {
                    ReturnType::Function(domain, _) => ReturnType::Set(Box::new(*domain)),
                    _ => bug!(
                        "Invalid preImage operation: expected the operand to be a function, got {self}: {subject}"
                    ),
                }
            }
            Expression::Restrict(_, function, new_domain) => {
                let subject = function.return_type();
                match subject {
                    ReturnType::Function(_, codomain) => {
                        ReturnType::Function(Box::new(new_domain.return_type()), codomain)
                    }
                    _ => bug!(
                        "Invalid preImage operation: expected the operand to be a function, got {self}: {subject}"
                    ),
                }
            }
            Expression::Inverse(..) => ReturnType::Bool,
            Expression::LexLt(..) => ReturnType::Bool,
            Expression::LexGt(..) => ReturnType::Bool,
            Expression::LexLeq(..) => ReturnType::Bool,
            Expression::LexGeq(..) => ReturnType::Bool,
            Expression::FlatLexLt(..) => ReturnType::Bool,
            Expression::FlatLexLeq(..) => ReturnType::Bool,
            Expression::Active(..) => ReturnType::Bool,
            Expression::ToSet(_, other) => {
                let subject = other.return_type();
                match subject {
                    ReturnType::Function(domain, codomain) => {
                        ReturnType::Set(Box::new(ReturnType::Tuple(vec![*domain, *codomain])))
                    }
                    ReturnType::Relation(domains) => {
                        ReturnType::Set(Box::new(ReturnType::Tuple(domains)))
                    }
                    ReturnType::MSet(domain) => ReturnType::Set(Box::new(*domain)),
                    ReturnType::Matrix(domain) => ReturnType::Set(Box::new(*domain)),
                    _ => bug!(
                        "Invalid toSet operation: expected the operand to be a mset, matrix, relation, or function, got {self}: {subject}"
                    ),
                }
            }
            Expression::ToMSet(_, other) => {
                let subject = other.return_type();
                match subject {
                    ReturnType::Function(domain, codomain) => {
                        ReturnType::MSet(Box::new(ReturnType::Tuple(vec![*domain, *codomain])))
                    }
                    ReturnType::Relation(domains) => {
                        ReturnType::MSet(Box::new(ReturnType::Tuple(domains)))
                    }
                    ReturnType::Set(domain) => ReturnType::MSet(Box::new(*domain)),
                    _ => bug!(
                        "Invalid toMSet operation: expected the operand to be a set, relation, or function, got {self}: {subject}"
                    ),
                }
            }
            Expression::ToRelation(_, function) => {
                let subject = function.return_type();
                match subject {
                    ReturnType::Function(domain, codomain) => {
                        ReturnType::Relation(vec![*domain, *codomain])
                    }
                    _ => bug!(
                        "Invalid toRelation operation: expected the operand to be a function, got {self}: {subject}"
                    ),
                }
            }
            Expression::RelationProj(_, relation, projections) => {
                let subject = relation.return_type();
                match subject {
                    ReturnType::Relation(domains) => {
                        let new_doms = domains
                            .iter()
                            .zip(projections.iter())
                            .filter_map(|(domain, included)| {
                                if included.is_none() {
                                    // The domains corresponding to projections which are None remain in the relation
                                    Some(domain.clone())
                                } else {
                                    None
                                }
                            })
                            .collect();
                        ReturnType::Relation(new_doms)
                    }
                    _ => bug!(
                        "Invalid RelationProj operation: expected the operand to be a relation, got {self}: {subject}"
                    ),
                }
            }
            Expression::Card(..) => ReturnType::Int,
            Expression::Subsequence(_, _, _) => ReturnType::Bool,
            Expression::Substring(_, _, _) => ReturnType::Bool,
        }
    }
}

impl Expression {
    /// Visit each direct `Expression` child by reference, without cloning.
    fn for_each_expr_child(&self, f: &mut impl FnMut(&Expression)) {
        match self {
            // Special Case
            Expression::AbstractLiteral(_, alit) => match alit {
                AbstractLiteral::Set(v) | AbstractLiteral::MSet(v) | AbstractLiteral::Tuple(v) => {
                    for expr in v {
                        f(expr);
                    }
                }
                AbstractLiteral::Partition(two_d_v) => {
                    for part in two_d_v {
                        for expr in part {
                            f(expr);
                        }
                    }
                }
                AbstractLiteral::Matrix(v, _domain) => {
                    for expr in v {
                        f(expr);
                    }
                }
                AbstractLiteral::Record(rs) => {
                    for r in rs {
                        f(&r.value);
                    }
                }
                AbstractLiteral::Sequence(v) => {
                    for expr in v {
                        f(expr);
                    }
                }
                AbstractLiteral::Function(vs) => {
                    for (a, b) in vs {
                        f(a);
                        f(b);
                    }
                }
                AbstractLiteral::Variant(v) => {
                    f(&v.value);
                }
                AbstractLiteral::Relation(vs) => {
                    for exprs in vs {
                        for expr in exprs {
                            f(expr);
                        }
                    }
                }
            },
            Expression::Root(_, vs) => {
                for expr in vs {
                    f(expr);
                }
            }

            // Moo<Expression>
            Expression::DominanceRelation(_, m1)
            | Expression::ToInt(_, m1)
            | Expression::Abs(_, m1)
            | Expression::Sum(_, m1)
            | Expression::Product(_, m1)
            | Expression::Min(_, m1)
            | Expression::Max(_, m1)
            | Expression::Not(_, m1)
            | Expression::Or(_, m1)
            | Expression::And(_, m1)
            | Expression::Neg(_, m1)
            | Expression::Defined(_, m1)
            | Expression::AllDiff(_, m1)
            | Expression::Factorial(_, m1)
            | Expression::Range(_, m1)
            | Expression::Participants(_, m1)
            | Expression::Parts(_, m1)
            | Expression::ToSet(_, m1)
            | Expression::ToMSet(_, m1)
            | Expression::ToRelation(_, m1)
            | Expression::Card(_, m1)
            | Expression::RecordField(_, m1, _)
            | Expression::Active(_, m1, _) => {
                f(m1);
            }

            // Moo<Expression> + Moo<Expression>
            Expression::Table(_, m1, m2)
            | Expression::NegativeTable(_, m1, m2)
            | Expression::Bubble(_, m1, m2)
            | Expression::Imply(_, m1, m2)
            | Expression::Iff(_, m1, m2)
            | Expression::Union(_, m1, m2)
            | Expression::In(_, m1, m2)
            | Expression::Intersect(_, m1, m2)
            | Expression::Supset(_, m1, m2)
            | Expression::SupsetEq(_, m1, m2)
            | Expression::Subset(_, m1, m2)
            | Expression::SubsetEq(_, m1, m2)
            | Expression::Eq(_, m1, m2)
            | Expression::Neq(_, m1, m2)
            | Expression::Geq(_, m1, m2)
            | Expression::Leq(_, m1, m2)
            | Expression::Gt(_, m1, m2)
            | Expression::Lt(_, m1, m2)
            | Expression::SafeDiv(_, m1, m2)
            | Expression::UnsafeDiv(_, m1, m2)
            | Expression::SafeMod(_, m1, m2)
            | Expression::UnsafeMod(_, m1, m2)
            | Expression::UnsafePow(_, m1, m2)
            | Expression::SafePow(_, m1, m2)
            | Expression::Minus(_, m1, m2)
            | Expression::PairwiseSum(_, m1, m2)
            | Expression::PairwiseProduct(_, m1, m2)
            | Expression::Image(_, m1, m2)
            | Expression::ImageSet(_, m1, m2)
            | Expression::PreImage(_, m1, m2)
            | Expression::Inverse(_, m1, m2)
            | Expression::Restrict(_, m1, m2)
            | Expression::Apart(_, m1, m2)
            | Expression::Together(_, m1, m2)
            | Expression::Party(_, m1, m2)
            | Expression::LexLt(_, m1, m2)
            | Expression::LexLeq(_, m1, m2)
            | Expression::LexGt(_, m1, m2)
            | Expression::LexGeq(_, m1, m2)
            | Expression::Subsequence(_, m1, m2)
            | Expression::Substring(_, m1, m2) => {
                f(m1);
                f(m2);
            }

            // Moo<Expression> + Vec<Expression>
            Expression::UnsafeIndex(_, m, vs) | Expression::SafeIndex(_, m, vs) => {
                f(m);
                for v in vs {
                    f(v);
                }
            }
            // Moo<Expression> + Vec<Option<Expression>>
            Expression::UnsafeSlice(_, m, vs)
            | Expression::SafeSlice(_, m, vs)
            | Expression::RelationProj(_, m, vs) => {
                f(m);
                for e in vs.iter().flatten() {
                    f(e);
                }
            }

            // Moo<Expression> + DomainPtr
            Expression::InDomain(_, m, _) => {
                f(m);
            }

            // Option<Moo<Expression>> + Moo<Expression>
            Expression::Flatten(_, opt, m) => {
                if let Some(e) = opt {
                    f(e);
                }
                f(m);
            }

            // Moo<Expression> + Atom
            Expression::MinionReify(_, m, _) | Expression::MinionReifyImply(_, m, _) => {
                f(m);
            }

            // Reference + Moo<Expression>
            Expression::AuxDeclaration(_, _, m) => {
                f(m);
            }

            // SATIntEncoding + Moo<Expression> + (i32, i32)
            Expression::SATInt(_, _, m, _) => {
                f(m);
            }

            // No Expression children
            Expression::Comprehension(_, _)
            | Expression::AbstractComprehension(_, _)
            | Expression::Atomic(_, _)
            | Expression::FromSolution(_, _)
            | Expression::Metavar(_, _)
            | Expression::FlatAbsEq(_, _, _)
            | Expression::FlatMinusEq(_, _, _)
            | Expression::FlatProductEq(_, _, _, _)
            | Expression::MinionDivEqUndefZero(_, _, _, _)
            | Expression::MinionModuloEqUndefZero(_, _, _, _)
            | Expression::MinionPow(_, _, _, _)
            | Expression::FlatAllDiff(_, _)
            | Expression::FlatSumGeq(_, _, _)
            | Expression::FlatSumLeq(_, _, _)
            | Expression::FlatIneq(_, _, _, _)
            | Expression::FlatWatchedLiteral(_, _, _)
            | Expression::FlatWeightedSumLeq(_, _, _, _)
            | Expression::FlatWeightedSumGeq(_, _, _, _)
            | Expression::MinionWInIntervalSet(_, _, _)
            | Expression::MinionWInSet(_, _, _)
            | Expression::MinionElementOne(_, _, _, _)
            | Expression::FlatLexLt(_, _, _)
            | Expression::FlatLexLeq(_, _, _) => {}
        }
    }
}

impl CacheHashable for Expression {
    fn invalidate_cache(&self) {
        self.meta_ref()
            .stored_hash
            .store(NO_HASH, Ordering::Relaxed);
    }

    fn invalidate_cache_recursive(&self) {
        self.invalidate_cache();
        self.for_each_expr_child(&mut |child| {
            child.invalidate_cache_recursive();
        });
    }

    fn get_cached_hash(&self) -> u64 {
        let stored = self.meta_ref().stored_hash.load(Ordering::Relaxed);
        if stored != NO_HASH {
            HASH_HITS.fetch_add(1, Ordering::Relaxed);
            return stored;
        }
        HASH_MISSES.fetch_add(1, Ordering::Relaxed);
        self.calculate_hash()
    }

    fn calculate_hash(&self) -> u64 {
        let mut hasher = DefaultHasher::new();
        std::mem::discriminant(self).hash(&mut hasher);
        match self {
            // Special Case
            Expression::AbstractLiteral(_, alit) => match alit {
                AbstractLiteral::Set(v)
                | AbstractLiteral::MSet(v)
                | AbstractLiteral::Tuple(v)
                | AbstractLiteral::Sequence(v) => {
                    for expr in v {
                        expr.get_cached_hash().hash(&mut hasher);
                    }
                }
                AbstractLiteral::Matrix(v, domain) => {
                    domain.hash(&mut hasher);
                    for expr in v {
                        expr.get_cached_hash().hash(&mut hasher);
                    }
                }
                AbstractLiteral::Record(rs) => {
                    for r in rs {
                        r.name.hash(&mut hasher);
                        r.value.get_cached_hash().hash(&mut hasher);
                    }
                }
                AbstractLiteral::Function(vs) => {
                    for (a, b) in vs {
                        a.get_cached_hash().hash(&mut hasher);
                        b.get_cached_hash().hash(&mut hasher);
                    }
                }
                AbstractLiteral::Variant(v) => {
                    v.name.hash(&mut hasher);
                    v.value.get_cached_hash().hash(&mut hasher);
                }
                AbstractLiteral::Relation(v) => {
                    for exprs in v {
                        for expr in exprs {
                            expr.get_cached_hash().hash(&mut hasher);
                        }
                    }
                }
                AbstractLiteral::Partition(v) => {
                    for exprs in v {
                        for expr in exprs {
                            expr.get_cached_hash().hash(&mut hasher);
                        }
                    }
                }
            },
            Expression::Root(_, vs) => {
                for expr in vs {
                    expr.get_cached_hash().hash(&mut hasher);
                }
            }

            // Moo<Expression>
            Expression::DominanceRelation(_, m1)
            | Expression::ToInt(_, m1)
            | Expression::Abs(_, m1)
            | Expression::Sum(_, m1)
            | Expression::Product(_, m1)
            | Expression::Min(_, m1)
            | Expression::Max(_, m1)
            | Expression::Not(_, m1)
            | Expression::Or(_, m1)
            | Expression::And(_, m1)
            | Expression::Neg(_, m1)
            | Expression::Defined(_, m1)
            | Expression::AllDiff(_, m1)
            | Expression::Factorial(_, m1)
            | Expression::Participants(_, m1)
            | Expression::Parts(_, m1)
            | Expression::Range(_, m1)
            | Expression::ToSet(_, m1)
            | Expression::ToMSet(_, m1)
            | Expression::ToRelation(_, m1)
            | Expression::Card(_, m1) => {
                m1.get_cached_hash().hash(&mut hasher);
            }

            // Moo<Expression> + Moo<Expression>
            Expression::Table(_, m1, m2)
            | Expression::NegativeTable(_, m1, m2)
            | Expression::Bubble(_, m1, m2)
            | Expression::Imply(_, m1, m2)
            | Expression::Iff(_, m1, m2)
            | Expression::Union(_, m1, m2)
            | Expression::In(_, m1, m2)
            | Expression::Intersect(_, m1, m2)
            | Expression::Supset(_, m1, m2)
            | Expression::SupsetEq(_, m1, m2)
            | Expression::Subset(_, m1, m2)
            | Expression::SubsetEq(_, m1, m2)
            | Expression::Eq(_, m1, m2)
            | Expression::Neq(_, m1, m2)
            | Expression::Geq(_, m1, m2)
            | Expression::Leq(_, m1, m2)
            | Expression::Gt(_, m1, m2)
            | Expression::Lt(_, m1, m2)
            | Expression::Apart(_, m1, m2)
            | Expression::Together(_, m1, m2)
            | Expression::Party(_, m1, m2)
            | Expression::SafeDiv(_, m1, m2)
            | Expression::UnsafeDiv(_, m1, m2)
            | Expression::SafeMod(_, m1, m2)
            | Expression::UnsafeMod(_, m1, m2)
            | Expression::UnsafePow(_, m1, m2)
            | Expression::SafePow(_, m1, m2)
            | Expression::Minus(_, m1, m2)
            | Expression::PairwiseSum(_, m1, m2)
            | Expression::PairwiseProduct(_, m1, m2)
            | Expression::Image(_, m1, m2)
            | Expression::ImageSet(_, m1, m2)
            | Expression::PreImage(_, m1, m2)
            | Expression::Inverse(_, m1, m2)
            | Expression::Restrict(_, m1, m2)
            | Expression::LexLt(_, m1, m2)
            | Expression::LexLeq(_, m1, m2)
            | Expression::LexGt(_, m1, m2)
            | Expression::LexGeq(_, m1, m2)
            | Expression::Subsequence(_, m1, m2)
            | Expression::Substring(_, m1, m2) => {
                m1.get_cached_hash().hash(&mut hasher);
                m2.get_cached_hash().hash(&mut hasher);
            }
            // Moo<Expression> + Vec<Expression>
            Expression::UnsafeIndex(_, m, vs) | Expression::SafeIndex(_, m, vs) => {
                m.get_cached_hash().hash(&mut hasher);
                for v in vs {
                    v.get_cached_hash().hash(&mut hasher);
                }
            }

            // Moo<Expression> + Vec<Option<Expression>>
            Expression::UnsafeSlice(_, m, vs)
            | Expression::SafeSlice(_, m, vs)
            | Expression::RelationProj(_, m, vs) => {
                m.get_cached_hash().hash(&mut hasher);
                for v in vs {
                    match v {
                        Some(e) => e.get_cached_hash().hash(&mut hasher),
                        None => 0u64.hash(&mut hasher),
                    }
                }
            }

            // Moo<Expression> + Name
            Expression::RecordField(_, m, n) | Expression::Active(_, m, n) => {
                m.get_cached_hash().hash(&mut hasher);
                n.hash(&mut hasher);
            }

            // Moo<Expression> + DomainPtr
            Expression::InDomain(_, m, d) => {
                m.get_cached_hash().hash(&mut hasher);
                d.hash(&mut hasher);
            }

            // Option<Moo<Expression>> + Moo<Expression>
            Expression::Flatten(_, opt, m) => {
                if let Some(e) = opt {
                    e.get_cached_hash().hash(&mut hasher);
                }
                m.get_cached_hash().hash(&mut hasher);
            }

            // Moo<Expression> + Atom
            Expression::MinionReify(_, m, a) | Expression::MinionReifyImply(_, m, a) => {
                m.get_cached_hash().hash(&mut hasher);
                a.hash(&mut hasher);
            }

            // Reference + Moo<Expression>
            Expression::AuxDeclaration(_, r, m) => {
                r.hash(&mut hasher);
                m.get_cached_hash().hash(&mut hasher);
            }

            // SATIntEncoding + Moo<Expression> + (i32, i32)
            Expression::SATInt(_, enc, m, bounds) => {
                enc.hash(&mut hasher);
                m.get_cached_hash().hash(&mut hasher);
                bounds.hash(&mut hasher);
            }

            // Non-Expression Moo types - hash normally
            Expression::Comprehension(_, c) => c.hash(&mut hasher),
            Expression::AbstractComprehension(_, c) => c.hash(&mut hasher),

            // Leaf types - no Expression children
            Expression::Atomic(_, a) => a.hash(&mut hasher),
            Expression::FromSolution(_, a) => a.hash(&mut hasher),
            Expression::Metavar(_, u) => u.hash(&mut hasher),

            // Two Moo<Atom>
            Expression::FlatAbsEq(_, a1, a2) | Expression::FlatMinusEq(_, a1, a2) => {
                a1.hash(&mut hasher);
                a2.hash(&mut hasher);
            }

            // Three Moo<Atom>
            Expression::FlatProductEq(_, a1, a2, a3)
            | Expression::MinionDivEqUndefZero(_, a1, a2, a3)
            | Expression::MinionModuloEqUndefZero(_, a1, a2, a3)
            | Expression::MinionPow(_, a1, a2, a3) => {
                a1.hash(&mut hasher);
                a2.hash(&mut hasher);
                a3.hash(&mut hasher);
            }

            // Vec<Atom>
            Expression::FlatAllDiff(_, vs) => {
                for v in vs {
                    v.hash(&mut hasher);
                }
            }

            // Vec<Atom> + Atom
            Expression::FlatSumGeq(_, vs, a) | Expression::FlatSumLeq(_, vs, a) => {
                for v in vs {
                    v.hash(&mut hasher);
                }
                a.hash(&mut hasher);
            }

            // Moo<Atom> + Moo<Atom> + Box<Literal>
            Expression::FlatIneq(_, a1, a2, lit) => {
                a1.hash(&mut hasher);
                a2.hash(&mut hasher);
                lit.hash(&mut hasher);
            }

            // Reference + Literal
            Expression::FlatWatchedLiteral(_, r, l) => {
                r.hash(&mut hasher);
                l.hash(&mut hasher);
            }

            // Vec<Literal> + Vec<Atom> + Moo<Atom>
            Expression::FlatWeightedSumLeq(_, lits, atoms, a)
            | Expression::FlatWeightedSumGeq(_, lits, atoms, a) => {
                for l in lits {
                    l.hash(&mut hasher);
                }
                for at in atoms {
                    at.hash(&mut hasher);
                }
                a.hash(&mut hasher);
            }

            // Atom + Vec<i32>
            Expression::MinionWInIntervalSet(_, a, vs) | Expression::MinionWInSet(_, a, vs) => {
                a.hash(&mut hasher);
                for v in vs {
                    v.hash(&mut hasher);
                }
            }

            // Vec<Atom> + Moo<Atom> + Moo<Atom>
            Expression::MinionElementOne(_, vs, a1, a2) => {
                for v in vs {
                    v.hash(&mut hasher);
                }
                a1.hash(&mut hasher);
                a2.hash(&mut hasher);
            }

            // Vec<Atom> + Vec<Atom>
            Expression::FlatLexLt(_, v1, v2) | Expression::FlatLexLeq(_, v1, v2) => {
                for v in v1 {
                    v.hash(&mut hasher);
                }
                for v in v2 {
                    v.hash(&mut hasher);
                }
            }
        };

        let result = hasher.finish();
        self.meta_ref().stored_hash.swap(result, Ordering::Relaxed);
        result
    }
}

#[cfg(test)]
mod tests {
    use crate::matrix_expr;

    use super::*;

    #[test]
    fn test_domain_of_constant_sum() {
        let c1 = Expression::Atomic(Metadata::new(), Atom::Literal(Literal::Int(1)));
        let c2 = Expression::Atomic(Metadata::new(), Atom::Literal(Literal::Int(2)));
        let sum = Expression::Sum(Metadata::new(), Moo::new(matrix_expr![c1, c2]));
        assert_eq!(sum.domain_of(), Some(Domain::int(vec![Range::Single(3)])));
    }

    #[test]
    fn test_domain_of_constant_invalid_type() {
        let c1 = Expression::Atomic(Metadata::new(), Atom::Literal(Literal::Int(1)));
        let c2 = Expression::Atomic(Metadata::new(), Atom::Literal(Literal::Bool(true)));
        let sum = Expression::Sum(Metadata::new(), Moo::new(matrix_expr![c1, c2]));
        assert_eq!(sum.domain_of(), None);
    }

    #[test]
    fn test_domain_of_empty_sum() {
        let sum = Expression::Sum(Metadata::new(), Moo::new(matrix_expr![]));
        assert_eq!(sum.domain_of(), None);
    }
}