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

34802

pub fn try_enumerate_indices(

43

34802

    index_domains: Vec<Moo<GroundDomain>>,

44

34802

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

45

34802

    let domains = index_domains

46

34802

        .into_iter()

47

49448

        .map(|x| x.values().map(|values| values.collect_vec()))

48

34802

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

49

34802

    Ok(domains.into_iter().multi_cartesian_product())

50

34802

51

52

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

53

///

54

/// See [`try_enumerate_indices`] for the fallible variant.

55

34282

pub fn enumerate_indices(

56

34282

    index_domains: Vec<Moo<GroundDomain>>,

57

34282

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

58

34282

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

59

34282

60

61

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

62

///

63

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

64

800

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

65

800

    let idx_dom_lengths = index_domains

66

800

        .iter()

67

1600

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

68

800

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

69

800

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

70

800

71

72

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

73

///

74

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

75

///

76

/// # Panics

77

///

78

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

79

///

80

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

81

11524

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

82

11524

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

83

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

84

};

85

86

11524

    flatten_1(elems)

87

11524

88

89

28180

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

90

56218

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

91

11524

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

92

11524

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

93

        } else {

94

44694

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

95

96

56218

})

97

28180

98

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

99

///

100

/// # Panics

101

///

102

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

103

///

104

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

105

///

106

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

107

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

108

16656

pub fn flatten_enumerate(

109

16656

    matrix: AbstractLiteral<Literal>,

110

16656

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

111

16656

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

112

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

113

};

114

115

16656

    let index_domains = index_domains(matrix);

116

117

16656

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

118

16656

119

120

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

121

///

122

/// # Panics

123

///

124

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

125

///

126

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

127

16736

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

128

16736

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

129

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

130

};

131

132

16736

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

133

34580

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

134

34580

        assert!(

135

34580

            child_index_domains.iter().all_equal(),

136

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

137

);

138

139

34580

        let child_index_domains = child_index_domains

140

34580

            .front()

141

34580

            .unwrap_or(&vec![])

142

34580

            .iter()

143

34580

            .cloned()

144

34580

            .collect_vec();

145

34580

        match element {

146

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

147

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

148

34580

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

149

34580

                let mut index_domains = vec![bound_index_domain_from_length(domain, elems.len())];

150

34580

                index_domains.extend(child_index_domains);

151

34580

                index_domains

152

153

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

154

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

155

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

156

157

34580

})

158

16736

159

160

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

161

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

162

400

pub fn enumerate_index_union_indices(

163

400

    a_domains: &[Moo<GroundDomain>],

164

400

    b_domains: &[Moo<GroundDomain>],

165

400

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

166

400

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

167

        return Err(DomainOpError::WrongType);

168

400

169

400

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

170

400

        .iter()

171

400

        .zip(b_domains.iter())

172

480

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

173

400

        .collect();

174

400

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

175

176

400

    try_enumerate_indices(idx_domains)

177

400

178

179

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

180

640

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

181

640

    let mut remaining = index;

182

640

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

183

184

640

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

185

640

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

186

640

187

188

640

    let mut coords = Vec::new();

189

1280

    for m in multipliers.iter() {

190

        // adjust for 1-based indexing

191

1280

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

192

1280

        remaining %= *m;

193

1280

194

195

640

    coords

196

640

197

198

/// Gets concrete index domains for a matrix expression.

199

///

200

/// For matrix literals, right-unbounded integer index domains like `int(1..)` are bounded using

201

/// the literal's realised size in that dimension. For non-literals, this falls back to the

202

/// expression's resolved domain.

203

31646

pub fn bound_index_domains_of_expr(expr: &Expr) -> Option<Vec<Moo<GroundDomain>>> {

204

31646

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

205

28366

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

206

26158

        return None;

207

};

208

209

2208

    let Some(dimension_lengths) = expr_matrix_dimension_lengths(expr) else {

210

2208

        return Some(index_domains.clone());

211

};

212

213

    assert_eq!(

214

        index_domains.len(),

215

        dimension_lengths.len(),

216

        "matrix literal domain rank should match its realised rank"

217

);

218

219

    Some(

220

        index_domains

221

            .iter()

222

            .cloned()

223

            .zip(dimension_lengths)

224

            .map(|(domain, len)| bound_index_domain_from_length(domain, len))

225

            .collect(),

226

227

31646

228

229

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

230

///

231

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

232

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

233

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

234

///

235

/// # Returns

236

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

237

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

238

3360

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

239

    match m {

240

1440

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

241

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

242

243

1440

            let mut idxs = idxs;

244

2160

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

245

1440

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

246

247

1440

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

248

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

249

};

250

251

1440

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

252

253

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

254

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

255

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

256

                return None;

257

};

258

259

            let index_domains = bound_index_domains_of_expr(inner.as_ref())?;

260

            if index_domains.iter().any(|domain| domain.length().is_err()) {

261

                return None;

262

263

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

264

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

265

266

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

267

268

1920

        _ => Some(Expr::SafeIndex(

269

1920

            Metadata::new(),

270

1920

            Moo::new(m),

271

1920

            vec![idx.into()],

272

1920

)),

273

274

3360

275

276

34580

fn bound_index_domain_from_length(mut domain: Moo<GroundDomain>, len: usize) -> Moo<GroundDomain> {

277

34580

    match Moo::make_mut(&mut domain) {

278

34580

        GroundDomain::Int(ranges) if ranges.len() == 1 && len > 0 => {

279

33860

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

280

                let end = start + (len as i32 - 1);

281

                ranges[0] = Range::Bounded(start, end);

282

33860

283

33860

            domain

284

285

720

        _ => domain,

286

287

34580

288

289

2208

fn expr_matrix_dimension_lengths(expr: &Expr) -> Option<Vec<usize>> {

290

2208

    let (elems, _) = expr.clone().unwrap_matrix_unchecked()?;

291

292

    let child_dimensions = elems

293

        .iter()

294

        .map(|elem| expr_matrix_dimension_lengths(elem).unwrap_or_default())

295

        .collect_vec();

296

297

    assert!(

298

        child_dimensions.iter().all_equal(),

299

        "each child of a matrix should have the same shape"

300

);

301

302

    let mut dimensions = vec![elems.len()];

303

    if let Some(child_dimensions) = child_dimensions.into_iter().next() {

304

        dimensions.extend(child_dimensions);

305

306

    Some(dimensions)

307

2208