Grcov report - matrix.rs

1

//! Utility functions for working with matrices.

2

3

// TODO: Georgiis essence macro would look really nice in these examples!

4

5

use std::collections::VecDeque;

6

7

use itertools::{Itertools, izip};

8

use uniplate::Uniplate as _;

9

10

use crate::ast::{DomainOpError, Expression as Expr, GroundDomain, Metadata, Moo, Range};

11

12

use super::{AbstractLiteral, Literal};

13

14

/// For some index domains, returns a list containing each of the possible indices.

15

///

16

/// Indices are traversed in row-major ordering.

17

///

18

/// This is an O(n^dim) operation, where dim is the number of dimensions in the matrix.

19

///

20

/// # Panics

21

///

22

/// + If any of the index domains are not finite or enumerable with [`Domain::values`].

23

///

24

/// # Example

25

///

26

/// ```

27

/// use std::collections::HashSet;

28

/// use conjure_cp_core::ast::{GroundDomain,Moo,Range,Literal,matrix};

29

/// let index_domains = vec![Moo::new(GroundDomain::Bool),Moo::new(GroundDomain::Int(vec![Range::Bounded(1,2)]))];

30

///

31

/// let expected_indices = HashSet::from([

32

///   vec![Literal::Bool(false),Literal::Int(1)],

33

///   vec![Literal::Bool(false),Literal::Int(2)],

34

///   vec![Literal::Bool(true),Literal::Int(1)],

35

///   vec![Literal::Bool(true),Literal::Int(2)]

36

///   ]);

37

///

38

/// let actual_indices: HashSet<_> = matrix::enumerate_indices(index_domains).collect();

39

///

40

/// assert_eq!(actual_indices, expected_indices);

41

/// ```

42

1680

pub fn enumerate_indices(

43

1680

    index_domains: Vec<Moo<GroundDomain>>,

44

1680

) -> impl Iterator<Item = Vec<Literal>> {

45

1680

    index_domains

46

1680

        .into_iter()

47

2000

        .map(|x| {

48

2000

            x.values()

49

2000

                .expect("index domain should be enumerable with .values()")

50

2000

                .collect_vec()

51

2000

})

52

1680

        .multi_cartesian_product()

53

1680

54

55

/// Returns the number of possible elements indexable by the given index domains.

56

///

57

/// In short, returns the product of the sizes of the given indices.

58

200

pub fn num_elements(index_domains: &[Moo<GroundDomain>]) -> Result<u64, DomainOpError> {

59

200

    let idx_dom_lengths = index_domains

60

200

        .iter()

61

400

        .map(|d| d.length())

62

200

        .collect::<Result<Vec<_>, _>>()?;

63

200

    Ok(idx_dom_lengths.iter().product())

64

200

65

66

/// Flattens a multi-dimensional matrix literal into a one-dimensional slice of its elements.

67

///

68

/// The elements of the matrix are returned in row-major ordering (see [`enumerate_indices`]).

69

///

70

/// # Panics

71

///

72

/// + If the number or type of elements in each dimension is inconsistent.

73

///

74

/// + If `matrix` is not a matrix.

75

140

pub fn flatten(matrix: AbstractLiteral<Literal>) -> impl Iterator<Item = Literal> {

76

140

    let AbstractLiteral::Matrix(elems, _) = matrix else {

77

        panic!("matrix should be a matrix");

78

};

79

80

140

    flatten_1(elems)

81

140

82

83

400

fn flatten_1(elems: Vec<Literal>) -> impl Iterator<Item = Literal> {

84

1160

    elems.into_iter().flat_map(|elem| {

85

140

        if let Literal::AbstractLiteral(m @ AbstractLiteral::Matrix(_, _)) = elem {

86

140

            Box::new(flatten(m)) as Box<dyn Iterator<Item = Literal>>

87

        } else {

88

1020

            Box::new(std::iter::once(elem)) as Box<dyn Iterator<Item = Literal>>

89

90

1160

})

91

400

92

/// Flattens a multi-dimensional matrix literal into an iterator over (indices,element).

93

///

94

/// # Panics

95

///

96

///   + If the number or type of elements in each dimension is inconsistent.

97

///

98

///   + If `matrix` is not a matrix.

99

///

100

///   + If any dimensions in the matrix are not finite or enumerable with [`Domain::values`].

101

///     However, index domains in the form `int(i..)` are supported.

102

260

pub fn flatten_enumerate(

103

260

    matrix: AbstractLiteral<Literal>,

104

260

) -> impl Iterator<Item = (Vec<Literal>, Literal)> {

105

260

    let AbstractLiteral::Matrix(elems, _) = matrix.clone() else {

106

        panic!("matrix should be a matrix");

107

};

108

109

260

    let index_domains = index_domains(matrix)

110

260

        .into_iter()

111

320

        .map(|mut x| match Moo::make_mut(&mut x) {

112

            // give unboundedr index domains an end

113

320

            GroundDomain::Int(ranges) if ranges.len() == 1 && !elems.is_empty() => {

114

280

                if let Range::UnboundedR(start) = ranges[0] {

115

120

                    ranges[0] = Range::Bounded(start, start + (elems.len() as i32 - 1));

116

160

};

117

280

118

119

40

            _ => x,

120

320

})

121

260

        .collect_vec();

122

123

260

    izip!(enumerate_indices(index_domains), flatten_1(elems))

124

260

125

126

/// Gets the index domains for a matrix literal.

127

///

128

/// # Panics

129

///

130

/// + If `matrix` is not a matrix.

131

///

132

/// + If the number or type of elements in each dimension is inconsistent.

133

280

pub fn index_domains(matrix: AbstractLiteral<Literal>) -> Vec<Moo<GroundDomain>> {

134

280

    let AbstractLiteral::Matrix(_, _) = matrix else {

135

        panic!("matrix should be a matrix");

136

};

137

138

280

    matrix.cata(&move |element: AbstractLiteral<Literal>,

139

520

                       child_index_domains: VecDeque<Vec<Moo<GroundDomain>>>| {

140

520

        assert!(

141

520

            child_index_domains.iter().all_equal(),

142

            "each child of a matrix should have the same index domain"

143

);

144

145

520

        let child_index_domains = child_index_domains

146

520

            .front()

147

520

            .unwrap_or(&vec![])

148

520

            .iter()

149

520

            .cloned()

150

520

            .collect_vec();

151

520

        match element {

152

            AbstractLiteral::Set(_) => vec![],

153

            AbstractLiteral::MSet(_) => vec![],

154

520

            AbstractLiteral::Matrix(_, domain) => {

155

520

                let mut index_domains = vec![domain];

156

520

                index_domains.extend(child_index_domains);

157

520

                index_domains

158

159

            AbstractLiteral::Tuple(_) => vec![],

160

            AbstractLiteral::Record(_) => vec![],

161

            AbstractLiteral::Function(_) => vec![],

162

163

520

})

164

280

165

166

/// See [`enumerate_indices`]. This function zips the two given lists of index domains, performs a

167

/// union on each pair, and returns an enumerating iterator over the new list of domains.

168

100

pub fn enumerate_index_union_indices(

169

100

    a_domains: &[Moo<GroundDomain>],

170

100

    b_domains: &[Moo<GroundDomain>],

171

100

) -> Result<impl Iterator<Item = Vec<Literal>>, DomainOpError> {

172

100

    if a_domains.len() != b_domains.len() {

173

        return Err(DomainOpError::WrongType);

174

100

175

100

    let idx_domains: Result<Vec<_>, _> = a_domains

176

100

        .iter()

177

100

        .zip(b_domains.iter())

178

120

        .map(|(a, b)| a.union(b))

179

100

        .collect();

180

100

    let idx_domains = idx_domains?.into_iter().map(Moo::new).collect();

181

182

100

    Ok(enumerate_indices(idx_domains))

183

100

184

185

// Given index domains for a multi-dimensional matrix and the nth index in the flattened matrix, find the coordinates in the original matrix

186

160

pub fn flat_index_to_full_index(index_domains: &[Moo<GroundDomain>], index: u64) -> Vec<Literal> {

187

160

    let mut remaining = index;

188

160

    let mut multipliers = vec![1; index_domains.len()];

189

190

160

    for i in (1..index_domains.len()).rev() {

191

160

        multipliers[i - 1] = multipliers[i] * index_domains[i].as_ref().length().unwrap();

192

160

193

194

160

    let mut coords = Vec::new();

195

320

    for m in multipliers.iter() {

196

        // adjust for 1-based indexing

197

320

        coords.push(((remaining / m + 1) as i32).into());

198

320

        remaining %= *m;

199

320

200

201

160

    coords

202

160

203

204

/// This is the same as `m[x]` except when `m` is of the forms:

205

///

206

/// - `n[..]`, then it produces n[x] instead of n[..][x]

207

/// - `flatten(n)`, then it produces `n[y]` instead of `flatten(n)[y]`,

208

///   where `y` is the full index corresponding to flat index `x`

209

///

210

/// # Returns

211

/// + `Some(expr)` if the safe indexing could be constructed

212

/// + `None` if it could not be constructed (e.g. invalid index type)

213

960

pub fn safe_index_optimised(m: Expr, idx: Literal) -> Option<Expr> {

214

    match m {

215

480

        Expr::SafeSlice(_, mat, idxs) => {

216

            // TODO: support >1 slice index (i.e. multidimensional slices)

217

218

480

            let mut idxs = idxs;

219

720

            let (slice_idx, _) = idxs.iter().find_position(|opt| opt.is_none())?;

220

480

            let _ = idxs[slice_idx].replace(idx.into());

221

222

480

            let Some(idxs) = idxs.into_iter().collect::<Option<Vec<_>>>() else {

223

                todo!("slice expression should not contain more than one unspecified index")

224

};

225

226

480

            Some(Expr::SafeIndex(Metadata::new(), mat, idxs))

227

228

        Expr::Flatten(_, None, inner) => {

229

            // Similar to indexed_flatten_matrix rule, but we don't care about out of bounds here

230

            let Literal::Int(index) = idx else {

231

                return None;

232

};

233

234

            let dom = inner.domain_of().and_then(|dom| dom.resolve())?;

235

            let GroundDomain::Matrix(_, index_domains) = dom.as_ref() else {

236

                return None;

237

};

238

            let flat_index = flat_index_to_full_index(index_domains, (index - 1) as u64);

239

            let flat_index: Vec<Expr> = flat_index.into_iter().map(Into::into).collect();

240

241

            Some(Expr::SafeIndex(Metadata::new(), inner, flat_index))

242

243

480

        _ => Some(Expr::SafeIndex(

244

480

            Metadata::new(),

245

480

            Moo::new(m),

246

480

            vec![idx.into()],

247

480

)),

248

249

960