conjure_cp_core/ast/
matrix.rs

1//! Utility functions for working with matrices.
2use std::collections::VecDeque;
3
4use crate::ast::literals::AbstractLiteralValue;
5use crate::ast::{
6    AbstractLiteral, Atom, DomainOpError, DomainPtr, Expression as Expr, GroundDomain, Literal,
7    Metadata, Moo, Range,
8};
9use crate::bug;
10use crate::utils::MatrixShape;
11
12use itertools::{Itertools, izip};
13use uniplate::Biplate;
14
15// ======================================================
16// = "Shape" operations for matrices and matrix domains =
17// ======================================================
18
19/// Given an [AbstractLiteral::Matrix], get its [MatrixShape]
20pub fn shape_of<T: MatrixValue>(matrix: &AbstractLiteral<T>) -> Option<MatrixShape<T::Dom>> {
21    // put the dimensions in correct order
22    let mut res = shape_of_inner(matrix)?;
23    res.strides.reverse();
24    res.dims.reverse();
25    res.idx_doms.reverse();
26    Some(res)
27}
28
29fn shape_of_inner<T: MatrixValue>(matrix: &AbstractLiteral<T>) -> Option<MatrixShape<T::Dom>> {
30    let AbstractLiteral::Matrix(elems, dom) = matrix else {
31        return None;
32    };
33
34    let sz = elems.len();
35    if sz == 0 {
36        return Some(MatrixShape {
37            size: 0,
38            strides: vec![0],
39            dims: vec![0],
40            idx_doms: vec![dom.clone()],
41        });
42    };
43
44    // get child shapes and check that they are all the same
45    let fst = elems[0].as_nested_matrix().and_then(shape_of_inner);
46    for elem in elems.iter().skip(1) {
47        debug_assert_eq!(
48            fst,
49            elem.as_nested_matrix().and_then(shape_of_inner),
50            "Expected matrix elements to be consistent"
51        );
52    }
53
54    Some(match fst {
55        // if child shape is None, we have reached the last dimension
56        None => MatrixShape {
57            size: sz,
58            strides: vec![1],
59            dims: vec![sz],
60            idx_doms: vec![dom.clone()],
61        },
62        // accumulate the next dimension
63        Some(mut res) => {
64            res.strides.push(res.size);
65            res.dims.push(sz);
66            res.idx_doms.push(dom.clone());
67            res.size = if res.size == 0 { sz } else { sz * res.size };
68            res
69        }
70    })
71}
72
73/// If this is a matrix expression (as defined by [Expr::unwrap_matrix_unchecked]),
74/// get its [MatrixShape]. See also: [shape_of].
75pub fn shape_of_matrix_expr(expr: &Expr) -> Option<MatrixShape<DomainPtr>> {
76    match expr {
77        Expr::Atomic(_, Atom::Literal(Literal::AbstractLiteral(lit))) => {
78            Some(shape_of(lit)?.into())
79        }
80        Expr::AbstractLiteral(_, lit) => shape_of(lit),
81        _ => None,
82    }
83}
84
85/// Same as [shape_of] but for a ground matrix domain
86pub fn shape_of_dom(
87    matrix_dom_gd: &GroundDomain,
88) -> Result<MatrixShape<Moo<GroundDomain>>, DomainOpError> {
89    let GroundDomain::Matrix(_, idx_doms) = matrix_dom_gd else {
90        return Err(DomainOpError::WrongType);
91    };
92
93    let len = idx_doms.len();
94    let mut strides = VecDeque::with_capacity(len);
95    let mut dimensions = VecDeque::with_capacity(len);
96
97    let mut size: usize = 1;
98    for gd in idx_doms.iter().rev() {
99        let gd_sz = gd.len_usize()?;
100        strides.push_front(size);
101        dimensions.push_front(gd_sz);
102        size = size.checked_mul(gd_sz).ok_or(DomainOpError::TooLarge)?;
103    }
104
105    Ok(MatrixShape {
106        size,
107        dims: dimensions.into(),
108        strides: strides.into(),
109        idx_doms: idx_doms.clone(),
110    })
111}
112
113// ================================
114// = Matrix flattening operations =
115// ================================
116
117/// Given a nested matrix, flatten its first `n+1` dimensions into one.
118/// The resulting matrix will be a list because otherwise index domains would get weird, fast...
119/// (unless we want to only support integer indices?)
120pub fn partial_flatten<T: MatrixValue>(n: usize, matrix: AbstractLiteral<T>) -> AbstractLiteral<T> {
121    if n == 0 {
122        return matrix;
123    }
124
125    let shape = shape_of(&matrix).unwrap_or_else(|| bug!("Expected a matrix, got: {matrix}"));
126    debug_assert!(
127        n > 0 && n < shape.dims.len(),
128        "Invalid number of dimensions to flatten"
129    );
130    let new_strides = Vec::from(&shape.strides[n..]);
131
132    let flattened = flatten_owned(matrix).collect_vec();
133    let res = unflatten_list(&flattened, &new_strides);
134
135    res.into_nested_matrix()
136        .unwrap_or_else(|e| bug!("Not a matrix: {e}"))
137}
138
139/// Flattens a multi-dimensional matrix into a one-dimensional slice of its elements.
140/// The elements are returned in row-major ordering (see [`enumerate_indices`]).
141/// The elements are borrowed; To consume the matrix and return owned values, see [flatten_owned].
142///
143/// # Panics
144/// + If the number or type of elements in each dimension is inconsistent.
145/// + If `matrix` is not a matrix.
146pub fn flatten<T: MatrixValue>(matrix: &AbstractLiteral<T>) -> impl Iterator<Item = &T> {
147    let AbstractLiteral::Matrix(elems, _) = matrix else {
148        panic!("expected a matrix");
149    };
150    flatten_inner(elems)
151}
152
153#[inline]
154fn flatten_inner<'a, T: MatrixValue>(elems: &'a [T]) -> impl Iterator<Item = &'a T> {
155    elems.iter().flat_map(|elem| {
156        if let Some(m) = elem.as_nested_matrix() {
157            Box::new(flatten(m)) as Box<dyn Iterator<Item = &'a T>>
158        } else {
159            Box::new(std::iter::once(elem)) as Box<dyn Iterator<Item = &'a T>>
160        }
161    })
162}
163
164/// Consumes a multi-dimensional matrix and returns a one-dimensional slice of its elements.
165/// The elements are returned in row-major ordering (see [`enumerate_indices`]).
166///
167/// # Panics
168/// + If the number or type of elements in each dimension is inconsistent.
169/// + If `matrix` is not a matrix.
170pub fn flatten_owned<T: MatrixValue>(matrix: AbstractLiteral<T>) -> impl Iterator<Item = T> {
171    let AbstractLiteral::Matrix(elems, _) = matrix else {
172        panic!("expected a matrix");
173    };
174    flatten_owned_inner(elems)
175}
176
177#[inline]
178fn flatten_owned_inner<T: MatrixValue>(elems: Vec<T>) -> impl Iterator<Item = T> {
179    elems
180        .into_iter()
181        .flat_map(|elem| match elem.into_nested_matrix() {
182            Ok(m) => Box::new(flatten_owned(m)) as Box<dyn Iterator<Item = T>>,
183            Err(leaf) => Box::new(std::iter::once(leaf)) as Box<dyn Iterator<Item = T>>,
184        })
185}
186
187// ====================================
188// = Matrix "unflattening" operations =
189// ====================================
190
191/// "Un-flatten" a slice of elements into a Matrix with the given index domains
192pub fn unflatten_matrix<T: MatrixValue>(
193    elems: &[T],
194    index_domains: &[T::Dom],
195    strides: &[usize],
196) -> T {
197    let dom = index_domains.first().expect("no index domains").clone();
198    let stride = *strides.first().expect("no strides");
199
200    if index_domains.len() == 1 {
201        return T::from(AbstractLiteral::Matrix(Vec::from(elems), dom));
202    }
203
204    let mut inners = Vec::<T>::with_capacity(stride);
205    let mut i_start: usize = 0;
206    while i_start < elems.len() {
207        let next = i_start + stride;
208        let elem = unflatten_matrix(&elems[i_start..next], &index_domains[1..], &strides[1..]);
209        inners.push(elem);
210        i_start = next;
211    }
212    T::from(AbstractLiteral::Matrix(inners, dom))
213}
214
215/// Same transformation as [unflatten_matrix], but all index domains become `int(1..)`
216pub fn unflatten_list<T: MatrixValue>(elems: &[T], strides: &[usize]) -> T {
217    let stride = *strides.first().expect("no strides");
218    if strides.len() == 1 {
219        return AbstractLiteral::matrix_implied_indices(Vec::from(elems)).into();
220    }
221
222    let mut inners = Vec::<T>::with_capacity(stride);
223    let mut i_start: usize = 0;
224    while i_start < elems.len() {
225        let next = i_start + stride;
226        let elem = unflatten_list(&elems[i_start..next], &strides[1..]);
227        inners.push(elem);
228        i_start = next;
229    }
230    AbstractLiteral::matrix_implied_indices(inners).into()
231}
232
233// =============================
234// = Matrix indexing utilities =
235// =============================
236
237/// Gets the index domains for a matrix literal.
238///
239/// # Panics
240///
241/// + If `matrix` is not a matrix.
242///
243/// + If the number or type of elements in each dimension is inconsistent.
244#[inline]
245pub fn index_domains<T: MatrixValue>(matrix: &AbstractLiteral<T>) -> Vec<T::Dom> {
246    shape_of(matrix)
247        .unwrap_or_else(|| bug!("Expected matrix, got: {matrix}"))
248        .idx_doms
249}
250
251/// Gets the index domains for a matrix expression and resolves them
252pub fn resolved_index_domains(
253    matrix: &AbstractLiteral<Expr>,
254) -> Result<Vec<Moo<GroundDomain>>, DomainOpError> {
255    index_domains(matrix)
256        .into_iter()
257        .map(|d| d.resolve())
258        .try_collect()
259}
260
261/// For some index domains, returns a list containing each of the possible indices.
262///
263/// Indices are traversed in row-major ordering.
264///
265/// This is an O(n^dim) operation, where dim is the number of dimensions in the matrix.
266///
267/// # Panics
268///
269/// + If any of the index domains are not finite or enumerable with [`Domain::values`].
270///
271/// # Example
272///
273/// ```
274/// use std::collections::HashSet;
275/// use conjure_cp_core::ast::{GroundDomain,Moo,Range,Literal,matrix};
276/// let index_domains = vec![Moo::new(GroundDomain::Bool),Moo::new(GroundDomain::Int(vec![Range::Bounded(1,2)]))];
277///
278/// let expected_indices = HashSet::from([
279///   vec![Literal::Bool(false),Literal::Int(1)],
280///   vec![Literal::Bool(false),Literal::Int(2)],
281///   vec![Literal::Bool(true),Literal::Int(1)],
282///   vec![Literal::Bool(true),Literal::Int(2)]
283///   ]);
284///
285/// let actual_indices: HashSet<_> = matrix::enumerate_indices(index_domains).collect();
286///
287/// assert_eq!(actual_indices, expected_indices);
288/// ```
289pub fn try_enumerate_indices(
290    index_domains: Vec<Moo<GroundDomain>>,
291) -> Result<impl Iterator<Item = Vec<Literal>>, DomainOpError> {
292    let domains = index_domains
293        .into_iter()
294        .map(|x| x.values().map(|values| values.collect_vec()))
295        .collect::<Result<Vec<_>, _>>()?;
296    Ok(domains.into_iter().multi_cartesian_product())
297}
298
299/// For some index domains, returns a list containing each of the possible indices.
300///
301/// See [`try_enumerate_indices`] for the fallible variant.
302#[inline]
303pub fn enumerate_indices(
304    index_domains: Vec<Moo<GroundDomain>>,
305) -> impl Iterator<Item = Vec<Literal>> {
306    try_enumerate_indices(index_domains).expect("index domain should be enumerable with .values()")
307}
308
309/// Returns the number of possible elements indexable by the given index domains.
310///
311/// In short, returns the product of the sizes of the given indices.
312pub fn num_elements(index_domains: &[Moo<GroundDomain>]) -> Result<u64, DomainOpError> {
313    let idx_dom_lengths = index_domains
314        .iter()
315        .map(|d| d.length())
316        .collect::<Result<Vec<_>, _>>()?;
317    Ok(idx_dom_lengths.iter().product())
318}
319
320/// Flattens a multi-dimensional matrix literal into an iterator over (indices,element).
321///
322/// # Panics
323///
324///   + If the number or type of elements in each dimension is inconsistent.
325///
326///   + If `matrix` is not a matrix.
327///
328///   + If any dimensions in the matrix are not finite or enumerable with [`Domain::values`].
329///     However, index domains in the form `int(i..)` are supported.
330pub fn flatten_enumerate(
331    matrix: AbstractLiteral<Literal>,
332) -> impl Iterator<Item = (Vec<Literal>, Literal)> {
333    let index_domains = index_domains(&matrix);
334    izip!(enumerate_indices(index_domains), flatten_owned(matrix))
335}
336
337/// See [`enumerate_indices`]. This function zips the two given lists of index domains, performs a
338/// union on each pair, and returns an enumerating iterator over the new list of domains.
339pub fn enumerate_index_union_indices(
340    a_domains: &[Moo<GroundDomain>],
341    b_domains: &[Moo<GroundDomain>],
342) -> Result<impl Iterator<Item = Vec<Literal>>, DomainOpError> {
343    if a_domains.len() != b_domains.len() {
344        return Err(DomainOpError::WrongType);
345    }
346    let idx_domains: Result<Vec<_>, _> = a_domains
347        .iter()
348        .zip(b_domains.iter())
349        .map(|(a, b)| a.union(b))
350        .collect();
351    let idx_domains = idx_domains?.into_iter().map(Moo::new).collect();
352
353    try_enumerate_indices(idx_domains)
354}
355
356/// Given index domains for a multi-dimensional matrix and
357/// the nth index in the flattened matrix, find the coordinates in the original matrix
358pub fn flat_index_to_full_index(index_domains: &[Moo<GroundDomain>], index: u64) -> Vec<Literal> {
359    let mut remaining = index;
360    let mut multipliers = vec![1; index_domains.len()];
361
362    for i in (1..index_domains.len()).rev() {
363        multipliers[i - 1] = multipliers[i] * index_domains[i].as_ref().length().unwrap();
364    }
365
366    let mut coords = Vec::new();
367    for m in multipliers.iter() {
368        // adjust for 1-based indexing
369        coords.push(((remaining / m + 1) as i32).into());
370        remaining %= *m;
371    }
372
373    coords
374}
375
376/// Gets concrete index domains for a matrix expression.
377///
378/// For matrix literals, right-unbounded integer index domains like `int(1..)` are bounded using
379/// the literal's realised size in that dimension. For non-literals, this falls back to the
380/// expression's resolved domain.
381pub fn bound_index_domains_of_expr(expr: &Expr) -> Option<Vec<Moo<GroundDomain>>> {
382    let dom = expr.domain_of().and_then(|dom| dom.resolve().ok())?;
383    let GroundDomain::Matrix(_, index_domains) = dom.as_ref() else {
384        return None;
385    };
386
387    let Some(dimension_lengths) = expr_matrix_dimension_lengths(expr) else {
388        return Some(index_domains.clone());
389    };
390
391    assert_eq!(
392        index_domains.len(),
393        dimension_lengths.len(),
394        "matrix literal domain rank should match its realised rank"
395    );
396
397    Some(
398        index_domains
399            .iter()
400            .cloned()
401            .zip(dimension_lengths)
402            .map(|(domain, len)| bound_index_domain_from_length(domain, len))
403            .collect(),
404    )
405}
406
407/// This is the same as `m[x]` except when `m` is of the forms:
408///
409/// - `n[..]`, then it produces n[x] instead of n[..][x]
410/// - `flatten(n)`, then it produces `n[y]` instead of `flatten(n)[y]`,
411///   where `y` is the full index corresponding to flat index `x`
412///
413/// # Returns
414/// + `Some(expr)` if the safe indexing could be constructed
415/// + `None` if it could not be constructed (e.g. invalid index type)
416pub fn safe_index_optimised(m: Expr, idx: Literal) -> Option<Expr> {
417    match m {
418        Expr::SafeSlice(_, mat, idxs) => {
419            // TODO: support >1 slice index (i.e. multidimensional slices)
420
421            let mut idxs = idxs;
422            let (slice_idx, _) = idxs.iter().find_position(|opt| opt.is_none())?;
423            let _ = idxs[slice_idx].replace(idx.into());
424
425            let Some(idxs) = idxs.into_iter().collect::<Option<Vec<_>>>() else {
426                todo!("slice expression should not contain more than one unspecified index")
427            };
428
429            Some(Expr::SafeIndex(Metadata::new(), mat, idxs))
430        }
431        Expr::Flatten(_, None, inner) => {
432            // Similar to indexed_flatten_matrix rule, but we don't care about out of bounds here
433            let Literal::Int(index) = idx else {
434                return None;
435            };
436
437            let index_domains = bound_index_domains_of_expr(inner.as_ref())?;
438            if index_domains.iter().any(|domain| domain.length().is_err()) {
439                return None;
440            }
441            let flat_index = flat_index_to_full_index(&index_domains, (index - 1) as u64);
442            let flat_index: Vec<Expr> = flat_index.into_iter().map(Into::into).collect();
443
444            Some(Expr::SafeIndex(Metadata::new(), inner, flat_index))
445        }
446        _ => Some(Expr::SafeIndex(
447            Metadata::new(),
448            Moo::new(m),
449            vec![idx.into()],
450        )),
451    }
452}
453
454// ====================
455// = Internal helpers =
456// ====================
457
458/// If this is a matrix expression, get sizes along its dimensions
459#[inline]
460fn expr_matrix_dimension_lengths(expr: &Expr) -> Option<Vec<usize>> {
461    Some(shape_of_matrix_expr(expr)?.dims)
462}
463
464/// Cap all `N..` ranges in an int domain to the given length
465#[inline]
466fn bound_index_domain_from_length(mut domain: Moo<GroundDomain>, len: usize) -> Moo<GroundDomain> {
467    match Moo::make_mut(&mut domain) {
468        GroundDomain::Int(ranges) if ranges.len() == 1 && len > 0 => {
469            if let Range::UnboundedR(start) = ranges[0] {
470                let end = start + (len as i32 - 1);
471                ranges[0] = Range::Bounded(start, end);
472            }
473            domain
474        }
475        _ => domain,
476    }
477}
478
479/// Things that can appear inside a matrix.
480///
481/// This is a helper trait to unify matrix operations on `Expression::AbstractLiteral`
482/// and `AbstractLiteral<Literal>`
483pub trait MatrixValue:
484    AbstractLiteralValue + Sized + From<AbstractLiteral<Self>> + Biplate<AbstractLiteral<Self>>
485{
486    /// If this element is a nested matrix, return a reference to it
487    fn as_nested_matrix(&self) -> Option<&AbstractLiteral<Self>>;
488    /// If this element is a nested matrix, consume it and return the matrix
489    fn into_nested_matrix(self) -> Result<AbstractLiteral<Self>, Self>;
490}
491
492impl MatrixValue for Literal {
493    #[inline]
494    fn as_nested_matrix(&self) -> Option<&AbstractLiteral<Literal>> {
495        match self {
496            Literal::AbstractLiteral(m @ AbstractLiteral::Matrix(..)) => Some(m),
497            _ => None,
498        }
499    }
500
501    #[inline]
502    fn into_nested_matrix(self) -> Result<AbstractLiteral<Literal>, Self> {
503        match self {
504            Literal::AbstractLiteral(m @ AbstractLiteral::Matrix(..)) => Ok(m),
505            other => Err(other),
506        }
507    }
508}
509
510impl MatrixValue for Expr {
511    #[inline]
512    fn as_nested_matrix(&self) -> Option<&AbstractLiteral<Expr>> {
513        match self {
514            Expr::AbstractLiteral(_, m @ AbstractLiteral::Matrix(..)) => Some(m),
515            _ => None,
516        }
517    }
518
519    #[inline]
520    fn into_nested_matrix(self) -> Result<AbstractLiteral<Expr>, Self> {
521        match self {
522            Expr::AbstractLiteral(_, m @ AbstractLiteral::Matrix(..)) => Ok(m),
523            other => Err(other),
524        }
525    }
526}
conjure_cp_core/ast/matrix.rs

conjure_cp_core/ast/
matrix.rs
