conjure_cp_core/ast/

literals.rs

use funcmap::FuncMap;
use itertools::Itertools;
use serde::{Deserialize, Serialize};
use std::fmt::{Debug, Display, Formatter};
use std::hash::Hash;
use ustr::Ustr;

use super::{
    Atom, Domain, DomainPtr, Expression, GroundDomain, Metadata, Moo, PartitionAttr, Range,
    ReturnType, SetAttr, Typeable, domains::HasDomain, domains::Int, records::Field,
};
use crate::ast::domains::{MSetAttr, SequenceAttr};
use crate::ast::pretty::pretty_vec;
use crate::bug;
use polyquine::Quine;
use uniplate::{Biplate, Tree, Uniplate};

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Uniplate, Hash, Quine)]
#[uniplate(walk_into=[AbstractLiteral<Literal>])]
#[biplate(to=Atom)]
#[biplate(to=AbstractLiteral<Literal>)]
#[biplate(to=AbstractLiteral<Expression>)]
#[biplate(to=Field<Literal>)]
#[biplate(to=Field<Expression>)]
#[biplate(to=Expression)]
#[path_prefix(conjure_cp::ast)]
/// A literal value, equivalent to constants in Conjure.
pub enum Literal {
    Int(i32),
    Bool(bool),
    //abstract literal variant ends in Literal, but that's ok
    #[allow(clippy::enum_variant_names)]
    AbstractLiteral(AbstractLiteral<Literal>),
}

impl HasDomain for Literal {
    fn domain_of(&self) -> DomainPtr {
        match self {
            Literal::Int(i) => Domain::int(vec![Range::Single(*i)]),
            Literal::Bool(_) => Domain::bool(),
            Literal::AbstractLiteral(abstract_literal) => abstract_literal.domain_of(),
        }
    }
}

// make possible values of an AbstractLiteral a closed world to make the trait bounds more sane (particularly in Uniplate instances!!)
pub trait AbstractLiteralValue:
    Clone + Eq + PartialEq + Display + Uniplate + Biplate<Field<Self>> + 'static
{
    type Dom: Clone
        + Eq
        + PartialEq
        + Debug
        + Display
        + Quine
        + From<GroundDomain>
        + Into<DomainPtr>;
}
impl AbstractLiteralValue for Expression {
    type Dom = DomainPtr;
}
impl AbstractLiteralValue for Literal {
    type Dom = Moo<GroundDomain>;
}

#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Quine)]
#[path_prefix(conjure_cp::ast)]
pub enum AbstractLiteral<T: AbstractLiteralValue> {
    Set(Vec<T>),

    MSet(Vec<T>),

    /// A 1 dimensional matrix slice with an index domain.
    Matrix(Vec<T>, T::Dom),

    // a tuple of literals
    Tuple(Vec<T>),

    Record(Vec<Field<T>>),

    Sequence(Vec<T>),

    Function(Vec<(T, T)>),

    // Variants only contain one of their name-domain pairs
    Variant(Moo<Field<T>>),

    // A list of partitions, each part has a set of values
    Partition(Vec<Vec<T>>),
    Relation(Vec<Vec<T>>),
}

// TODO: use HasDomain instead once Expression::domain_of returns Domain not Option<Domain>
impl AbstractLiteral<Expression> {
    pub fn domain_of(&self) -> Option<DomainPtr> {
        match self {
            AbstractLiteral::Set(items) => {
                // ensure that all items have a domain, or return None
                let item_domains: Vec<DomainPtr> = items
                    .iter()
                    .map(|x| x.domain_of())
                    .collect::<Option<Vec<DomainPtr>>>()?;

                // union all item domains together
                let mut item_domain_iter = item_domains.iter().cloned();
                let first_item = item_domain_iter.next()?;
                let item_domain = item_domains
                    .iter()
                    .try_fold(first_item, |x, y| x.union(y))
                    .expect("taking the union of all item domains of a set literal should succeed");

                Some(Domain::set(SetAttr::<Int>::default(), item_domain))
            }

            AbstractLiteral::MSet(items) => {
                // ensure that all items have a domain, or return None
                let item_domains: Vec<DomainPtr> = items
                    .iter()
                    .map(|x| x.domain_of())
                    .collect::<Option<Vec<DomainPtr>>>()?;

                // union all item domains together
                let mut item_domain_iter = item_domains.iter().cloned();
                let first_item = item_domain_iter.next()?;
                let item_domain = item_domains
                    .iter()
                    .try_fold(first_item, |x, y| x.union(y))
                    .expect("taking the union of all item domains of a set literal should succeed");

                Some(Domain::mset(MSetAttr::<Int>::default(), item_domain))
            }

            AbstractLiteral::Sequence(elems) => {
                let item_domains: Vec<DomainPtr> = elems
                    .iter()
                    .map(|x| x.domain_of())
                    .collect::<Option<Vec<DomainPtr>>>()?;

                // Get the union of all domains in the sequence.
                // i.e. if <(1..3), (1..3), (5), (8..9)> then seq dom is (1..3, 5, 8..9)
                let mut item_domain_iter = item_domains.iter().cloned();
                let first_item = item_domain_iter.next()?;
                let item_domain = item_domains
                    .iter()
                    .try_fold(first_item, |x, y| x.union(y))
                    .expect("taking the union of all item domains of a set literal should succeed");

                Some(Domain::sequence(
                    SequenceAttr::<Int>::default(),
                    item_domain,
                ))
            }

            AbstractLiteral::Partition(items) => {
                // Flatten the Vec<Vec< into a single vec
                // ensure that all elemes in each part have a domain, or return None

                let item_domains: Vec<DomainPtr> = items
                    .iter()
                    .flatten()
                    .map(|x| x.domain_of())
                    .collect::<Option<Vec<DomainPtr>>>()?;

                // union all item domains together
                let mut item_domain_iter = item_domains.iter().cloned();
                let first_item = item_domain_iter.next()?;
                let item_domain = item_domains
                    .iter()
                    .try_fold(first_item, |x, y| x.union(y))
                    .expect("taking the union of all item domains of a partition literal should succeed");

                Some(Domain::partition(
                    PartitionAttr::<Int>::default(),
                    item_domain,
                ))
            }

            AbstractLiteral::Matrix(items, _) => {
                // ensure that all items have a domain, or return None
                let item_domains = items
                    .iter()
                    .map(|x| x.domain_of())
                    .collect::<Option<Vec<DomainPtr>>>()?;

                // union all item domains together
                let mut item_domain_iter = item_domains.iter().cloned();

                let first_item = item_domain_iter.next()?;

                let item_domain = item_domains
                    .iter()
                    .try_fold(first_item, |x, y| x.union(y))
                    .expect(
                        "taking the union of all item domains of a matrix literal should succeed",
                    );

                let mut new_index_domain = vec![];

                // flatten index domains of n-d matrix into list
                let mut e = Expression::AbstractLiteral(Metadata::new(), self.clone());
                while let Expression::AbstractLiteral(_, AbstractLiteral::Matrix(elems, idx)) = e {
                    assert!(
                        idx.as_matrix().is_none(),
                        "n-dimensional matrix literals should be represented as a matrix inside a matrix, got {idx}"
                    );
                    new_index_domain.push(idx);
                    e = elems[0].clone();
                }
                Some(Domain::matrix(item_domain, new_index_domain))
            }
            AbstractLiteral::Tuple(_) => None,
            AbstractLiteral::Record(_) => None,
            AbstractLiteral::Function(_) => None,
            AbstractLiteral::Variant(_) => None,
            AbstractLiteral::Relation(_) => None,
        }
    }
}

impl HasDomain for AbstractLiteral<Literal> {
    fn domain_of(&self) -> DomainPtr {
        Domain::from_literal_vec(&[Literal::AbstractLiteral(self.clone())])
            .expect("abstract literals should be correctly typed")
    }
}

impl Typeable for AbstractLiteral<Expression> {
    fn return_type(&self) -> ReturnType {
        match self {
            AbstractLiteral::Set(items) if items.is_empty() => {
                ReturnType::Set(Box::new(ReturnType::Unknown))
            }
            AbstractLiteral::Set(items) => {
                let item_type = items[0].return_type();

                // if any items do not have a type, return none.
                let item_types: Vec<ReturnType> = items.iter().map(|x| x.return_type()).collect();

                assert!(
                    item_types.iter().all(|x| x == &item_type),
                    "all items in a set should have the same type"
                );

                ReturnType::Set(Box::new(item_type))
            }
            AbstractLiteral::MSet(items) if items.is_empty() => {
                ReturnType::MSet(Box::new(ReturnType::Unknown))
            }
            AbstractLiteral::MSet(items) => {
                let item_type = items[0].return_type();

                // if any items do not have a type, return none.
                let item_types: Vec<ReturnType> = items.iter().map(|x| x.return_type()).collect();

                assert!(
                    item_types.iter().all(|x| x == &item_type),
                    "all items in a set should have the same type"
                );

                ReturnType::MSet(Box::new(item_type))
            }
            AbstractLiteral::Sequence(items) if items.is_empty() => {
                ReturnType::Sequence(Box::new(ReturnType::Unknown))
            }
            AbstractLiteral::Sequence(items) => {
                let item_type = items[0].return_type();

                // if any items do not have a type, return none.
                let item_types: Vec<ReturnType> = items.iter().map(|x| x.return_type()).collect();

                assert!(
                    item_types.iter().all(|x| x == &item_type),
                    "all items in a sequence should have the same type"
                );

                ReturnType::Sequence(Box::new(item_type))
            }
            AbstractLiteral::Partition(items) if items.is_empty() || items[0].is_empty() => {
                ReturnType::Partition(Box::new(ReturnType::Unknown))
            }
            AbstractLiteral::Partition(items) => {
                let item_type = items[0][0].return_type();

                // if any items do not have a type, return none.
                let item_types: Vec<ReturnType> =
                    items.iter().flatten().map(|x| x.return_type()).collect();

                assert!(
                    item_types.iter().all(|x| x == &item_type),
                    "all items in every part of a partition should have the same type"
                );

                ReturnType::Partition(Box::new(item_type))
            }
            AbstractLiteral::Matrix(items, _) if items.is_empty() => {
                ReturnType::Matrix(Box::new(ReturnType::Unknown))
            }
            AbstractLiteral::Matrix(items, _) => {
                let item_type = items[0].return_type();

                // if any items do not have a type, return none.
                let item_types: Vec<ReturnType> = items.iter().map(|x| x.return_type()).collect();

                assert!(
                    item_types.iter().all(|x| x == &item_type),
                    "all items in a matrix should have the same type. items: {items} types: {types:#?}",
                    items = pretty_vec(items),
                    types = items
                        .iter()
                        .map(|x| x.return_type())
                        .collect::<Vec<ReturnType>>()
                );

                ReturnType::Matrix(Box::new(item_type))
            }
            AbstractLiteral::Tuple(items) => {
                let mut item_types = vec![];
                for item in items {
                    item_types.push(item.return_type());
                }
                ReturnType::Tuple(item_types)
            }
            AbstractLiteral::Record(items) => {
                let mut item_types = vec![];
                for item in items {
                    item_types.push(item.clone().func_map(|x| x.return_type()));
                }
                ReturnType::Record(item_types)
            }
            AbstractLiteral::Function(items) => {
                if items.is_empty() {
                    return ReturnType::Function(
                        Box::new(ReturnType::Unknown),
                        Box::new(ReturnType::Unknown),
                    );
                }

                // Check that all items have the same return type
                let (x1, y1) = &items[0];
                let (t1, t2) = (x1.return_type(), y1.return_type());
                for (x, y) in items {
                    let (tx, ty) = (x.return_type(), y.return_type());
                    if tx != t1 {
                        bug!("Expected {t1}, got {x}: {tx}");
                    }
                    if ty != t2 {
                        bug!("Expected {t2}, got {y}: {ty}");
                    }
                }

                ReturnType::Function(Box::new(t1), Box::new(t2))
            }
            AbstractLiteral::Variant(item) => {
                // Variants hold multiple possible types. In the case of a literal we know which type it chose
                ReturnType::Variant(vec![item.as_ref().clone().func_map(|x| x.return_type())])
            }
            AbstractLiteral::Relation(items) => {
                if items.is_empty() {
                    return ReturnType::Relation(vec![ReturnType::Unknown]);
                }
                let mut item_types = vec![];
                let x1 = &items[0];
                let size = x1.len();
                for item in x1 {
                    item_types.push(item.return_type());
                }
                for x in items {
                    if x.len() != size {
                        let strs = item_types.iter().map(|x| format!("{}", x)).join(",");
                        bug!("Expected ({strs}) with length {size}, got size {}", x.len());
                    }
                    for i in 1..size {
                        if let Some(new_type) = x.get(i)
                            && let Some(old_type) = item_types.get(i)
                            && new_type.return_type() != *old_type
                        {
                            bug!("Expected {old_type}, got {new_type}");
                        }
                    }
                }
                ReturnType::Relation(item_types)
            }
        }
    }
}

impl<T> AbstractLiteral<T>
where
    T: AbstractLiteralValue,
{
    /// Creates a matrix with elements `elems`, with domain `int(1..)`.
    ///
    /// This acts as a variable sized list.
    pub fn matrix_implied_indices(elems: Vec<T>) -> Self {
        AbstractLiteral::Matrix(elems, GroundDomain::Int(vec![Range::UnboundedR(1)]).into())
    }

    /// If the AbstractLiteral is a list, returns its elements.
    ///
    /// 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<T>> {
        let AbstractLiteral::Matrix(elems, domain) = self else {
            return None;
        };

        let domain: DomainPtr = domain.clone().into();
        let Some(GroundDomain::Int(ranges)) = domain.as_ground() else {
            return None;
        };

        let [Range::UnboundedR(1)] = ranges[..] else {
            return None;
        };

        Some(elems)
    }
}

impl<T> Display for AbstractLiteral<T>
where
    T: AbstractLiteralValue,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            AbstractLiteral::Set(elems) => {
                let elems_str: String = elems.iter().map(|x| format!("{x}")).join(",");
                write!(f, "{{{elems_str}}}")
            }
            AbstractLiteral::MSet(elems) => {
                let elems_str: String = elems.iter().map(|x| format!("{x}")).join(",");
                write!(f, "mset({elems_str})")
            }
            AbstractLiteral::Matrix(elems, index_domain) => {
                let elems_str: String = elems.iter().map(|x| format!("{x}")).join(",");
                write!(f, "[{elems_str};{index_domain}]")
            }
            AbstractLiteral::Tuple(elems) => {
                let elems_str: String = elems.iter().map(|x| format!("{x}")).join(",");
                write!(f, "({elems_str})")
            }
            AbstractLiteral::Sequence(elems) => {
                let elems_str: String = elems.iter().map(|x| format!("{x}")).join(",");
                write!(f, "sequence({elems_str})")
            }
            AbstractLiteral::Partition(parts) => {
                let elems_str: String = parts
                    .iter()
                    .map(|inner| {
                        let elems_str = inner.iter().map(|x| format!("{x}")).join(",");
                        format!("{{{}}}", elems_str)
                    })
                    .join(", ");

                write!(f, "partition({elems_str})")
            }
            AbstractLiteral::Record(entries) => {
                let entries_str: String = entries
                    .iter()
                    .map(|entry| format!("{} = {}", entry.name, entry.value))
                    .join(",");
                write!(f, "record {{{entries_str}}}")
            }
            AbstractLiteral::Function(entries) => {
                let entries_str: String = entries
                    .iter()
                    .map(|entry| format!("{} --> {}", entry.0, entry.1))
                    .join(",");
                write!(f, "function({entries_str})")
            }
            AbstractLiteral::Variant(entry) => {
                write!(f, "variant{{{} = {}}}", entry.name, entry.value)
            }
            AbstractLiteral::Relation(elems) => {
                let elems_str: String = elems
                    .iter()
                    .map(|x| format!("({})", x.iter().map(|x| format!("{x}")).join(",")))
                    .join(",");
                write!(f, "relation({elems_str})")
            }
        }
    }
}

impl<T> Uniplate for AbstractLiteral<T>
where
    T: AbstractLiteralValue + Biplate<AbstractLiteral<T>>,
{
    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
        // walking into T
        match self {
            AbstractLiteral::Set(vec) => {
                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(vec);
                (f1_tree, Box::new(move |x| AbstractLiteral::Set(f1_ctx(x))))
            }
            AbstractLiteral::MSet(vec) => {
                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(vec);
                (f1_tree, Box::new(move |x| AbstractLiteral::MSet(f1_ctx(x))))
            }
            AbstractLiteral::Matrix(elems, index_domain) => {
                let index_domain = index_domain.clone();
                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(elems);
                (
                    f1_tree,
                    Box::new(move |x| AbstractLiteral::Matrix(f1_ctx(x), index_domain.clone())),
                )
            }
            AbstractLiteral::Sequence(vec) => {
                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(vec);
                (
                    f1_tree,
                    Box::new(move |x| AbstractLiteral::Sequence(f1_ctx(x))),
                )
            }
            AbstractLiteral::Tuple(elems) => {
                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(elems);
                (
                    f1_tree,
                    Box::new(move |x| AbstractLiteral::Tuple(f1_ctx(x))),
                )
            }
            AbstractLiteral::Record(entries) => {
                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(entries);
                (
                    f1_tree,
                    Box::new(move |x| AbstractLiteral::Record(f1_ctx(x))),
                )
            }
            AbstractLiteral::Function(entries) => {
                let entry_count = entries.len();
                let flattened: Vec<T> = entries
                    .iter()
                    .flat_map(|(lhs, rhs)| [lhs.clone(), rhs.clone()])
                    .collect();

                let (f1_tree, f1_ctx) =
                    <Vec<T> as Biplate<AbstractLiteral<T>>>::biplate(&flattened);
                (
                    f1_tree,
                    Box::new(move |x| {
                        let rebuilt = f1_ctx(x);
                        assert_eq!(
                            rebuilt.len(),
                            entry_count * 2,
                            "number of function literal children should remain unchanged"
                        );

                        let mut iter = rebuilt.into_iter();
                        let mut pairs = Vec::with_capacity(entry_count);
                        while let (Some(lhs), Some(rhs)) = (iter.next(), iter.next()) {
                            pairs.push((lhs, rhs));
                        }

                        AbstractLiteral::Function(pairs)
                    }),
                )
            }
            AbstractLiteral::Variant(entries) => {
                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(entries);
                (
                    f1_tree,
                    Box::new(move |x| AbstractLiteral::Variant(f1_ctx(x))),
                )
            }
            AbstractLiteral::Relation(elems) => {
                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(elems);
                (
                    f1_tree,
                    Box::new(move |x| AbstractLiteral::Relation(f1_ctx(x))),
                )
            }
            AbstractLiteral::Partition(elems) => {
                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(elems);
                (
                    f1_tree,
                    Box::new(move |x| AbstractLiteral::Partition(f1_ctx(x))),
                )
            }
        }
    }
}

impl<U, To> Biplate<To> for AbstractLiteral<U>
where
    To: Uniplate,
    U: AbstractLiteralValue + Biplate<AbstractLiteral<U>> + Biplate<To>,
    Field<U>: Biplate<AbstractLiteral<U>> + Biplate<To>,
{
    fn biplate(&self) -> (Tree<To>, Box<dyn Fn(Tree<To>) -> Self>) {
        if std::any::TypeId::of::<To>() == std::any::TypeId::of::<AbstractLiteral<U>>() {
            // To ==From => return One(self)

            unsafe {
                // SAFETY: asserted the type equality above
                let self_to = std::mem::transmute::<&AbstractLiteral<U>, &To>(self).clone();
                let tree = Tree::One(self_to);
                let ctx = Box::new(move |x| {
                    let Tree::One(x) = x else {
                        panic!();
                    };

                    std::mem::transmute::<&To, &AbstractLiteral<U>>(&x).clone()
                });

                (tree, ctx)
            }
        } else {
            // walking into T
            match self {
                AbstractLiteral::Set(vec) => {
                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(vec);
                    (f1_tree, Box::new(move |x| AbstractLiteral::Set(f1_ctx(x))))
                }
                AbstractLiteral::MSet(vec) => {
                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(vec);
                    (f1_tree, Box::new(move |x| AbstractLiteral::MSet(f1_ctx(x))))
                }
                AbstractLiteral::Matrix(elems, index_domain) => {
                    let index_domain = index_domain.clone();
                    let (f1_tree, f1_ctx) = <Vec<U> as Biplate<To>>::biplate(elems);
                    (
                        f1_tree,
                        Box::new(move |x| AbstractLiteral::Matrix(f1_ctx(x), index_domain.clone())),
                    )
                }
                AbstractLiteral::Sequence(vec) => {
                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(vec);
                    (
                        f1_tree,
                        Box::new(move |x| AbstractLiteral::Sequence(f1_ctx(x))),
                    )
                }
                AbstractLiteral::Tuple(elems) => {
                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(elems);
                    (
                        f1_tree,
                        Box::new(move |x| AbstractLiteral::Tuple(f1_ctx(x))),
                    )
                }
                AbstractLiteral::Record(entries) => {
                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(entries);
                    (
                        f1_tree,
                        Box::new(move |x| AbstractLiteral::Record(f1_ctx(x))),
                    )
                }
                AbstractLiteral::Function(entries) => {
                    let entry_count = entries.len();
                    let flattened: Vec<U> = entries
                        .iter()
                        .flat_map(|(lhs, rhs)| [lhs.clone(), rhs.clone()])
                        .collect();

                    let (f1_tree, f1_ctx) = <Vec<U> as Biplate<To>>::biplate(&flattened);
                    (
                        f1_tree,
                        Box::new(move |x| {
                            let rebuilt = f1_ctx(x);
                            assert_eq!(
                                rebuilt.len(),
                                entry_count * 2,
                                "number of function literal children should remain unchanged"
                            );

                            let mut iter = rebuilt.into_iter();
                            let mut pairs = Vec::with_capacity(entry_count);
                            while let (Some(lhs), Some(rhs)) = (iter.next(), iter.next()) {
                                pairs.push((lhs, rhs));
                            }

                            AbstractLiteral::Function(pairs)
                        }),
                    )
                }
                AbstractLiteral::Variant(entries) => {
                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(entries);
                    (
                        f1_tree,
                        Box::new(move |x| AbstractLiteral::Variant(f1_ctx(x))),
                    )
                }
                AbstractLiteral::Relation(elems) => {
                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(elems);
                    (
                        f1_tree,
                        Box::new(move |x| AbstractLiteral::Relation(f1_ctx(x))),
                    )
                }
                AbstractLiteral::Partition(elems) => {
                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(elems);
                    (
                        f1_tree,
                        Box::new(move |x| AbstractLiteral::Partition(f1_ctx(x))),
                    )
                }
            }
        }
    }
}

impl TryFrom<Literal> for i32 {
    type Error = &'static str;

    fn try_from(value: Literal) -> Result<Self, Self::Error> {
        match value {
            Literal::Int(i) => Ok(i),
            _ => Err("Cannot convert non-i32 literal to i32"),
        }
    }
}

impl TryFrom<Box<Literal>> for i32 {
    type Error = &'static str;

    fn try_from(value: Box<Literal>) -> Result<Self, Self::Error> {
        (*value).try_into()
    }
}

impl TryFrom<&Box<Literal>> for i32 {
    type Error = &'static str;

    fn try_from(value: &Box<Literal>) -> Result<Self, Self::Error> {
        TryFrom::<&Literal>::try_from(value.as_ref())
    }
}

impl TryFrom<&Moo<Literal>> for i32 {
    type Error = &'static str;

    fn try_from(value: &Moo<Literal>) -> Result<Self, Self::Error> {
        TryFrom::<&Literal>::try_from(value.as_ref())
    }
}

impl TryFrom<&Literal> for i32 {
    type Error = &'static str;

    fn try_from(value: &Literal) -> Result<Self, Self::Error> {
        match value {
            Literal::Int(i) => Ok(*i),
            _ => Err("Cannot convert non-i32 literal to i32"),
        }
    }
}

impl TryFrom<Literal> for bool {
    type Error = &'static str;

    fn try_from(value: Literal) -> Result<Self, Self::Error> {
        match value {
            Literal::Bool(b) => Ok(b),
            _ => Err("Cannot convert non-bool literal to bool"),
        }
    }
}

impl TryFrom<&Literal> for bool {
    type Error = &'static str;

    fn try_from(value: &Literal) -> Result<Self, Self::Error> {
        match value {
            Literal::Bool(b) => Ok(*b),
            _ => Err("Cannot convert non-bool literal to bool"),
        }
    }
}

impl From<i32> for Literal {
    fn from(i: i32) -> Self {
        Literal::Int(i)
    }
}

impl From<bool> for Literal {
    fn from(b: bool) -> Self {
        Literal::Bool(b)
    }
}

impl From<Literal> for Ustr {
    fn from(value: Literal) -> Self {
        // TODO: avoid the temporary-allocation of a string by format! here?
        Ustr::from(&format!("{value}"))
    }
}

impl From<AbstractLiteral<Literal>> for Literal {
    fn from(literal: AbstractLiteral<Literal>) -> Self {
        Literal::AbstractLiteral(literal)
    }
}

impl AbstractLiteral<Expression> {
    /// If all the elements are literals, returns this as an AbstractLiteral<Literal>.
    /// Otherwise, returns `None`.
    pub fn into_literals(self) -> Option<AbstractLiteral<Literal>> {
        match self {
            AbstractLiteral::Set(elements) => {
                let literals = elements
                    .into_iter()
                    .map(|expr| match expr {
                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
                        Expression::AbstractLiteral(_, abslit) => {
                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
                        }
                        _ => None,
                    })
                    .collect::<Option<Vec<_>>>()?;
                Some(AbstractLiteral::Set(literals))
            }
            AbstractLiteral::MSet(elements) => {
                let literals = elements
                    .into_iter()
                    .map(|expr| match expr {
                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
                        Expression::AbstractLiteral(_, abslit) => {
                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
                        }
                        _ => None,
                    })
                    .collect::<Option<Vec<_>>>()?;
                Some(AbstractLiteral::MSet(literals))
            }
            AbstractLiteral::Partition(elems) => {
                // want to ascertain if every elem in Vec<Vec<Expr>> is a literal. If any are not, return none
                // otherwise confirm it is an abslit<lit>
                let mut partition: Vec<Vec<_>> = Vec::new();

                for part in elems {
                    let literals = part
                        .into_iter()
                        .map(|expr| match expr {
                            Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
                            Expression::AbstractLiteral(_, abslit) => {
                                Some(Literal::AbstractLiteral(abslit.into_literals()?))
                            }
                            _ => None,
                        })
                        .collect::<Option<Vec<_>>>()?;

                    partition.push(literals);
                }

                Some(AbstractLiteral::Partition(partition))
            }
            AbstractLiteral::Matrix(items, domain) => {
                let mut literals = vec![];
                for item in items {
                    let literal = match item {
                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
                        Expression::AbstractLiteral(_, abslit) => {
                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
                        }
                        _ => None,
                    }?;
                    literals.push(literal);
                }

                Some(AbstractLiteral::Matrix(literals, domain.resolve().ok()?))
            }
            AbstractLiteral::Sequence(elements) => {
                let literals = elements
                    .into_iter()
                    .map(|expr| match expr {
                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
                        Expression::AbstractLiteral(_, abslit) => {
                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
                        }
                        _ => None,
                    })
                    .collect::<Option<Vec<_>>>()?;
                Some(AbstractLiteral::Sequence(literals))
            }
            AbstractLiteral::Tuple(items) => {
                let mut literals = vec![];
                for item in items {
                    let literal = match item {
                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
                        Expression::AbstractLiteral(_, abslit) => {
                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
                        }
                        _ => None,
                    }?;
                    literals.push(literal);
                }

                Some(AbstractLiteral::Tuple(literals))
            }
            AbstractLiteral::Record(entries) => {
                let mut literals = vec![];
                for entry in entries {
                    let literal = match entry.value {
                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
                        Expression::AbstractLiteral(_, abslit) => {
                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
                        }
                        _ => None,
                    }?;

                    literals.push((entry.name, literal));
                }
                Some(AbstractLiteral::Record(
                    literals
                        .into_iter()
                        .map(|(name, literal)| Field {
                            name,
                            value: literal,
                        })
                        .collect(),
                ))
            }
            AbstractLiteral::Function(_) => todo!("Implement into_literals for functions"),
            AbstractLiteral::Variant(entry) => {
                let literal = match entry.value.clone() {
                    Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
                    Expression::AbstractLiteral(_, abslit) => {
                        Some(Literal::AbstractLiteral(abslit.into_literals()?))
                    }
                    _ => None,
                }?;
                Some(AbstractLiteral::Variant(Moo::new(Field {
                    name: entry.name.clone(),
                    value: literal,
                })))
            }
            AbstractLiteral::Relation(_) => todo!("Implement into_literals for relations"),
        }
    }
}

// need display implementations for other types as well
impl Display for Literal {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match &self {
            Literal::Int(i) => write!(f, "{i}"),
            Literal::Bool(b) => write!(f, "{b}"),
            Literal::AbstractLiteral(l) => write!(f, "{l}"),
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use crate::ast::matrix::{flatten, partial_flatten, shape_of};
    use crate::{domain_int_ground, into_matrix, matrix, matrix_lit, range};
    use uniplate::Uniplate;

    #[test]
    fn matrix_uniplate_universe() {
        // Can we traverse through matrices with uniplate?
        let my_matrix: AbstractLiteral<Literal> = into_matrix![
            vec![Literal::AbstractLiteral(matrix![Literal::Bool(true);Moo::new(GroundDomain::Bool)]); 5];
            Moo::new(GroundDomain::Bool)
        ];

        let expected_index_domains = vec![Moo::new(GroundDomain::Bool); 6];
        let actual_index_domains: Vec<Moo<GroundDomain>> =
            my_matrix.cata(&move |elem, children| {
                let mut res = vec![];
                res.extend(children.into_iter().flatten());
                if let AbstractLiteral::Matrix(_, index_domain) = elem {
                    res.push(index_domain);
                }

                res
            });

        assert_eq!(actual_index_domains, expected_index_domains);
    }

    #[test]
    fn matrix_flatten() {
        let tensor: AbstractLiteral<Literal> = matrix![
            [
                // batch 1
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12]
            ],
            [
                // batch 2
                [13, 14, 15, 16],
                [17, 18, 19, 20],
                [21, 22, 23, 24]
            ]
        ];

        let actual_elems: Vec<Literal> = flatten(&tensor).cloned().collect();
        let expected_elems = (1..25).map(Literal::from).collect::<Vec<_>>();
        assert_eq!(actual_elems, expected_elems);
    }

    #[test]
    fn matrix_domain_1d() {
        let matrix = matrix_lit![10, 11, 12, 13; domain_int_ground!(1..4)];
        let dom = matrix.domain_of();

        let (inner_dom, idx_doms) = dom.as_matrix_ground().expect("must be ground matrix");
        assert_eq!(inner_dom, &domain_int_ground!(10..13));
        assert_eq!(idx_doms.len(), 1);
        assert_eq!(&idx_doms[0], &domain_int_ground!(1..4));
    }

    #[test]
    fn matrix_domain_2d() {
        let matrix = matrix_lit![
            [1, 2, 3, 4],
            [5, 6, 7, 8];
            [
                domain_int_ground!(1..2),
                domain_int_ground!(1..4)
            ]
        ];
        let dom = matrix.domain_of();

        let (inner_dom, idx_doms) = dom.as_matrix_ground().expect("must be ground matrix");
        assert_eq!(inner_dom, &domain_int_ground!(1..8));
        assert_eq!(idx_doms.len(), 2);
        assert_eq!(&idx_doms[0], &domain_int_ground!(1..2));
        assert_eq!(&idx_doms[1], &domain_int_ground!(1..4));
    }

    #[test]
    fn matrix_shape_3d() {
        let tensor: AbstractLiteral<Literal> = matrix![
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12]
            ],
            [
                [13, 14, 15, 16],
                [17, 18, 19, 20],
                [21, 22, 23, 24]
            ];
            [
                domain_int_ground!(1..2),
                domain_int_ground!(1..3),
                domain_int_ground!(1..4)
            ]
        ];
        let shape = shape_of(&tensor).expect("shape_of to work on a 3D matrix");

        assert_eq!(shape.size, 24);
        assert_eq!(shape.dims, vec![2, 3, 4]);
        assert_eq!(shape.strides, vec![12, 4, 1]);
        assert_eq!(
            shape.idx_doms,
            vec![
                domain_int_ground!(1..2),
                domain_int_ground!(1..3),
                domain_int_ground!(1..4)
            ]
        );
    }

    #[test]
    fn matrix_partial_flatten() {
        let tensor: AbstractLiteral<Literal> = matrix![
            [
                // batch 1
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12]
            ],
            [
                // batch 2
                [13, 14, 15, 16],
                [17, 18, 19, 20],
                [21, 22, 23, 24]
            ]
        ];
        assert_eq!(partial_flatten(0, tensor.clone()), tensor);

        let expected_flatten_1: AbstractLiteral<Literal> = matrix![
            [1, 2, 3, 4],
            [5, 6, 7, 8],
            [9, 10, 11, 12],
            [13, 14, 15, 16],
            [17, 18, 19, 20],
            [21, 22, 23, 24]
        ];
        assert_eq!(partial_flatten(1, tensor.clone()), expected_flatten_1);

        let expected_flatten_2 =
            AbstractLiteral::matrix_implied_indices((1..25).map(Literal::from).collect());
        assert_eq!(partial_flatten(2, tensor), expected_flatten_2);
    }
}