1
//! Utility functions for working with matrices.
2
use std::collections::VecDeque;
3

            
4
use crate::ast::literals::AbstractLiteralValue;
5
use crate::ast::{
6
    AbstractLiteral, Atom, DomainOpError, DomainPtr, Expression as Expr, GroundDomain, Literal,
7
    Metadata, Moo, Range,
8
};
9
use crate::bug;
10
use crate::utils::MatrixShape;
11

            
12
use itertools::{Itertools, izip};
13
use uniplate::Biplate;
14

            
15
// ======================================================
16
// = "Shape" operations for matrices and matrix domains =
17
// ======================================================
18

            
19
/// Given an [AbstractLiteral::Matrix], get its [MatrixShape]
20
10017595
pub fn shape_of<T: MatrixValue>(matrix: &AbstractLiteral<T>) -> Option<MatrixShape<T::Dom>> {
21
    // put the dimensions in correct order
22
10017595
    let mut res = shape_of_inner(matrix)?;
23
10017595
    res.strides.reverse();
24
10017595
    res.dims.reverse();
25
10017595
    res.idx_doms.reverse();
26
10017595
    Some(res)
27
10017595
}
28

            
29
13460015
fn shape_of_inner<T: MatrixValue>(matrix: &AbstractLiteral<T>) -> Option<MatrixShape<T::Dom>> {
30
13460015
    let AbstractLiteral::Matrix(elems, dom) = matrix else {
31
        return None;
32
    };
33

            
34
13460015
    let sz = elems.len();
35
13460015
    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
13460015
    };
43

            
44
    // get child shapes and check that they are all the same
45
13460015
    let fst = elems[0].as_nested_matrix().and_then(shape_of_inner);
46
13460015
    for elem in elems.iter().skip(1) {
47
61718999
        debug_assert_eq!(
48
61718999
            fst,
49
61718999
            elem.as_nested_matrix().and_then(shape_of_inner),
50
            "Expected matrix elements to be consistent"
51
        );
52
    }
53

            
54
13460015
    Some(match fst {
55
        // if child shape is None, we have reached the last dimension
56
12952149
        None => MatrixShape {
57
12952149
            size: sz,
58
12952149
            strides: vec![1],
59
12952149
            dims: vec![sz],
60
12952149
            idx_doms: vec![dom.clone()],
61
12952149
        },
62
        // accumulate the next dimension
63
        Some(mut res) => {
64
507866
            res.strides.push(res.size);
65
507866
            res.dims.push(sz);
66
507866
            res.idx_doms.push(dom.clone());
67
507866
            res.size = if res.size == 0 { sz } else { sz * res.size };
68
507866
            res
69
        }
70
    })
71
13460015
}
72

            
73
/// If this is a matrix expression (as defined by [Expr::unwrap_matrix_unchecked]),
74
/// get its [MatrixShape]. See also: [shape_of].
75
288173
pub fn shape_of_matrix_expr(expr: &Expr) -> Option<MatrixShape<DomainPtr>> {
76
288173
    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
288173
        _ => None,
82
    }
83
288173
}
84

            
85
/// Same as [shape_of] but for a ground matrix domain
86
75314
pub fn shape_of_dom(
87
75314
    matrix_dom_gd: &GroundDomain,
88
75314
) -> Result<MatrixShape<Moo<GroundDomain>>, DomainOpError> {
89
75314
    let GroundDomain::Matrix(_, idx_doms, _) = matrix_dom_gd else {
90
        return Err(DomainOpError::WrongType);
91
    };
92

            
93
75314
    let len = idx_doms.len();
94
75314
    let mut strides = VecDeque::with_capacity(len);
95
75314
    let mut dimensions = VecDeque::with_capacity(len);
96

            
97
75314
    let mut size: usize = 1;
98
75314
    for gd in idx_doms.iter().rev() {
99
78506
        let gd_sz = gd.len_usize()?;
100
78506
        strides.push_front(size);
101
78506
        dimensions.push_front(gd_sz);
102
78506
        size = size.checked_mul(gd_sz).ok_or(DomainOpError::TooLarge)?;
103
    }
104

            
105
75314
    Ok(MatrixShape {
106
75314
        size,
107
75314
        dims: dimensions.into(),
108
75314
        strides: strides.into(),
109
75314
        idx_doms: idx_doms.clone(),
110
75314
    })
111
75314
}
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?)
120
3
pub fn partial_flatten<T: MatrixValue>(n: usize, matrix: AbstractLiteral<T>) -> AbstractLiteral<T> {
121
3
    if n == 0 {
122
1
        return matrix;
123
2
    }
124

            
125
    let shape = shape_of(&matrix).unwrap_or_else(|| bug!("Expected a matrix, got: {matrix}"));
126
    debug_assert!(
127
2
        n > 0 && n < shape.dims.len(),
128
        "Invalid number of dimensions to flatten"
129
    );
130
2
    let new_strides = Vec::from(&shape.strides[n..]);
131

            
132
2
    let flattened = flatten_owned(matrix).collect_vec();
133
2
    let res = unflatten_list(&flattened, &new_strides);
134

            
135
2
    res.into_nested_matrix()
136
        .unwrap_or_else(|e| bug!("Not a matrix: {e}"))
137
3
}
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.
146
78516
pub fn flatten<T: MatrixValue>(matrix: &AbstractLiteral<T>) -> impl Iterator<Item = &T> {
147
78516
    let AbstractLiteral::Matrix(elems, _) = matrix else {
148
        panic!("expected a matrix");
149
    };
150
78516
    flatten_inner(elems)
151
78516
}
152

            
153
#[inline]
154
78516
fn flatten_inner<'a, T: MatrixValue>(elems: &'a [T]) -> impl Iterator<Item = &'a T> {
155
223464
    elems.iter().flat_map(|elem| {
156
223464
        if let Some(m) = elem.as_nested_matrix() {
157
54270
            Box::new(flatten(m)) as Box<dyn Iterator<Item = &'a T>>
158
        } else {
159
169194
            Box::new(std::iter::once(elem)) as Box<dyn Iterator<Item = &'a T>>
160
        }
161
223464
    })
162
78516
}
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.
170
11880650
pub fn flatten_owned<T: MatrixValue>(matrix: AbstractLiteral<T>) -> impl Iterator<Item = T> {
171
11880650
    let AbstractLiteral::Matrix(elems, _) = matrix else {
172
        panic!("expected a matrix");
173
    };
174
11880650
    flatten_owned_inner(elems)
175
11880650
}
176

            
177
#[inline]
178
11880650
fn flatten_owned_inner<T: MatrixValue>(elems: Vec<T>) -> impl Iterator<Item = T> {
179
11880650
    elems
180
11880650
        .into_iter()
181
38110961
        .flat_map(|elem| match elem.into_nested_matrix() {
182
1865173
            Ok(m) => Box::new(flatten_owned(m)) as Box<dyn Iterator<Item = T>>,
183
36245788
            Err(leaf) => Box::new(std::iter::once(leaf)) as Box<dyn Iterator<Item = T>>,
184
38110961
        })
185
11880650
}
186

            
187
// ====================================
188
// = Matrix "unflattening" operations =
189
// ====================================
190

            
191
/// "Un-flatten" a slice of elements into a Matrix with the given index domains
192
13034843
pub fn unflatten_matrix<T: MatrixValue>(
193
13034843
    elems: &[T],
194
13034843
    index_domains: &[T::Dom],
195
13034843
    strides: &[usize],
196
13034843
) -> T {
197
13034843
    let dom = index_domains.first().expect("no index domains").clone();
198
13034843
    let stride = *strides.first().expect("no strides");
199

            
200
13034843
    if index_domains.len() == 1 {
201
12628943
        return T::from(AbstractLiteral::Matrix(Vec::from(elems), dom));
202
405900
    }
203

            
204
405900
    let mut inners = Vec::<T>::with_capacity(stride);
205
405900
    let mut i_start: usize = 0;
206
1387607
    while i_start < elems.len() {
207
981707
        let next = i_start + stride;
208
981707
        let elem = unflatten_matrix(&elems[i_start..next], &index_domains[1..], &strides[1..]);
209
981707
        inners.push(elem);
210
981707
        i_start = next;
211
    }
212
405900
    T::from(AbstractLiteral::Matrix(inners, dom))
213
13034843
}
214

            
215
/// Same transformation as [unflatten_matrix], but all index domains become `int(1..)`
216
8
pub fn unflatten_list<T: MatrixValue>(elems: &[T], strides: &[usize]) -> T {
217
8
    let stride = *strides.first().expect("no strides");
218
8
    if strides.len() == 1 {
219
7
        return AbstractLiteral::matrix_implied_indices(Vec::from(elems)).into();
220
1
    }
221

            
222
1
    let mut inners = Vec::<T>::with_capacity(stride);
223
1
    let mut i_start: usize = 0;
224
7
    while i_start < elems.len() {
225
6
        let next = i_start + stride;
226
6
        let elem = unflatten_list(&elems[i_start..next], &strides[1..]);
227
6
        inners.push(elem);
228
6
        i_start = next;
229
    }
230
1
    AbstractLiteral::matrix_implied_indices(inners).into()
231
8
}
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]
245
2117
pub fn index_domains<T: MatrixValue>(matrix: &AbstractLiteral<T>) -> Vec<T::Dom> {
246
2117
    shape_of(matrix)
247
        .unwrap_or_else(|| bug!("Expected matrix, got: {matrix}"))
248
        .idx_doms
249
2117
}
250

            
251
/// Gets the index domains for a matrix expression and resolves them
252
pub 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)], None))];
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
/// ```
289
10023421
pub fn try_enumerate_indices(
290
10023421
    index_domains: Vec<Moo<GroundDomain>>,
291
10023421
) -> Result<impl Iterator<Item = Vec<Literal>>, DomainOpError> {
292
10023421
    let domains = index_domains
293
10023421
        .into_iter()
294
10529451
        .map(|x| x.values().map(|values| values.collect_vec()))
295
        .collect::<Result<Vec<_>, _>>()?;
296
10023421
    Ok(domains.into_iter().multi_cartesian_product())
297
10023421
}
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]
303
10015742
pub fn enumerate_indices(
304
10015742
    index_domains: Vec<Moo<GroundDomain>>,
305
10015742
) -> impl Iterator<Item = Vec<Literal>> {
306
10015742
    try_enumerate_indices(index_domains).expect("index domain should be enumerable with .values()")
307
10015742
}
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.
312
9831
pub fn num_elements(index_domains: &[Moo<GroundDomain>]) -> Result<u64, DomainOpError> {
313
9831
    let idx_dom_lengths = index_domains
314
9831
        .iter()
315
18096
        .map(|d| d.length())
316
        .collect::<Result<Vec<_>, _>>()?;
317
9831
    Ok(idx_dom_lengths.iter().product())
318
9831
}
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
10015475
pub fn flatten_enumerate(
330
10015475
    matrix: AbstractLiteral<Literal>,
331
10015475
) -> impl Iterator<Item = (Vec<Literal>, Literal)> {
332
    let shape = shape_of(&matrix).unwrap_or_else(|| bug!("Expected matrix, got: {matrix}"));
333
10015475
    let index_domains: Vec<Moo<GroundDomain>> = shape
334
10015475
        .idx_doms
335
10015475
        .into_iter()
336
10015475
        .zip(shape.dims)
337
10520258
        .map(|(domain, len)| bound_index_domain_from_length(domain, len))
338
        .collect();
339
10015475
    izip!(enumerate_indices(index_domains), flatten_owned(matrix))
340
10015475
}
341

            
342
/// See [`enumerate_indices`]. This function zips the two given lists of index domains, performs a
343
/// union on each pair, and returns an enumerating iterator over the new list of domains.
344
6786
pub fn enumerate_index_union_indices(
345
6786
    a_domains: &[Moo<GroundDomain>],
346
6786
    b_domains: &[Moo<GroundDomain>],
347
6786
) -> Result<impl Iterator<Item = Vec<Literal>>, DomainOpError> {
348
6786
    if a_domains.len() != b_domains.len() {
349
        return Err(DomainOpError::WrongType);
350
6786
    }
351
6786
    let idx_domains: Result<Vec<_>, _> = a_domains
352
6786
        .iter()
353
6786
        .zip(b_domains.iter())
354
6844
        .map(|(a, b)| a.union(b))
355
        .collect();
356
6786
    let idx_domains = idx_domains?.into_iter().map(Moo::new).collect();
357

            
358
6786
    try_enumerate_indices(idx_domains)
359
6786
}
360

            
361
/// Given index domains for a multi-dimensional matrix and
362
/// the nth index in the flattened matrix, find the coordinates in the original matrix
363
271382
pub fn flat_index_to_full_index(index_domains: &[Moo<GroundDomain>], index: u64) -> Vec<Literal> {
364
271382
    let mut remaining = index;
365
271382
    let mut multipliers = vec![1; index_domains.len()];
366

            
367
271382
    for i in (1..index_domains.len()).rev() {
368
263552
        multipliers[i - 1] = multipliers[i] * index_domains[i].as_ref().length().unwrap();
369
    }
370

            
371
271382
    let mut coords = Vec::new();
372
271382
    for m in multipliers.iter() {
373
        // adjust for 1-based indexing
374
534934
        coords.push(((remaining / m + 1) as i32).into());
375
534934
        remaining %= *m;
376
    }
377

            
378
271382
    coords
379
271382
}
380

            
381
/// Gets concrete index domains for a matrix expression.
382
///
383
/// For matrix literals, right-unbounded integer index domains like `int(1..)` are bounded using
384
/// the literal's realised size in that dimension. For non-literals, this falls back to the
385
/// expression's resolved domain.
386
979193
pub fn bound_index_domains_of_expr(expr: &Expr) -> Option<Vec<Moo<GroundDomain>>> {
387
979193
    let dom = expr.domain_of().and_then(|dom| dom.resolve().ok())?;
388
978034
    let GroundDomain::Matrix(_, index_domains, _) = dom.as_ref() else {
389
689861
        return None;
390
    };
391

            
392
288173
    let Some(dimension_lengths) = expr_matrix_dimension_lengths(expr) else {
393
288173
        return Some(index_domains.clone());
394
    };
395

            
396
    assert_eq!(
397
        index_domains.len(),
398
        dimension_lengths.len(),
399
        "matrix literal domain rank should match its realised rank"
400
    );
401

            
402
    Some(
403
        index_domains
404
            .iter()
405
            .cloned()
406
            .zip(dimension_lengths)
407
            .map(|(domain, len)| bound_index_domain_from_length(domain, len))
408
            .collect(),
409
    )
410
979193
}
411

            
412
/// This is the same as `m[x]` except when `m` is of the forms:
413
///
414
/// - `n[..]`, then it produces n[x] instead of n[..][x]
415
/// - `flatten(n)`, then it produces `n[y]` instead of `flatten(n)[y]`,
416
///   where `y` is the full index corresponding to flat index `x`
417
///
418
/// # Returns
419
/// + `Some(expr)` if the safe indexing could be constructed
420
/// + `None` if it could not be constructed (e.g. invalid index type)
421
633563
pub fn safe_index_optimised(m: Expr, idx: Literal) -> Option<Expr> {
422
633563
    match m {
423
        Expr::SafeSlice(_, mat, idxs) => {
424
            // TODO: support >1 slice index (i.e. multidimensional slices)
425

            
426
113970
            let mut idxs = idxs;
427
168983
            let (slice_idx, _) = idxs.iter().find_position(|opt| opt.is_none())?;
428
113970
            let _ = idxs[slice_idx].replace(idx.into());
429

            
430
113970
            let Some(idxs) = idxs.into_iter().collect::<Option<Vec<_>>>() else {
431
                todo!("slice expression should not contain more than one unspecified index")
432
            };
433

            
434
113970
            Some(Expr::SafeIndex(Metadata::new(), mat, idxs))
435
        }
436
        Expr::Flatten(_, None, inner) => {
437
            // Similar to indexed_flatten_matrix rule, but we don't care about out of bounds here
438
7830
            let Literal::Int(index) = idx else {
439
                return None;
440
            };
441

            
442
7830
            let index_domains = bound_index_domains_of_expr(inner.as_ref())?;
443
7830
            if index_domains.iter().any(|domain| domain.length().is_err()) {
444
                return None;
445
7830
            }
446
7830
            let flat_index = flat_index_to_full_index(&index_domains, (index - 1) as u64);
447
7830
            let flat_index: Vec<Expr> = flat_index.into_iter().map(Into::into).collect();
448

            
449
7830
            Some(Expr::SafeIndex(Metadata::new(), inner, flat_index))
450
        }
451
511763
        _ => Some(Expr::SafeIndex(
452
511763
            Metadata::new(),
453
511763
            Moo::new(m),
454
511763
            vec![idx.into()],
455
511763
        )),
456
    }
457
633563
}
458

            
459
// ====================
460
// = Internal helpers =
461
// ====================
462

            
463
/// If this is a matrix expression, get sizes along its dimensions
464
#[inline]
465
288173
fn expr_matrix_dimension_lengths(expr: &Expr) -> Option<Vec<usize>> {
466
288173
    Some(shape_of_matrix_expr(expr)?.dims)
467
288173
}
468

            
469
/// Cap unbounded integer index domains to the given matrix dimension length.
470
#[inline]
471
10520258
fn bound_index_domain_from_length(mut domain: Moo<GroundDomain>, len: usize) -> Moo<GroundDomain> {
472
10520258
    match Moo::make_mut(&mut domain) {
473
10519475
        GroundDomain::Int(ranges, _) if ranges.len() == 1 && len > 0 => {
474
10407825
            ranges[0] = match &ranges[0] {
475
456
                Range::UnboundedR(start) => Range::Bounded(*start, start + (len as i32 - 1)),
476
                Range::Unbounded | Range::UnboundedL(_) => Range::Bounded(1, len as i32),
477
10387533
                Range::Bounded(low, high) => Range::Bounded(*low, *high),
478
19836
                Range::Single(value) => Range::Single(*value),
479
            };
480
10407825
            domain
481
        }
482
112433
        _ => domain,
483
    }
484
10520258
}
485

            
486
/// Things that can appear inside a matrix.
487
///
488
/// This is a helper trait to unify matrix operations on `Expression::AbstractLiteral`
489
/// and `AbstractLiteral<Literal>`
490
pub trait MatrixValue:
491
    AbstractLiteralValue + Sized + From<AbstractLiteral<Self>> + Biplate<AbstractLiteral<Self>>
492
{
493
    /// If this element is a nested matrix, return a reference to it
494
    fn as_nested_matrix(&self) -> Option<&AbstractLiteral<Self>>;
495
    /// If this element is a nested matrix, consume it and return the matrix
496
    fn into_nested_matrix(self) -> Result<AbstractLiteral<Self>, Self>;
497
}
498

            
499
impl MatrixValue for Literal {
500
    #[inline]
501
75402478
    fn as_nested_matrix(&self) -> Option<&AbstractLiteral<Literal>> {
502
75402478
        match self {
503
3496690
            Literal::AbstractLiteral(m @ AbstractLiteral::Matrix(..)) => Some(m),
504
71905788
            _ => None,
505
        }
506
75402478
    }
507

            
508
    #[inline]
509
38110963
    fn into_nested_matrix(self) -> Result<AbstractLiteral<Literal>, Self> {
510
38110963
        match self {
511
1865175
            Literal::AbstractLiteral(m @ AbstractLiteral::Matrix(..)) => Ok(m),
512
36245788
            other => Err(other),
513
        }
514
38110963
    }
515
}
516

            
517
impl MatrixValue for Expr {
518
    #[inline]
519
    fn as_nested_matrix(&self) -> Option<&AbstractLiteral<Expr>> {
520
        match self {
521
            Expr::AbstractLiteral(_, m @ AbstractLiteral::Matrix(..)) => Some(m),
522
            _ => None,
523
        }
524
    }
525

            
526
    #[inline]
527
    fn into_nested_matrix(self) -> Result<AbstractLiteral<Expr>, Self> {
528
        match self {
529
            Expr::AbstractLiteral(_, m @ AbstractLiteral::Matrix(..)) => Ok(m),
530
            other => Err(other),
531
        }
532
    }
533
}