Skip to main content

conjure_cp_core/ast/
literals.rs

1use funcmap::FuncMap;
2use itertools::Itertools;
3use serde::{Deserialize, Serialize};
4use std::cmp::Ordering;
5use std::fmt::{Debug, Display, Formatter};
6use std::hash::Hash;
7use ustr::Ustr;
8
9use super::collection_domains::{SizeBounds, bounded_set};
10use super::eval::eval_constant;
11
12use super::{
13    Atom, Domain, DomainPtr, Expression, GroundDomain, Moo, PartitionAttr, PermutationAttr, Range,
14    ReturnType, SetAttr, Typeable, domains::HasDomain, domains::Int, records::Field,
15};
16use crate::ast::domains::{MSetAttr, SequenceAttr};
17use crate::ast::enumerated::EnumVariant;
18use crate::ast::pretty::pretty_vec;
19use crate::bug;
20use crate::bug_assert;
21use polyquine::Quine;
22use uniplate::{Biplate, Tree, Uniplate};
23
24#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Uniplate, Hash, Quine)]
25#[uniplate(walk_into=[AbstractLiteral<Literal>])]
26#[biplate(to=Atom)]
27#[biplate(to=AbstractLiteral<Literal>)]
28#[biplate(to=AbstractLiteral<Expression>)]
29#[biplate(to=Field<Literal>)]
30#[biplate(to=Field<Expression>)]
31#[biplate(to=Expression)]
32#[path_prefix(conjure_cp::ast)]
33/// A literal value, equivalent to constants in Conjure.
34pub enum Literal {
35    Int(i32),
36    Bool(bool),
37    EnumVariant(EnumVariant),
38    Unnamed(crate::ast::unnamed::UnnamedValue),
39    //abstract literal variant ends in Literal, but that's ok
40    #[allow(clippy::enum_variant_names)]
41    AbstractLiteral(AbstractLiteral<Literal>),
42}
43
44impl HasDomain for Literal {
45    fn domain_of(&self) -> DomainPtr {
46        match self {
47            Literal::Int(i) => Domain::int(vec![Range::Single(*i)]),
48            Literal::Bool(_) => Domain::bool(),
49            Literal::EnumVariant(ev) => {
50                Domain::enumerated_ground(ev.ty.clone(), vec![Range::Single(ev.variant)])
51            }
52            Literal::Unnamed(value) => {
53                Moo::new(crate::ast::GroundDomain::UnnamedType(value.ty.clone())).into()
54            }
55            Literal::AbstractLiteral(abstract_literal) => abstract_literal.domain_of(),
56        }
57    }
58}
59
60impl Literal {
61    /// Compare values using Essence-aware value ordering.
62    ///
63    /// Booleans and integers use their natural order, tuple-like values use
64    /// lexicographic order, and sets use lexicographic occurrence order over
65    /// ascending element values (`false < true`).
66    pub fn essence_cmp(&self, other: &Self) -> Ordering {
67        match (self, other) {
68            (Literal::Unnamed(lhs), Literal::Unnamed(rhs)) => {
69                (&lhs.ty.name, lhs.ty.size, lhs.value).cmp(&(&rhs.ty.name, rhs.ty.size, rhs.value))
70            }
71            (Literal::Bool(lhs), Literal::Bool(rhs)) => lhs.cmp(rhs),
72            (Literal::Int(lhs), Literal::Int(rhs)) => lhs.cmp(rhs),
73            (Literal::EnumVariant(lhs), Literal::EnumVariant(rhs)) => (
74                &lhs.ty.name,
75                &lhs.ty.variants,
76                lhs.variant,
77            )
78                .cmp(&(&rhs.ty.name, &rhs.ty.variants, rhs.variant)),
79            (Literal::AbstractLiteral(lhs), Literal::AbstractLiteral(rhs)) => {
80                abstract_literal_essence_cmp(lhs, rhs)
81            }
82            _ => literal_kind(self).cmp(&literal_kind(other)),
83        }
84    }
85}
86
87fn literal_kind(literal: &Literal) -> u8 {
88    match literal {
89        Literal::Bool(_) => 0,
90        Literal::Int(_) => 1,
91        Literal::AbstractLiteral(_) => 2,
92        Literal::EnumVariant(_) => 3,
93        Literal::Unnamed(_) => 4,
94    }
95}
96
97fn abstract_literal_essence_cmp(
98    lhs: &AbstractLiteral<Literal>,
99    rhs: &AbstractLiteral<Literal>,
100) -> Ordering {
101    match (lhs, rhs) {
102        (AbstractLiteral::Set(lhs), AbstractLiteral::Set(rhs)) => set_essence_cmp(lhs, rhs),
103        (AbstractLiteral::MSet(lhs), AbstractLiteral::MSet(rhs)) => sorted_literals_cmp(lhs, rhs),
104        (AbstractLiteral::Matrix(lhs, _), AbstractLiteral::Matrix(rhs, _))
105        | (AbstractLiteral::Tuple(lhs), AbstractLiteral::Tuple(rhs))
106        | (AbstractLiteral::Sequence(lhs), AbstractLiteral::Sequence(rhs)) => {
107            literal_slice_cmp(lhs, rhs)
108        }
109        (AbstractLiteral::Record(lhs), AbstractLiteral::Record(rhs)) => lhs
110            .iter()
111            .zip(rhs)
112            .find_map(|(lhs, rhs)| {
113                let ordering = lhs.name.to_string().cmp(&rhs.name.to_string());
114                (ordering != Ordering::Equal)
115                    .then_some(ordering)
116                    .or_else(|| {
117                        let ordering = lhs.value.essence_cmp(&rhs.value);
118                        (ordering != Ordering::Equal).then_some(ordering)
119                    })
120            })
121            .unwrap_or_else(|| lhs.len().cmp(&rhs.len())),
122        (AbstractLiteral::Function(lhs), AbstractLiteral::Function(rhs)) => lhs
123            .iter()
124            .zip(rhs)
125            .find_map(|((lhs_from, lhs_to), (rhs_from, rhs_to))| {
126                let ordering = lhs_from.essence_cmp(rhs_from);
127                (ordering != Ordering::Equal)
128                    .then_some(ordering)
129                    .or_else(|| {
130                        let ordering = lhs_to.essence_cmp(rhs_to);
131                        (ordering != Ordering::Equal).then_some(ordering)
132                    })
133            })
134            .unwrap_or_else(|| lhs.len().cmp(&rhs.len())),
135        (AbstractLiteral::Variant(lhs), AbstractLiteral::Variant(rhs)) => lhs
136            .name
137            .to_string()
138            .cmp(&rhs.name.to_string())
139            .then_with(|| lhs.value.essence_cmp(&rhs.value)),
140        (AbstractLiteral::Partition(lhs), AbstractLiteral::Partition(rhs))
141        | (AbstractLiteral::Relation(lhs), AbstractLiteral::Relation(rhs))
142        | (AbstractLiteral::Permutation(lhs), AbstractLiteral::Permutation(rhs)) => lhs
143            .iter()
144            .zip(rhs)
145            .find_map(|(lhs, rhs)| {
146                let ordering = literal_slice_cmp(lhs, rhs);
147                (ordering != Ordering::Equal).then_some(ordering)
148            })
149            .unwrap_or_else(|| lhs.len().cmp(&rhs.len())),
150        _ => abstract_literal_kind(lhs).cmp(&abstract_literal_kind(rhs)),
151    }
152}
153
154fn abstract_literal_kind(literal: &AbstractLiteral<Literal>) -> u8 {
155    match literal {
156        AbstractLiteral::Set(_) => 0,
157        AbstractLiteral::MSet(_) => 1,
158        AbstractLiteral::Matrix(..) => 2,
159        AbstractLiteral::Tuple(_) => 3,
160        AbstractLiteral::Record(_) => 4,
161        AbstractLiteral::Sequence(_) => 5,
162        AbstractLiteral::Function(_) => 6,
163        AbstractLiteral::Variant(_) => 7,
164        AbstractLiteral::Partition(_) => 8,
165        AbstractLiteral::Relation(_) => 9,
166        AbstractLiteral::Permutation(_) => 10,
167    }
168}
169
170fn literal_slice_cmp(lhs: &[Literal], rhs: &[Literal]) -> Ordering {
171    lhs.iter()
172        .zip(rhs)
173        .find_map(|(lhs, rhs)| {
174            let ordering = lhs.essence_cmp(rhs);
175            (ordering != Ordering::Equal).then_some(ordering)
176        })
177        .unwrap_or_else(|| lhs.len().cmp(&rhs.len()))
178}
179
180fn sorted_literals_cmp(lhs: &[Literal], rhs: &[Literal]) -> Ordering {
181    let mut lhs = lhs.iter().collect::<Vec<_>>();
182    let mut rhs = rhs.iter().collect::<Vec<_>>();
183    lhs.sort_by(|lhs, rhs| lhs.essence_cmp(rhs));
184    rhs.sort_by(|lhs, rhs| lhs.essence_cmp(rhs));
185    lhs.iter()
186        .zip(&rhs)
187        .find_map(|(lhs, rhs)| {
188            let ordering = lhs.essence_cmp(rhs);
189            (ordering != Ordering::Equal).then_some(ordering)
190        })
191        .unwrap_or_else(|| lhs.len().cmp(&rhs.len()))
192}
193
194/// Compare sets as occurrence vectors over the ordered union of their elements.
195fn set_essence_cmp(lhs: &[Literal], rhs: &[Literal]) -> Ordering {
196    let mut lhs = lhs.iter().collect::<Vec<_>>();
197    let mut rhs = rhs.iter().collect::<Vec<_>>();
198    lhs.sort_by(|lhs, rhs| lhs.essence_cmp(rhs));
199    rhs.sort_by(|lhs, rhs| lhs.essence_cmp(rhs));
200
201    let (mut lhs_index, mut rhs_index) = (0, 0);
202    while lhs_index < lhs.len() && rhs_index < rhs.len() {
203        match lhs[lhs_index].essence_cmp(rhs[rhs_index]) {
204            Ordering::Equal => {
205                lhs_index += 1;
206                rhs_index += 1;
207            }
208            // `lhs` contains the least differing element and `rhs` does not.
209            Ordering::Less => return Ordering::Greater,
210            // `rhs` contains the least differing element and `lhs` does not.
211            Ordering::Greater => return Ordering::Less,
212        }
213    }
214    match (lhs_index < lhs.len(), rhs_index < rhs.len()) {
215        (false, false) => Ordering::Equal,
216        (true, false) => Ordering::Greater,
217        (false, true) => Ordering::Less,
218        (true, true) => unreachable!(),
219    }
220}
221
222// make possible values of an AbstractLiteral a closed world to make the trait bounds more sane (particularly in Uniplate instances!!)
223pub trait AbstractLiteralValue:
224    Clone + Eq + PartialEq + Display + Uniplate + Biplate<Field<Self>> + 'static
225{
226    type Dom: Clone
227        + Eq
228        + PartialEq
229        + Debug
230        + Display
231        + Quine
232        + From<GroundDomain>
233        + Into<DomainPtr>;
234
235    /// Returns whether `domain` is the implicit one-based list index domain.
236    fn has_implied_list_domain(domain: &Self::Dom) -> bool;
237}
238impl AbstractLiteralValue for Expression {
239    type Dom = DomainPtr;
240
241    fn has_implied_list_domain(domain: &Self::Dom) -> bool {
242        let Domain::Ground(domain) = domain.as_ref() else {
243            return false;
244        };
245        matches!(domain.as_ref(), GroundDomain::Int(ranges, _) if ranges.as_slice() == [Range::UnboundedR(1)])
246    }
247}
248impl AbstractLiteralValue for Literal {
249    type Dom = Moo<GroundDomain>;
250
251    fn has_implied_list_domain(domain: &Self::Dom) -> bool {
252        matches!(domain.as_ref(), GroundDomain::Int(ranges, _) if ranges.as_slice() == [Range::UnboundedR(1)])
253    }
254}
255
256#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Quine)]
257#[path_prefix(conjure_cp::ast)]
258pub enum AbstractLiteral<T: AbstractLiteralValue> {
259    Set(Vec<T>),
260
261    MSet(Vec<T>),
262
263    /// A 1 dimensional matrix slice with an index domain.
264    Matrix(Vec<T>, T::Dom),
265
266    // a tuple of literals
267    Tuple(Vec<T>),
268
269    Record(Vec<Field<T>>),
270
271    Sequence(Vec<T>),
272
273    Function(Vec<(T, T)>),
274
275    // Variants only contain one of their name-domain pairs
276    Variant(Moo<Field<T>>),
277
278    // A list of partitions, each part has a set of values
279    Partition(Vec<Vec<T>>),
280    Relation(Vec<Vec<T>>),
281
282    /// Cycle notation for a permutation: each inner vec is one cycle. Unlike `Partition`, this is
283    /// *sparse* -- any element of the permutation's domain not mentioned in any cycle is an
284    /// implicit fixed point (maps to itself), rather than every element needing to be covered.
285    Permutation(Vec<Vec<T>>),
286}
287
288// TODO: use HasDomain instead once Expression::domain_of returns Domain not Option<Domain>
289fn union_item_domains(item_domains: Vec<DomainPtr>, literal_kind: &str) -> Option<DomainPtr> {
290    let mut item_domain_iter = item_domains.into_iter();
291    let first_item = item_domain_iter.next()?;
292
293    Some(
294        item_domain_iter
295            .try_fold(first_item, |x, y| x.union(&y))
296            .unwrap_or_else(|error| {
297                bug!(
298                    "taking the union of all item domains of a {literal_kind} literal should succeed: {error}"
299                )
300            }),
301    )
302}
303
304impl AbstractLiteral<Expression> {
305    pub fn domain_of(&self) -> Option<DomainPtr> {
306        match self {
307            AbstractLiteral::Set(items) => {
308                if items.is_empty() {
309                    return Some(Domain::set(
310                        SetAttr::new_size(0),
311                        Domain::empty(ReturnType::Unknown),
312                    ));
313                }
314                // ensure that all items have a domain, or return None
315                let item_domains: Vec<DomainPtr> = items
316                    .iter()
317                    .map(|x| x.domain_of())
318                    .collect::<Option<Vec<DomainPtr>>>()?;
319
320                // union all item domains together
321                let item_domain = union_item_domains(item_domains, "set")?;
322
323                let values = items.iter().map(eval_constant).collect::<Option<Vec<_>>>();
324                let high = i64::try_from(items.len()).ok()?;
325                let bounds = match values {
326                    Some(values) => {
327                        let n = i64::try_from(distinct_literal_count(&values)).ok()?;
328                        SizeBounds::new(n, Some(n))
329                    }
330                    None => SizeBounds::new(1, Some(high)),
331                };
332                Some(bounded_set(bounds, item_domain))
333            }
334
335            AbstractLiteral::MSet(items) => {
336                if items.is_empty() {
337                    return Some(Domain::mset(
338                        MSetAttr::new(Range::Single(0), Range::Single(0)),
339                        Domain::empty(ReturnType::Unknown),
340                    ));
341                }
342                let cardinality = i32::try_from(items.len()).ok()?;
343                // ensure that all items have a domain, or return None
344                let item_domains: Vec<DomainPtr> = items
345                    .iter()
346                    .map(|x| x.domain_of())
347                    .collect::<Option<Vec<DomainPtr>>>()?;
348
349                // union all item domains together
350                let item_domain = union_item_domains(item_domains, "mset")?;
351
352                let occurrence = if cardinality == 0 {
353                    Range::Single(0)
354                } else {
355                    Range::Bounded(1, cardinality)
356                };
357                Some(Domain::mset(
358                    MSetAttr::new(Range::Single(cardinality), occurrence),
359                    item_domain,
360                ))
361            }
362
363            AbstractLiteral::Sequence(elems) => {
364                let item_domains: Vec<DomainPtr> = elems
365                    .iter()
366                    .map(|x| x.domain_of())
367                    .collect::<Option<Vec<DomainPtr>>>()?;
368
369                // Get the union of all domains in the sequence.
370                // i.e. if <(1..3), (1..3), (5), (8..9)> then seq dom is (1..3, 5, 8..9)
371                let item_domain = union_item_domains(item_domains, "sequence")?;
372
373                // The literal's own length is its size: without it nothing downstream can tell
374                // how many positions the sequence has.
375                let size = Range::Single(i32::try_from(elems.len()).ok()?);
376                Some(Domain::sequence(
377                    SequenceAttr::<Int> {
378                        size,
379                        ..SequenceAttr::default()
380                    },
381                    item_domain,
382                ))
383            }
384
385            AbstractLiteral::Partition(items) => {
386                // Flatten the Vec<Vec< into a single vec
387                // ensure that all elemes in each part have a domain, or return None
388
389                let item_domains: Vec<DomainPtr> = items
390                    .iter()
391                    .flatten()
392                    .map(|x| x.domain_of())
393                    .collect::<Option<Vec<DomainPtr>>>()?;
394
395                // union all item domains together
396                let item_domain = union_item_domains(item_domains, "partition")?;
397
398                Some(Domain::partition(
399                    PartitionAttr::<Int>::default(),
400                    item_domain,
401                ))
402            }
403
404            AbstractLiteral::Permutation(cycles) => {
405                // Flatten the Vec<Vec< into a single vec; unlike partition, elements not
406                // mentioned in any cycle are implicit fixed points, so an empty literal (or one
407                // with no domain-bearing elements) has no way to infer an inner domain.
408                let item_domains: Vec<DomainPtr> = cycles
409                    .iter()
410                    .flatten()
411                    .map(|x| x.domain_of())
412                    .collect::<Option<Vec<DomainPtr>>>()?;
413
414                let item_domain = union_item_domains(item_domains, "permutation")?;
415
416                Some(Domain::permutation(
417                    PermutationAttr::<Int>::default(),
418                    item_domain,
419                ))
420            }
421
422            AbstractLiteral::Matrix(items, _) => {
423                let item_domains = items
424                    .iter()
425                    .map(|x| x.domain_of())
426                    .collect::<Option<Vec<DomainPtr>>>()?;
427
428                // union all item domains together
429                let mut item_domain = union_item_domains(item_domains, "matrix")?;
430
431                let mut new_index_domain = vec![];
432
433                // flatten index domains of n-d matrix into list
434                let mut matrix = self;
435                while let AbstractLiteral::Matrix(elems, idx) = matrix {
436                    bug_assert!(
437                        idx.as_matrix().is_none(),
438                        "n-dimensional matrix literals should be represented as a matrix inside a matrix, got {idx}"
439                    );
440                    new_index_domain.push(idx.clone());
441                    let Some(Expression::AbstractLiteral(_, inner @ AbstractLiteral::Matrix(..))) =
442                        elems.first()
443                    else {
444                        break;
445                    };
446                    matrix = inner;
447                }
448                // Each flattened literal dimension has already contributed to the row domain.
449                // Peel those dimensions so they are not represented twice. Matrix-valued
450                // expression leaves retain their own domains.
451                let mut remaining = new_index_domain.len() - 1;
452                while remaining > 0 {
453                    let (inner, indices) = item_domain.as_matrix()?;
454                    if indices.is_empty() {
455                        return None;
456                    }
457                    remaining = remaining.checked_sub(indices.len())?;
458                    item_domain = inner;
459                }
460                Some(Domain::matrix(item_domain, new_index_domain))
461            }
462            AbstractLiteral::Tuple(items) => Some(Domain::tuple(
463                items
464                    .iter()
465                    .map(Expression::domain_of)
466                    .collect::<Option<Vec<_>>>()?,
467            )),
468            AbstractLiteral::Record(items) => Some(Domain::record(
469                items
470                    .iter()
471                    .map(|field| {
472                        Some(super::Field {
473                            name: field.name.clone(),
474                            value: field.value.domain_of()?,
475                        })
476                    })
477                    .collect::<Option<Vec<_>>>()?,
478            )),
479            AbstractLiteral::Function(_) => None,
480            AbstractLiteral::Variant(_) => None,
481            AbstractLiteral::Relation(_) => None,
482        }
483    }
484}
485
486pub(super) fn distinct_literal_count(items: &[Literal]) -> usize {
487    let mut sorted = items.iter().collect::<Vec<_>>();
488    sorted.sort_by(|a, b| a.essence_cmp(b));
489    sorted.dedup_by(|a, b| a.essence_cmp(b).is_eq());
490    sorted.len()
491}
492
493impl HasDomain for AbstractLiteral<Literal> {
494    fn domain_of(&self) -> DomainPtr {
495        let domain = Domain::from_literal_vec(&[Literal::AbstractLiteral(self.clone())])
496            .expect("abstract literals should be correctly typed");
497        match self {
498            AbstractLiteral::Set(items) => {
499                let (_, inner) = domain.as_set().expect("set literal domain");
500                let n =
501                    i64::try_from(distinct_literal_count(items)).expect("literal fits in memory");
502                bounded_set(SizeBounds::new(n, Some(n)), inner)
503            }
504            AbstractLiteral::MSet(items) => {
505                let (_, inner) = domain.as_mset().expect("multiset literal domain");
506                let n = i32::try_from(items.len()).expect("literal cardinality fits i32");
507                let occurrence = if n == 0 {
508                    Range::Single(0)
509                } else {
510                    Range::Bounded(1, n)
511                };
512                Domain::mset(MSetAttr::new(Range::Single(n), occurrence), inner)
513            }
514            _ => domain,
515        }
516    }
517}
518
519impl Typeable for AbstractLiteral<Expression> {
520    fn return_type(&self) -> ReturnType {
521        match self {
522            AbstractLiteral::Set(items) if items.is_empty() => {
523                ReturnType::Set(Box::new(ReturnType::Unknown))
524            }
525            AbstractLiteral::Set(items) => {
526                let item_type = items[0].return_type();
527
528                // if any items do not have a type, return none.
529                let item_types: Vec<ReturnType> = items.iter().map(|x| x.return_type()).collect();
530
531                bug_assert!(
532                    item_types.iter().all(|x| x == &item_type),
533                    "all items in a set should have the same type"
534                );
535
536                ReturnType::Set(Box::new(item_type))
537            }
538            AbstractLiteral::MSet(items) if items.is_empty() => {
539                ReturnType::MSet(Box::new(ReturnType::Unknown))
540            }
541            AbstractLiteral::MSet(items) => {
542                let item_type = items[0].return_type();
543
544                // if any items do not have a type, return none.
545                let item_types: Vec<ReturnType> = items.iter().map(|x| x.return_type()).collect();
546
547                bug_assert!(
548                    item_types.iter().all(|x| x == &item_type),
549                    "all items in a set should have the same type"
550                );
551
552                ReturnType::MSet(Box::new(item_type))
553            }
554            AbstractLiteral::Sequence(items) if items.is_empty() => {
555                ReturnType::Sequence(Box::new(ReturnType::Unknown))
556            }
557            AbstractLiteral::Sequence(items) => {
558                let item_type = items[0].return_type();
559
560                // if any items do not have a type, return none.
561                let item_types: Vec<ReturnType> = items.iter().map(|x| x.return_type()).collect();
562
563                bug_assert!(
564                    item_types.iter().all(|x| x == &item_type),
565                    "all items in a sequence should have the same type"
566                );
567
568                ReturnType::Sequence(Box::new(item_type))
569            }
570            AbstractLiteral::Partition(items) if items.is_empty() || items[0].is_empty() => {
571                ReturnType::Partition(Box::new(ReturnType::Unknown))
572            }
573            AbstractLiteral::Partition(items) => {
574                let item_type = items[0][0].return_type();
575
576                // if any items do not have a type, return none.
577                let item_types: Vec<ReturnType> =
578                    items.iter().flatten().map(|x| x.return_type()).collect();
579
580                bug_assert!(
581                    item_types.iter().all(|x| x == &item_type),
582                    "all items in every part of a partition should have the same type"
583                );
584
585                ReturnType::Partition(Box::new(item_type))
586            }
587            AbstractLiteral::Permutation(items) if items.is_empty() || items[0].is_empty() => {
588                ReturnType::Permutation(Box::new(ReturnType::Unknown))
589            }
590            AbstractLiteral::Permutation(items) => {
591                let item_type = items[0][0].return_type();
592
593                let item_types: Vec<ReturnType> =
594                    items.iter().flatten().map(|x| x.return_type()).collect();
595
596                bug_assert!(
597                    item_types.iter().all(|x| x == &item_type),
598                    "all items in every cycle of a permutation should have the same type"
599                );
600
601                ReturnType::Permutation(Box::new(item_type))
602            }
603            AbstractLiteral::Matrix(items, _) if items.is_empty() => {
604                ReturnType::Matrix(Box::new(ReturnType::Unknown))
605            }
606            AbstractLiteral::Matrix(items, _) => {
607                let item_type = items[0].return_type();
608
609                // if any items do not have a type, return none.
610                let item_types: Vec<ReturnType> = items.iter().map(|x| x.return_type()).collect();
611
612                bug_assert!(
613                    item_types.iter().all(|x| x == &item_type),
614                    "all items in a matrix should have the same type. items: {items} types: {types:#?}",
615                    items = pretty_vec(items),
616                    types = items
617                        .iter()
618                        .map(|x| x.return_type())
619                        .collect::<Vec<ReturnType>>()
620                );
621
622                ReturnType::Matrix(Box::new(item_type))
623            }
624            AbstractLiteral::Tuple(items) => {
625                let mut item_types = vec![];
626                for item in items {
627                    item_types.push(item.return_type());
628                }
629                ReturnType::Tuple(item_types)
630            }
631            AbstractLiteral::Record(items) => {
632                let mut item_types = vec![];
633                for item in items {
634                    item_types.push(item.clone().func_map(|x| x.return_type()));
635                }
636                item_types.sort();
637                ReturnType::Record(item_types)
638            }
639            AbstractLiteral::Function(items) => {
640                if items.is_empty() {
641                    return ReturnType::Function(
642                        Box::new(ReturnType::Unknown),
643                        Box::new(ReturnType::Unknown),
644                    );
645                }
646
647                // Check that all items have the same return type
648                let (x1, y1) = &items[0];
649                let (t1, t2) = (x1.return_type(), y1.return_type());
650                for (x, y) in items {
651                    let (tx, ty) = (x.return_type(), y.return_type());
652                    if tx != t1 {
653                        bug!("Expected {t1}, got {x}: {tx}");
654                    }
655                    if ty != t2 {
656                        bug!("Expected {t2}, got {y}: {ty}");
657                    }
658                }
659
660                ReturnType::Function(Box::new(t1), Box::new(t2))
661            }
662            AbstractLiteral::Variant(item) => {
663                // Variants hold multiple possible types. In the case of a literal we know which type it chose
664                ReturnType::Variant(vec![item.as_ref().clone().func_map(|x| x.return_type())])
665            }
666            AbstractLiteral::Relation(items) => {
667                if items.is_empty() {
668                    return ReturnType::Relation(vec![ReturnType::Unknown]);
669                }
670                let mut item_types = vec![];
671                let x1 = &items[0];
672                let size = x1.len();
673                for item in x1 {
674                    item_types.push(item.return_type());
675                }
676                for x in items {
677                    if x.len() != size {
678                        let strs = item_types.iter().map(|x| format!("{}", x)).join(",");
679                        bug!("Expected ({strs}) with length {size}, got size {}", x.len());
680                    }
681                    for i in 1..size {
682                        if let Some(new_type) = x.get(i)
683                            && let Some(old_type) = item_types.get(i)
684                            && new_type.return_type() != *old_type
685                        {
686                            bug!("Expected {old_type}, got {new_type}");
687                        }
688                    }
689                }
690                ReturnType::Relation(item_types)
691            }
692        }
693    }
694}
695
696impl<T> AbstractLiteral<T>
697where
698    T: AbstractLiteralValue,
699{
700    /// Creates a matrix with elements `elems`, with domain `int(1..)`.
701    ///
702    /// This acts as a variable sized list.
703    pub fn matrix_implied_indices(elems: Vec<T>) -> Self {
704        AbstractLiteral::Matrix(
705            elems,
706            GroundDomain::Int(vec![Range::UnboundedR(1)], None).into(),
707        )
708    }
709
710    /// If the AbstractLiteral is a list, returns its elements.
711    ///
712    /// A list is any a matrix with the domain `int(1..)`. This includes matrix literals without
713    /// any explicitly specified domain.
714    pub fn unwrap_list(&self) -> Option<&Vec<T>> {
715        let AbstractLiteral::Matrix(elems, domain) = self else {
716            return None;
717        };
718
719        T::has_implied_list_domain(domain).then_some(elems)
720    }
721
722    /// If this abstract literal is a list, consumes it and returns its elements.
723    pub fn into_list(self) -> Option<Vec<T>> {
724        let AbstractLiteral::Matrix(elems, domain) = self else {
725            return None;
726        };
727
728        T::has_implied_list_domain(&domain).then_some(elems)
729    }
730}
731
732impl<T> Display for AbstractLiteral<T>
733where
734    T: AbstractLiteralValue,
735{
736    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
737        match self {
738            AbstractLiteral::Set(elems) => {
739                let elems_str: String = elems.iter().map(|x| format!("{x}")).join(",");
740                write!(f, "{{{elems_str}}}")
741            }
742            AbstractLiteral::MSet(elems) => {
743                let elems_str: String = elems.iter().map(|x| format!("{x}")).join(",");
744                write!(f, "mset({elems_str})")
745            }
746            AbstractLiteral::Matrix(elems, index_domain) => {
747                let elems_str: String = elems.iter().map(|x| format!("{x}")).join(",");
748                write!(f, "[{elems_str};{index_domain}]")
749            }
750            AbstractLiteral::Tuple(elems) => {
751                let elems_str: String = elems.iter().map(|x| format!("{x}")).join(",");
752                write!(f, "({elems_str})")
753            }
754            AbstractLiteral::Sequence(elems) => {
755                let elems_str: String = elems.iter().map(|x| format!("{x}")).join(",");
756                write!(f, "sequence({elems_str})")
757            }
758            AbstractLiteral::Partition(parts) => {
759                let elems_str: String = parts
760                    .iter()
761                    .map(|inner| {
762                        let elems_str = inner.iter().map(|x| format!("{x}")).join(",");
763                        format!("{{{}}}", elems_str)
764                    })
765                    .join(", ");
766
767                write!(f, "partition({elems_str})")
768            }
769            AbstractLiteral::Permutation(cycles) => {
770                let cycles_str: String = cycles
771                    .iter()
772                    .map(|cycle| {
773                        let elems_str = cycle.iter().map(|x| format!("{x}")).join(",");
774                        format!("({elems_str})")
775                    })
776                    .join("");
777
778                write!(f, "permutation{cycles_str}")
779            }
780            AbstractLiteral::Record(entries) => {
781                let entries_str: String = entries
782                    .iter()
783                    .map(|entry| format!("{} = {}", entry.name, entry.value))
784                    .join(",");
785                write!(f, "record {{{entries_str}}}")
786            }
787            AbstractLiteral::Function(entries) => {
788                let entries_str: String = entries
789                    .iter()
790                    .map(|entry| format!("{} --> {}", entry.0, entry.1))
791                    .join(",");
792                write!(f, "function({entries_str})")
793            }
794            AbstractLiteral::Variant(entry) => {
795                write!(f, "variant{{{} = {}}}", entry.name, entry.value)
796            }
797            AbstractLiteral::Relation(elems) => {
798                let elems_str: String = elems
799                    .iter()
800                    .map(|x| format!("({})", x.iter().map(|x| format!("{x}")).join(",")))
801                    .join(",");
802                write!(f, "relation({elems_str})")
803            }
804        }
805    }
806}
807
808impl<T> Uniplate for AbstractLiteral<T>
809where
810    T: AbstractLiteralValue + Biplate<AbstractLiteral<T>>,
811{
812    fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
813        // walking into T
814        match self {
815            AbstractLiteral::Set(vec) => {
816                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(vec);
817                (f1_tree, Box::new(move |x| AbstractLiteral::Set(f1_ctx(x))))
818            }
819            AbstractLiteral::MSet(vec) => {
820                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(vec);
821                (f1_tree, Box::new(move |x| AbstractLiteral::MSet(f1_ctx(x))))
822            }
823            AbstractLiteral::Matrix(elems, index_domain) => {
824                let index_domain = index_domain.clone();
825                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(elems);
826                (
827                    f1_tree,
828                    Box::new(move |x| AbstractLiteral::Matrix(f1_ctx(x), index_domain.clone())),
829                )
830            }
831            AbstractLiteral::Sequence(vec) => {
832                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(vec);
833                (
834                    f1_tree,
835                    Box::new(move |x| AbstractLiteral::Sequence(f1_ctx(x))),
836                )
837            }
838            AbstractLiteral::Tuple(elems) => {
839                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(elems);
840                (
841                    f1_tree,
842                    Box::new(move |x| AbstractLiteral::Tuple(f1_ctx(x))),
843                )
844            }
845            AbstractLiteral::Record(entries) => {
846                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(entries);
847                (
848                    f1_tree,
849                    Box::new(move |x| AbstractLiteral::Record(f1_ctx(x))),
850                )
851            }
852            AbstractLiteral::Function(entries) => {
853                let entry_count = entries.len();
854                let flattened: Vec<T> = entries
855                    .iter()
856                    .flat_map(|(lhs, rhs)| [lhs.clone(), rhs.clone()])
857                    .collect();
858
859                let (f1_tree, f1_ctx) =
860                    <Vec<T> as Biplate<AbstractLiteral<T>>>::biplate(&flattened);
861                (
862                    f1_tree,
863                    Box::new(move |x| {
864                        let rebuilt = f1_ctx(x);
865                        assert_eq!(
866                            rebuilt.len(),
867                            entry_count * 2,
868                            "number of function literal children should remain unchanged"
869                        );
870
871                        let mut iter = rebuilt.into_iter();
872                        let mut pairs = Vec::with_capacity(entry_count);
873                        while let (Some(lhs), Some(rhs)) = (iter.next(), iter.next()) {
874                            pairs.push((lhs, rhs));
875                        }
876
877                        AbstractLiteral::Function(pairs)
878                    }),
879                )
880            }
881            AbstractLiteral::Variant(entries) => {
882                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(entries);
883                (
884                    f1_tree,
885                    Box::new(move |x| AbstractLiteral::Variant(f1_ctx(x))),
886                )
887            }
888            AbstractLiteral::Relation(elems) => {
889                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(elems);
890                (
891                    f1_tree,
892                    Box::new(move |x| AbstractLiteral::Relation(f1_ctx(x))),
893                )
894            }
895            AbstractLiteral::Partition(elems) => {
896                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(elems);
897                (
898                    f1_tree,
899                    Box::new(move |x| AbstractLiteral::Partition(f1_ctx(x))),
900                )
901            }
902            AbstractLiteral::Permutation(elems) => {
903                let (f1_tree, f1_ctx) = <_ as Biplate<AbstractLiteral<T>>>::biplate(elems);
904                (
905                    f1_tree,
906                    Box::new(move |x| AbstractLiteral::Permutation(f1_ctx(x))),
907                )
908            }
909        }
910    }
911}
912
913impl<U, To> Biplate<To> for AbstractLiteral<U>
914where
915    To: Uniplate,
916    U: AbstractLiteralValue + Biplate<AbstractLiteral<U>> + Biplate<To>,
917    Field<U>: Biplate<AbstractLiteral<U>> + Biplate<To>,
918{
919    fn biplate(&self) -> (Tree<To>, Box<dyn Fn(Tree<To>) -> Self>) {
920        if std::any::TypeId::of::<To>() == std::any::TypeId::of::<AbstractLiteral<U>>() {
921            // To ==From => return One(self)
922
923            unsafe {
924                // SAFETY: asserted the type equality above
925                let self_to = std::mem::transmute::<&AbstractLiteral<U>, &To>(self).clone();
926                let tree = Tree::One(self_to);
927                let ctx = Box::new(move |x| {
928                    let Tree::One(x) = x else {
929                        panic!();
930                    };
931
932                    std::mem::transmute::<&To, &AbstractLiteral<U>>(&x).clone()
933                });
934
935                (tree, ctx)
936            }
937        } else {
938            // walking into T
939            match self {
940                AbstractLiteral::Set(vec) => {
941                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(vec);
942                    (f1_tree, Box::new(move |x| AbstractLiteral::Set(f1_ctx(x))))
943                }
944                AbstractLiteral::MSet(vec) => {
945                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(vec);
946                    (f1_tree, Box::new(move |x| AbstractLiteral::MSet(f1_ctx(x))))
947                }
948                AbstractLiteral::Matrix(elems, index_domain) => {
949                    let index_domain = index_domain.clone();
950                    let (f1_tree, f1_ctx) = <Vec<U> as Biplate<To>>::biplate(elems);
951                    (
952                        f1_tree,
953                        Box::new(move |x| AbstractLiteral::Matrix(f1_ctx(x), index_domain.clone())),
954                    )
955                }
956                AbstractLiteral::Sequence(vec) => {
957                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(vec);
958                    (
959                        f1_tree,
960                        Box::new(move |x| AbstractLiteral::Sequence(f1_ctx(x))),
961                    )
962                }
963                AbstractLiteral::Tuple(elems) => {
964                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(elems);
965                    (
966                        f1_tree,
967                        Box::new(move |x| AbstractLiteral::Tuple(f1_ctx(x))),
968                    )
969                }
970                AbstractLiteral::Record(entries) => {
971                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(entries);
972                    (
973                        f1_tree,
974                        Box::new(move |x| AbstractLiteral::Record(f1_ctx(x))),
975                    )
976                }
977                AbstractLiteral::Function(entries) => {
978                    let entry_count = entries.len();
979                    let flattened: Vec<U> = entries
980                        .iter()
981                        .flat_map(|(lhs, rhs)| [lhs.clone(), rhs.clone()])
982                        .collect();
983
984                    let (f1_tree, f1_ctx) = <Vec<U> as Biplate<To>>::biplate(&flattened);
985                    (
986                        f1_tree,
987                        Box::new(move |x| {
988                            let rebuilt = f1_ctx(x);
989                            assert_eq!(
990                                rebuilt.len(),
991                                entry_count * 2,
992                                "number of function literal children should remain unchanged"
993                            );
994
995                            let mut iter = rebuilt.into_iter();
996                            let mut pairs = Vec::with_capacity(entry_count);
997                            while let (Some(lhs), Some(rhs)) = (iter.next(), iter.next()) {
998                                pairs.push((lhs, rhs));
999                            }
1000
1001                            AbstractLiteral::Function(pairs)
1002                        }),
1003                    )
1004                }
1005                AbstractLiteral::Variant(entries) => {
1006                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(entries);
1007                    (
1008                        f1_tree,
1009                        Box::new(move |x| AbstractLiteral::Variant(f1_ctx(x))),
1010                    )
1011                }
1012                AbstractLiteral::Relation(elems) => {
1013                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(elems);
1014                    (
1015                        f1_tree,
1016                        Box::new(move |x| AbstractLiteral::Relation(f1_ctx(x))),
1017                    )
1018                }
1019                AbstractLiteral::Partition(elems) => {
1020                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(elems);
1021                    (
1022                        f1_tree,
1023                        Box::new(move |x| AbstractLiteral::Partition(f1_ctx(x))),
1024                    )
1025                }
1026                AbstractLiteral::Permutation(elems) => {
1027                    let (f1_tree, f1_ctx) = <_ as Biplate<To>>::biplate(elems);
1028                    (
1029                        f1_tree,
1030                        Box::new(move |x| AbstractLiteral::Permutation(f1_ctx(x))),
1031                    )
1032                }
1033            }
1034        }
1035    }
1036
1037    fn children_bi_count(&self) -> usize {
1038        // Manual `biplate` (not derive): delegate counts to the plated containers so wide
1039        // matrices/lists stay O(1) instead of materialising `children_bi`.
1040        if std::any::TypeId::of::<To>() == std::any::TypeId::of::<AbstractLiteral<U>>() {
1041            return 1;
1042        }
1043        match self {
1044            AbstractLiteral::Set(v)
1045            | AbstractLiteral::MSet(v)
1046            | AbstractLiteral::Sequence(v)
1047            | AbstractLiteral::Tuple(v)
1048            | AbstractLiteral::Matrix(v, _) => <Vec<U> as Biplate<To>>::children_bi_count(v),
1049            AbstractLiteral::Record(entries) => {
1050                <Vec<Field<U>> as Biplate<To>>::children_bi_count(entries)
1051            }
1052            AbstractLiteral::Variant(entry) => {
1053                <Moo<Field<U>> as Biplate<To>>::children_bi_count(entry)
1054            }
1055            AbstractLiteral::Relation(elems)
1056            | AbstractLiteral::Partition(elems)
1057            | AbstractLiteral::Permutation(elems) => {
1058                <Vec<Vec<U>> as Biplate<To>>::children_bi_count(elems)
1059            }
1060            AbstractLiteral::Function(_) => <Self as Biplate<To>>::children_bi(self).len(),
1061        }
1062    }
1063
1064    fn try_replace_child_at_bi(&mut self, index: usize, child: To) -> bool {
1065        // Same as `children_bi_count`: keep in-place updates for owned vectors (lee-distance).
1066        if std::any::TypeId::of::<To>() == std::any::TypeId::of::<AbstractLiteral<U>>() {
1067            if index != 0 {
1068                return false;
1069            }
1070            // SAFETY: TypeId equality means To and AbstractLiteral<U> are the same type.
1071            unsafe {
1072                let child_as_self = std::mem::transmute_copy::<To, AbstractLiteral<U>>(&child);
1073                std::mem::forget(child);
1074                *self = child_as_self;
1075            }
1076            return true;
1077        }
1078        match self {
1079            AbstractLiteral::Set(v)
1080            | AbstractLiteral::MSet(v)
1081            | AbstractLiteral::Sequence(v)
1082            | AbstractLiteral::Tuple(v)
1083            | AbstractLiteral::Matrix(v, _) => {
1084                <Vec<U> as Biplate<To>>::try_replace_child_at_bi(v, index, child)
1085            }
1086            AbstractLiteral::Record(entries) => {
1087                <Vec<Field<U>> as Biplate<To>>::try_replace_child_at_bi(entries, index, child)
1088            }
1089            AbstractLiteral::Variant(entry) => {
1090                <Moo<Field<U>> as Biplate<To>>::try_replace_child_at_bi(entry, index, child)
1091            }
1092            AbstractLiteral::Relation(elems)
1093            | AbstractLiteral::Partition(elems)
1094            | AbstractLiteral::Permutation(elems) => {
1095                <Vec<Vec<U>> as Biplate<To>>::try_replace_child_at_bi(elems, index, child)
1096            }
1097            AbstractLiteral::Function(_) => {
1098                let mut children = <Self as Biplate<To>>::children_bi(self);
1099                if index >= children.len() {
1100                    return false;
1101                }
1102                children[index] = child;
1103                *self = self.with_children_bi(children);
1104                true
1105            }
1106        }
1107    }
1108}
1109
1110impl TryFrom<Literal> for i32 {
1111    type Error = &'static str;
1112
1113    fn try_from(value: Literal) -> Result<Self, Self::Error> {
1114        match value {
1115            Literal::Int(i) => Ok(i),
1116            _ => Err("Cannot convert non-i32 literal to i32"),
1117        }
1118    }
1119}
1120
1121impl TryFrom<Box<Literal>> for i32 {
1122    type Error = &'static str;
1123
1124    fn try_from(value: Box<Literal>) -> Result<Self, Self::Error> {
1125        (*value).try_into()
1126    }
1127}
1128
1129impl TryFrom<&Box<Literal>> for i32 {
1130    type Error = &'static str;
1131
1132    fn try_from(value: &Box<Literal>) -> Result<Self, Self::Error> {
1133        TryFrom::<&Literal>::try_from(value.as_ref())
1134    }
1135}
1136
1137impl TryFrom<&Moo<Literal>> for i32 {
1138    type Error = &'static str;
1139
1140    fn try_from(value: &Moo<Literal>) -> Result<Self, Self::Error> {
1141        TryFrom::<&Literal>::try_from(value.as_ref())
1142    }
1143}
1144
1145impl TryFrom<&Literal> for i32 {
1146    type Error = &'static str;
1147
1148    fn try_from(value: &Literal) -> Result<Self, Self::Error> {
1149        match value {
1150            Literal::Int(i) => Ok(*i),
1151            _ => Err("Cannot convert non-i32 literal to i32"),
1152        }
1153    }
1154}
1155
1156impl TryFrom<Literal> for bool {
1157    type Error = &'static str;
1158
1159    fn try_from(value: Literal) -> Result<Self, Self::Error> {
1160        match value {
1161            Literal::Bool(b) => Ok(b),
1162            _ => Err("Cannot convert non-bool literal to bool"),
1163        }
1164    }
1165}
1166
1167impl TryFrom<&Literal> for bool {
1168    type Error = &'static str;
1169
1170    fn try_from(value: &Literal) -> Result<Self, Self::Error> {
1171        match value {
1172            Literal::Bool(b) => Ok(*b),
1173            _ => Err("Cannot convert non-bool literal to bool"),
1174        }
1175    }
1176}
1177
1178impl TryFrom<Literal> for EnumVariant {
1179    type Error = &'static str;
1180
1181    fn try_from(value: Literal) -> Result<Self, Self::Error> {
1182        match value {
1183            Literal::EnumVariant(v) => Ok(v),
1184            _ => Err("Cannot convert non-enum literal to enum variant"),
1185        }
1186    }
1187}
1188
1189impl TryFrom<&Literal> for EnumVariant {
1190    type Error = &'static str;
1191
1192    fn try_from(value: &Literal) -> Result<Self, Self::Error> {
1193        match value {
1194            Literal::EnumVariant(v) => Ok(v.clone()),
1195            _ => Err("Cannot convert non-enum literal to enum variant"),
1196        }
1197    }
1198}
1199
1200impl From<i32> for Literal {
1201    fn from(i: i32) -> Self {
1202        Literal::Int(i)
1203    }
1204}
1205
1206impl From<bool> for Literal {
1207    fn from(b: bool) -> Self {
1208        Literal::Bool(b)
1209    }
1210}
1211
1212impl From<Literal> for Ustr {
1213    fn from(value: Literal) -> Self {
1214        // TODO: avoid the temporary-allocation of a string by format! here?
1215        Ustr::from(&format!("{value}"))
1216    }
1217}
1218
1219impl From<AbstractLiteral<Literal>> for Literal {
1220    fn from(literal: AbstractLiteral<Literal>) -> Self {
1221        Literal::AbstractLiteral(literal)
1222    }
1223}
1224
1225impl AbstractLiteral<Expression> {
1226    /// If all the elements are literals, returns this as an AbstractLiteral<Literal>.
1227    /// Otherwise, returns `None`.
1228    pub fn into_literals(self) -> Option<AbstractLiteral<Literal>> {
1229        match self {
1230            AbstractLiteral::Set(elements) => {
1231                let literals = elements
1232                    .into_iter()
1233                    .map(|expr| match expr {
1234                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
1235                        Expression::AbstractLiteral(_, abslit) => {
1236                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
1237                        }
1238                        _ => None,
1239                    })
1240                    .collect::<Option<Vec<_>>>()?;
1241                Some(AbstractLiteral::Set(literals))
1242            }
1243            AbstractLiteral::MSet(elements) => {
1244                let literals = elements
1245                    .into_iter()
1246                    .map(|expr| match expr {
1247                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
1248                        Expression::AbstractLiteral(_, abslit) => {
1249                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
1250                        }
1251                        _ => None,
1252                    })
1253                    .collect::<Option<Vec<_>>>()?;
1254                Some(AbstractLiteral::MSet(literals))
1255            }
1256            AbstractLiteral::Partition(elems) => {
1257                // want to ascertain if every elem in Vec<Vec<Expr>> is a literal. If any are not, return none
1258                // otherwise confirm it is an abslit<lit>
1259                let mut partition: Vec<Vec<_>> = Vec::new();
1260
1261                for part in elems {
1262                    let literals = part
1263                        .into_iter()
1264                        .map(|expr| match expr {
1265                            Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
1266                            Expression::AbstractLiteral(_, abslit) => {
1267                                Some(Literal::AbstractLiteral(abslit.into_literals()?))
1268                            }
1269                            _ => None,
1270                        })
1271                        .collect::<Option<Vec<_>>>()?;
1272
1273                    partition.push(literals);
1274                }
1275
1276                Some(AbstractLiteral::Partition(partition))
1277            }
1278            AbstractLiteral::Permutation(elems) => {
1279                let mut permutation: Vec<Vec<_>> = Vec::new();
1280
1281                for cycle in elems {
1282                    let literals = cycle
1283                        .into_iter()
1284                        .map(|expr| match expr {
1285                            Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
1286                            Expression::AbstractLiteral(_, abslit) => {
1287                                Some(Literal::AbstractLiteral(abslit.into_literals()?))
1288                            }
1289                            _ => None,
1290                        })
1291                        .collect::<Option<Vec<_>>>()?;
1292
1293                    permutation.push(literals);
1294                }
1295
1296                Some(AbstractLiteral::Permutation(permutation))
1297            }
1298            AbstractLiteral::Matrix(items, domain) => {
1299                let mut literals = vec![];
1300                for item in items {
1301                    let literal = match item {
1302                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
1303                        Expression::AbstractLiteral(_, abslit) => {
1304                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
1305                        }
1306                        _ => None,
1307                    }?;
1308                    literals.push(literal);
1309                }
1310
1311                Some(AbstractLiteral::Matrix(literals, domain.resolve().ok()?))
1312            }
1313            AbstractLiteral::Sequence(elements) => {
1314                let literals = elements
1315                    .into_iter()
1316                    .map(|expr| match expr {
1317                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
1318                        Expression::AbstractLiteral(_, abslit) => {
1319                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
1320                        }
1321                        _ => None,
1322                    })
1323                    .collect::<Option<Vec<_>>>()?;
1324                Some(AbstractLiteral::Sequence(literals))
1325            }
1326            AbstractLiteral::Tuple(items) => {
1327                let mut literals = vec![];
1328                for item in items {
1329                    let literal = match item {
1330                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
1331                        Expression::AbstractLiteral(_, abslit) => {
1332                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
1333                        }
1334                        _ => None,
1335                    }?;
1336                    literals.push(literal);
1337                }
1338
1339                Some(AbstractLiteral::Tuple(literals))
1340            }
1341            AbstractLiteral::Record(entries) => {
1342                let mut literals = vec![];
1343                for entry in entries {
1344                    let literal = match entry.value {
1345                        Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
1346                        Expression::AbstractLiteral(_, abslit) => {
1347                            Some(Literal::AbstractLiteral(abslit.into_literals()?))
1348                        }
1349                        _ => None,
1350                    }?;
1351
1352                    literals.push((entry.name, literal));
1353                }
1354                Some(AbstractLiteral::Record(
1355                    literals
1356                        .into_iter()
1357                        .map(|(name, literal)| Field {
1358                            name,
1359                            value: literal,
1360                        })
1361                        .collect(),
1362                ))
1363            }
1364            AbstractLiteral::Function(_) => todo!("Implement into_literals for functions"),
1365            AbstractLiteral::Variant(entry) => {
1366                let literal = match entry.value.clone() {
1367                    Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
1368                    Expression::AbstractLiteral(_, abslit) => {
1369                        Some(Literal::AbstractLiteral(abslit.into_literals()?))
1370                    }
1371                    _ => None,
1372                }?;
1373                Some(AbstractLiteral::Variant(Moo::new(Field {
1374                    name: entry.name.clone(),
1375                    value: literal,
1376                })))
1377            }
1378            AbstractLiteral::Relation(tuples) => {
1379                let mut literal_tuples = Vec::with_capacity(tuples.len());
1380                for fields in tuples {
1381                    let literals = fields
1382                        .into_iter()
1383                        .map(|expr| match expr {
1384                            Expression::Atomic(_, Atom::Literal(lit)) => Some(lit),
1385                            Expression::AbstractLiteral(_, abslit) => {
1386                                Some(Literal::AbstractLiteral(abslit.into_literals()?))
1387                            }
1388                            _ => None,
1389                        })
1390                        .collect::<Option<Vec<_>>>()?;
1391                    literal_tuples.push(literals);
1392                }
1393                Some(AbstractLiteral::Relation(literal_tuples))
1394            }
1395        }
1396    }
1397}
1398
1399// need display implementations for other types as well
1400impl Display for Literal {
1401    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1402        match &self {
1403            Literal::Int(i) => write!(f, "{i}"),
1404            Literal::Bool(b) => write!(f, "{b}"),
1405            Literal::EnumVariant(ev) => write!(f, "{}", ev.name()),
1406            Literal::Unnamed(value) => write!(f, "{}_{}", value.ty.name, value.value),
1407            Literal::AbstractLiteral(l) => write!(f, "{l}"),
1408        }
1409    }
1410}
1411
1412#[cfg(test)]
1413mod tests {
1414
1415    use super::*;
1416    use crate::ast::matrix::{flatten, partial_flatten, shape_of};
1417    use crate::ast::{DeclarationPtr, Metadata, Name};
1418    use crate::{domain_int_ground, into_matrix, matrix, matrix_lit, range};
1419    use uniplate::Uniplate;
1420
1421    #[test]
1422    fn essence_value_ordering_uses_occurrence_lex_for_sets() {
1423        let set = |values: &[i32]| {
1424            Literal::AbstractLiteral(AbstractLiteral::Set(
1425                values.iter().copied().map(Literal::Int).collect(),
1426            ))
1427        };
1428        let ordered = [
1429            set(&[]),
1430            set(&[3]),
1431            set(&[2]),
1432            set(&[2, 3]),
1433            set(&[1]),
1434            set(&[1, 3]),
1435            set(&[1, 2]),
1436            set(&[1, 2, 3]),
1437        ];
1438
1439        for pair in ordered.windows(2) {
1440            assert_eq!(pair[0].essence_cmp(&pair[1]), Ordering::Less);
1441        }
1442    }
1443
1444    #[test]
1445    fn essence_value_ordering_is_lexicographic_for_tuples_and_matrices() {
1446        let tuple = |values: &[i32]| {
1447            Literal::AbstractLiteral(AbstractLiteral::Tuple(
1448                values.iter().copied().map(Literal::Int).collect(),
1449            ))
1450        };
1451        assert_eq!(tuple(&[1, 2]).essence_cmp(&tuple(&[1, 3])), Ordering::Less);
1452        assert_eq!(
1453            tuple(&[1, 2]).essence_cmp(&tuple(&[1, 2, 0])),
1454            Ordering::Less
1455        );
1456
1457        let matrix = |values: &[i32]| {
1458            Literal::AbstractLiteral(AbstractLiteral::Matrix(
1459                values.iter().copied().map(Literal::Int).collect(),
1460                Moo::new(GroundDomain::Int(
1461                    vec![Range::Bounded(1, values.len() as i32)],
1462                    None,
1463                )),
1464            ))
1465        };
1466        assert_eq!(
1467            matrix(&[1, 2]).essence_cmp(&matrix(&[1, 3])),
1468            Ordering::Less
1469        );
1470    }
1471
1472    #[test]
1473    fn mset_domain_accepts_domain_letting_references() {
1474        let declaration = DeclarationPtr::new_domain_letting(
1475            Name::user("NUM"),
1476            Domain::int(vec![Range::Bounded(1, 999)]),
1477        );
1478        let alias = Domain::reference(declaration).unwrap();
1479        let item = Expression::DomainAnnotation(
1480            Metadata::new(),
1481            Moo::new(Expression::Atomic(
1482                Metadata::new(),
1483                Atom::Literal(Literal::Int(1)),
1484            )),
1485            alias,
1486        );
1487
1488        let domain = AbstractLiteral::MSet(vec![item.clone(), item])
1489            .domain_of()
1490            .unwrap();
1491        let (_, item_domain) = domain.as_mset().unwrap();
1492
1493        assert_eq!(
1494            item_domain.resolve(),
1495            Ok(Moo::new(GroundDomain::Int(
1496                vec![Range::Bounded(1, 999)],
1497                None
1498            )))
1499        );
1500    }
1501
1502    #[test]
1503    fn matrix_uniplate_universe() {
1504        // Can we traverse through matrices with uniplate?
1505        let my_matrix: AbstractLiteral<Literal> = into_matrix![
1506            vec![Literal::AbstractLiteral(matrix![Literal::Bool(true);Moo::new(GroundDomain::Bool)]); 5];
1507            Moo::new(GroundDomain::Bool)
1508        ];
1509
1510        let expected_index_domains = vec![Moo::new(GroundDomain::Bool); 6];
1511        let actual_index_domains: Vec<Moo<GroundDomain>> =
1512            my_matrix.cata(&move |elem, children| {
1513                let mut res = vec![];
1514                res.extend(children.into_iter().flatten());
1515                if let AbstractLiteral::Matrix(_, index_domain) = elem {
1516                    res.push(index_domain);
1517                }
1518
1519                res
1520            });
1521
1522        assert_eq!(actual_index_domains, expected_index_domains);
1523    }
1524
1525    #[test]
1526    fn matrix_flatten() {
1527        let tensor: AbstractLiteral<Literal> = matrix![
1528            [
1529                // batch 1
1530                [1, 2, 3, 4],
1531                [5, 6, 7, 8],
1532                [9, 10, 11, 12]
1533            ],
1534            [
1535                // batch 2
1536                [13, 14, 15, 16],
1537                [17, 18, 19, 20],
1538                [21, 22, 23, 24]
1539            ]
1540        ];
1541
1542        let actual_elems: Vec<Literal> = flatten(&tensor).cloned().collect();
1543        let expected_elems = (1..25).map(Literal::from).collect::<Vec<_>>();
1544        assert_eq!(actual_elems, expected_elems);
1545    }
1546
1547    #[test]
1548    fn matrix_domain_1d() {
1549        let matrix = matrix_lit![10, 11, 12, 13; domain_int_ground!(1..4)];
1550        let dom = matrix.domain_of();
1551
1552        let (inner_dom, idx_doms) = dom.as_matrix_ground().expect("must be ground matrix");
1553        assert_eq!(inner_dom, &domain_int_ground!(10..13));
1554        assert_eq!(idx_doms.len(), 1);
1555        assert_eq!(&idx_doms[0], &domain_int_ground!(1..4));
1556    }
1557
1558    #[test]
1559    fn matrix_domain_2d() {
1560        let matrix = matrix_lit![
1561            [1, 2, 3, 4],
1562            [5, 6, 7, 8];
1563            [
1564                domain_int_ground!(1..2),
1565                domain_int_ground!(1..4)
1566            ]
1567        ];
1568        let dom = matrix.domain_of();
1569
1570        let (inner_dom, idx_doms) = dom.as_matrix_ground().expect("must be ground matrix");
1571        assert_eq!(inner_dom, &domain_int_ground!(1..8));
1572        assert_eq!(idx_doms.len(), 2);
1573        assert_eq!(&idx_doms[0], &domain_int_ground!(1..2));
1574        assert_eq!(&idx_doms[1], &domain_int_ground!(1..4));
1575    }
1576
1577    #[test]
1578    fn matrix_shape_3d() {
1579        let tensor: AbstractLiteral<Literal> = matrix![
1580            [
1581                [1, 2, 3, 4],
1582                [5, 6, 7, 8],
1583                [9, 10, 11, 12]
1584            ],
1585            [
1586                [13, 14, 15, 16],
1587                [17, 18, 19, 20],
1588                [21, 22, 23, 24]
1589            ];
1590            [
1591                domain_int_ground!(1..2),
1592                domain_int_ground!(1..3),
1593                domain_int_ground!(1..4)
1594            ]
1595        ];
1596        let shape = shape_of(&tensor).expect("shape_of to work on a 3D matrix");
1597
1598        assert_eq!(shape.size, 24);
1599        assert_eq!(shape.dims, vec![2, 3, 4]);
1600        assert_eq!(shape.strides, vec![12, 4, 1]);
1601        assert_eq!(
1602            shape.idx_doms,
1603            vec![
1604                domain_int_ground!(1..2),
1605                domain_int_ground!(1..3),
1606                domain_int_ground!(1..4)
1607            ]
1608        );
1609    }
1610
1611    #[test]
1612    fn matrix_partial_flatten() {
1613        let tensor: AbstractLiteral<Literal> = matrix![
1614            [
1615                // batch 1
1616                [1, 2, 3, 4],
1617                [5, 6, 7, 8],
1618                [9, 10, 11, 12]
1619            ],
1620            [
1621                // batch 2
1622                [13, 14, 15, 16],
1623                [17, 18, 19, 20],
1624                [21, 22, 23, 24]
1625            ]
1626        ];
1627        assert_eq!(partial_flatten(0, tensor.clone()), tensor);
1628
1629        let expected_flatten_1: AbstractLiteral<Literal> = matrix![
1630            [1, 2, 3, 4],
1631            [5, 6, 7, 8],
1632            [9, 10, 11, 12],
1633            [13, 14, 15, 16],
1634            [17, 18, 19, 20],
1635            [21, 22, 23, 24]
1636        ];
1637        assert_eq!(partial_flatten(1, tensor.clone()), expected_flatten_1);
1638
1639        let expected_flatten_2 =
1640            AbstractLiteral::matrix_implied_indices((1..25).map(Literal::from).collect());
1641        assert_eq!(partial_flatten(2, tensor), expected_flatten_2);
1642    }
1643}