1
use std::collections::BTreeMap;
2
use std::path::PathBuf;
3
use std::string::ToString;
4
use std::sync::{Arc, Mutex, RwLock};
5

            
6
use conjure_cp::ast::{Atom, DeclarationKind, Expression, GroundDomain, Literal, Metadata, Name};
7
use conjure_cp::bug;
8
use conjure_cp::context::Context;
9

            
10
use serde_json::{Map, Value as JsonValue};
11

            
12
use itertools::Itertools as _;
13
use tempfile::tempdir;
14

            
15
use crate::utils::json::sort_json_object;
16
use conjure_cp::Model;
17
use conjure_cp::parse::tree_sitter::parse_essence_file;
18
use conjure_cp::solver::Solver;
19

            
20
use glob::glob;
21

            
22
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
23
use uniplate::Uniplate;
24
3676

            
25
7364
fn literal_for_reference_domain(
26
7364
    reference_domain: Option<GroundDomain>,
27
7364
    value: &Literal,
28
7364
) -> Option<Literal> {
29
7364
    if matches!(reference_domain, Some(GroundDomain::Bool)) {
30
5684
        return match value {
31
            Literal::Bool(x) => Some(Literal::Bool(*x)),
32
4674
            Literal::Int(1) => Some(Literal::Bool(true)),
33
1010
            Literal::Int(0) => Some(Literal::Bool(false)),
34
            _ => None,
35
3676
        };
36
1680
    }
37
3676

            
38
1680
    Some(value.clone())
39
7364
}
40

            
41
9516
fn substitute_from_solution(
42
9516
    expr: &Expression,
43
9516
    previous_solution: &BTreeMap<Name, Literal>,
44
9516
) -> Option<Expression> {
45
9516
    match expr {
46
5520
        Expression::FromSolution(_, atom_expr) => {
47
1844
            let Atom::Reference(reference) = atom_expr.as_ref() else {
48
3676
                return Some(expr.clone());
49
            };
50

            
51
5520
            let name = reference.name();
52
5520
            let value = previous_solution.get(&name)?;
53
1844
            let reference_domain = reference.resolved_domain().map(|x| x.as_ref().clone());
54
5520
            let value = literal_for_reference_domain(reference_domain, value)?;
55
1844
            Some(Expression::Atomic(Metadata::new(), Atom::Literal(value)))
56
1
        }
57
7672
        _ => Some(expr.clone()),
58
    }
59
9517
}
60
1

            
61
9517
fn substitute_current_solution_refs(
62
9517
    expr: &Expression,
63
9517
    candidate_solution: &BTreeMap<Name, Literal>,
64
9517
) -> Option<Expression> {
65
3688
    match expr {
66
1845
        Expression::Atomic(_, Atom::Reference(reference)) => {
67
1845
            let name = reference.name();
68
1845
            let value = candidate_solution.get(&name)?;
69
1844
            let reference_domain = reference.resolved_domain().map(|x| x.as_ref().clone());
70
1844
            let value = literal_for_reference_domain(reference_domain, value)?;
71
1844
            Some(Expression::Atomic(Metadata::new(), Atom::Literal(value)))
72
3675
        }
73
115397
        _ => Some(expr.clone()),
74
107725
    }
75
117241
}
76
107725

            
77
108078
fn does_solution_dominate(
78
4028
    dominance_expression: &Expression,
79
353
    candidate_solution: &BTreeMap<Name, Literal>,
80
353
    previous_solution: &BTreeMap<Name, Literal>,
81
4029
) -> bool {
82
353
    let expr = dominance_expression
83
9516
        .rewrite(&|e| substitute_from_solution(&e, previous_solution))
84
9516
        .rewrite(&|e| substitute_current_solution_refs(&e, candidate_solution));
85
3676

            
86
3737
    matches!(
87
4029
        conjure_cp::ast::eval::eval_constant(&expr),
88
        Some(Literal::Bool(true))
89
    )
90
353
}
91
3676

            
92
3692
fn retroactively_prune_dominated(
93
3692
    solutions: Vec<BTreeMap<Name, Literal>>,
94
166702
    dominance_expression: &Expression,
95
165788
) -> Vec<BTreeMap<Name, Literal>> {
96
165788
    solutions
97
163804
        .iter()
98
2000
        .enumerate()
99
2079
        .filter_map(|(i, solution)| {
100
2384
            let dominated = solutions.iter().enumerate().any(|(j, candidate)| {
101
400
                i != j && does_solution_dominate(dominance_expression, candidate, solution)
102
400
            });
103

            
104
2079
            if dominated {
105
2045
                None
106
            } else {
107
34
                Some(solution.clone())
108
            }
109
2079
        })
110
165788
        .collect()
111
3692
}
112

            
113
109125
pub fn get_solutions(
114
1399
    solver: Solver,
115
135469
    model: Model,
116
135447
    num_sols: i32,
117
135447
    solver_input_file: &Option<PathBuf>,
118
135447
) -> Result<Vec<BTreeMap<Name, Literal>>, anyhow::Error> {
119
1399
    let dominance_expression = model.dominance.as_ref().map(|expr| match expr {
120
15
        Expression::DominanceRelation(_, inner) => inner.as_ref().clone(),
121
        _ => expr.clone(),
122
107741
    });
123
107726

            
124
109125
    let adaptor_name = solver.get_name();
125
26342398

            
126
26343797
    eprintln!("Building {adaptor_name} model...");
127
26342398

            
128
    // Create for later since we consume the model when loading it
129
1399
    let symbols_ptr = model.symbols_ptr_unchecked().clone();
130

            
131
109125
    let solver = solver.load_model(model)?;
132
3676

            
133
5075
    if let Some(solver_input_file) = solver_input_file {
134
        eprintln!(
135
            "Writing solver input file to {}",
136
1
            solver_input_file.display()
137
1
        );
138
1
        let file = Box::new(std::fs::File::create(solver_input_file)?);
139
1
        solver.write_solver_input_file(&mut (file as Box<dyn std::io::Write>))?;
140
1400
    }
141
1

            
142
1399
    eprintln!("Running {adaptor_name}...");
143
1

            
144
    // Create two arcs, one to pass into the solver callback, one to get solutions out later
145
1400
    let all_solutions_ref = Arc::new(Mutex::<Vec<BTreeMap<Name, Literal>>>::new(vec![]));
146
1400
    let all_solutions_ref_2 = all_solutions_ref.clone();
147
1

            
148
1400
    let solver = if num_sols > 0 {
149
        // Get num_sols solutions
150
1
        let sols_left = Mutex::new(num_sols);
151

            
152
1
        #[allow(clippy::unwrap_used)]
153
1
        solver
154
1
            .solve(Box::new(move |sols| {
155
                let mut all_solutions = (*all_solutions_ref_2).lock().unwrap();
156
1
                (*all_solutions).push(sols.into_iter().collect());
157
                let mut sols_left = sols_left.lock().unwrap();
158
1
                *sols_left -= 1;
159

            
160
                *sols_left != 0
161
            }))
162
            .unwrap()
163
    } else {
164
        // Get all solutions
165
1
        #[allow(clippy::unwrap_used)]
166
1399
        solver
167
40503
            .solve(Box::new(move |sols| {
168
40503
                let mut all_solutions = (*all_solutions_ref_2).lock().unwrap();
169
40502
                (*all_solutions).push(sols.into_iter().collect());
170
40503
                true
171
40503
            }))
172
1400
            .unwrap()
173
1
    };
174
1

            
175
1400
    solver.save_stats_to_context();
176

            
177
    // Get the collections of solutions and model symbols
178
1
    #[allow(clippy::unwrap_used)]
179
1400
    let mut sols_guard = (*all_solutions_ref).lock().unwrap();
180
1400
    let sols = &mut *sols_guard;
181
1400
    let symbols = symbols_ptr.read();
182
1

            
183
    // Get the representations for each variable by name, since some variables are
184
    // divided into multiple auxiliary variables(see crate::representation::Representation)
185
1400
    let names = symbols.clone().into_iter().map(|x| x.0).collect_vec();
186
1400
    let representations = names
187
1399
        .into_iter()
188
62872
        .filter_map(|x| symbols.representations_for(&x).map(|repr| (x, repr)))
189
62530
        .filter_map(|(name, reprs)| {
190
62530
            if reprs.is_empty() {
191
61780
                return None;
192
751
            }
193
751
            assert!(
194
751
                reprs.len() <= 1,
195
                "multiple representations for a variable is not yet implemented"
196
1
            );
197
1

            
198
751
            assert_eq!(
199
751
                reprs[0].len(),
200
1
                1,
201
                "nested representations are not yet implemented"
202
3675
            );
203
4425
            Some((name, reprs[0][0].clone()))
204
170253
        })
205
109122
        .collect_vec();
206
356779

            
207
397281
    for sol in sols.iter_mut() {
208
        // Get the value of complex variables using their auxiliary variables
209
407145
        for (name, representation) in representations.iter() {
210
158069
            let value = representation.value_up(sol).unwrap();
211
50346
            sol.insert(name.clone(), value);
212
54021
        }
213
3675

            
214
        // Remove auxiliary variables since we've found the value of the
215
        // variable they represent
216
40502
        *sol = sol
217
40502
            .clone()
218
40502
            .into_iter()
219
9882116
            .filter(|(name, _)| {
220
9882116
                !matches!(name, Name::Represented(_)) && !matches!(name, Name::Machine(_))
221
9882116
            })
222
40502
            .collect();
223
    }
224

            
225
40502
    sols.retain(|x| !x.is_empty());
226
1399
    if let Some(dominance_expression) = dominance_expression.as_ref() {
227
15
        let pre_prune_len = sols.len();
228
15
        let pruned = retroactively_prune_dominated(sols.clone(), dominance_expression);
229
15
        let post_prune_len = pruned.len();
230
15

            
231
15
        eprintln!("Dominance pruning retained {post_prune_len} of {pre_prune_len} solutions.");
232
15

            
233
15
        *sols = pruned;
234
1384
    }
235

            
236
1399
    Ok(sols.clone())
237
1399
}
238

            
239
#[allow(clippy::unwrap_used)]
240
pub fn get_solutions_from_conjure(
241
    essence_file: &str,
242
    context: Arc<RwLock<Context<'static>>>,
243
) -> Result<Vec<BTreeMap<Name, Literal>>, anyhow::Error> {
244
    let tmp_dir = tempdir()?;
245

            
246
    let mut cmd = std::process::Command::new("conjure");
247
    let output = cmd
248
        .arg("solve")
249
        .arg("--number-of-solutions=all")
250
        .arg("--copy-solutions=no")
251
        .arg("-o")
252
        .arg(tmp_dir.path())
253
        .arg(essence_file)
254
        .output()?;
255

            
256
    if !output.status.success() {
257
        let stderr =
258
            String::from_utf8(output.stderr).unwrap_or_else(|e| e.utf8_error().to_string());
259
        return Err(anyhow::Error::msg(format!(
260
            "Error: `conjure solve` exited with code {}; stderr: {}",
261
            output.status, stderr
262
        )));
263
    }
264

            
265
    let solutions_files: Vec<_> =
266
        glob(&format!("{}/*.solution", tmp_dir.path().display()))?.collect();
267

            
268
    let solutions_set: Vec<_> = solutions_files
269
        .par_iter()
270
        .map(|solutions_file| {
271
            let solutions_file = solutions_file.as_ref().unwrap();
272
            let model = parse_essence_file(solutions_file.to_str().unwrap(), Arc::clone(&context))
273
                .expect("conjure solutions files to be parsable");
274

            
275
            let mut solutions = BTreeMap::new();
276
            for (name, decl) in model.symbols().clone().into_iter() {
277
                match &decl.kind() as &DeclarationKind {
278
                    conjure_cp::ast::DeclarationKind::ValueLetting(expression) => {
279
                        let literal = expression
280
                            .clone()
281
                            .into_literal()
282
                            .expect("lettings in a solution should only contain literals");
283
                        solutions.insert(name, literal);
284
                    }
285
                    _ => {
286
                        bug!("only expect value letting declarations in solutions")
287
                    }
288
                }
289
            }
290
            solutions
291
        })
292
        .collect();
293

            
294
    Ok(solutions_set
295
        .into_iter()
296
        .filter(|x| !x.is_empty())
297
        .collect())
298
}
299

            
300
2794
pub fn solutions_to_json(solutions: &Vec<BTreeMap<Name, Literal>>) -> JsonValue {
301
2794
    let mut json_solutions = Vec::new();
302
80876
    for solution in solutions {
303
80876
        let mut json_solution = Map::new();
304
267836
        for (var_name, constant) in solution {
305
267836
            let serialized_constant = serde_json::to_value(constant).unwrap();
306
267836
            json_solution.insert(var_name.to_string(), serialized_constant);
307
267836
        }
308
80876
        json_solutions.push(JsonValue::Object(json_solution));
309
    }
310
2794
    let ans = JsonValue::Array(json_solutions);
311
2794
    sort_json_object(&ans, true)
312
2794
}
313

            
314
#[cfg(test)]
315
mod tests {
316
    use super::*;
317
    use conjure_cp::ast::{DeclarationPtr, Domain, Moo, Reference};
318

            
319
    #[test]
320
1
    fn retroactive_pruning_removes_dominated_prior_solution() {
321
1
        let x = Name::User("x".into());
322
1
        let x_ref = Expression::Atomic(
323
1
            Metadata::new(),
324
1
            Atom::Reference(Reference::new(DeclarationPtr::new_find(
325
1
                x.clone(),
326
1
                Domain::bool(),
327
1
            ))),
328
1
        );
329
1
        let dominance_expression = Expression::Imply(
330
1
            Metadata::new(),
331
1
            Moo::new(x_ref.clone()),
332
1
            Moo::new(Expression::FromSolution(
333
1
                Metadata::new(),
334
1
                Moo::new(Atom::Reference(Reference::new(DeclarationPtr::new_find(
335
1
                    x.clone(),
336
1
                    Domain::bool(),
337
1
                )))),
338
1
            )),
339
1
        );
340

            
341
1
        let mut sol_true = BTreeMap::new();
342
1
        sol_true.insert(x.clone(), Literal::Int(1));
343
1
        let mut sol_false = BTreeMap::new();
344
1
        sol_false.insert(x.clone(), Literal::Int(0));
345

            
346
1
        let pruned =
347
1
            retroactively_prune_dominated(vec![sol_true, sol_false.clone()], &dominance_expression);
348

            
349
1
        assert_eq!(pruned, vec![sol_false]);
350
1
    }
351
}