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
use conjure_cp::settings::{configured_rule_trace_enabled, set_rule_trace_enabled};
10

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

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

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

            
21
use glob::glob;
22

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

            
26
/// Coerces a literal into the type expected by a reference domain when possible.
27
///
28
/// This is currently used to turn `0`/`1` solver outputs back into boolean
29
/// literals before substituting them into dominance expressions.
30
161572
fn literal_for_reference_domain(
31
161572
    reference_domain: Option<GroundDomain>,
32
161572
    value: &Literal,
33
161572
) -> Option<Literal> {
34
161572
    if matches!(reference_domain, Some(GroundDomain::Bool)) {
35
82516
        return match value {
36
23040
            Literal::Bool(x) => Some(Literal::Bool(*x)),
37
29342
            Literal::Int(1) => Some(Literal::Bool(true)),
38
30134
            Literal::Int(0) => Some(Literal::Bool(false)),
39
            _ => None,
40
        };
41
79056
    }
42

            
43
79056
    Some(value.clone())
44
161572
}
45

            
46
/// Replaces `fromSolution(x)` occurrences with the value of `x` from a previous solution.
47
440538
fn substitute_from_solution(
48
440538
    expr: &Expression,
49
440538
    previous_solution: &BTreeMap<Name, Literal>,
50
440538
) -> Option<Expression> {
51
440538
    match expr {
52
80786
        Expression::FromSolution(_, atom_expr) => {
53
80786
            let Atom::Reference(reference) = atom_expr.as_ref() else {
54
                return Some(expr.clone());
55
            };
56

            
57
80786
            let name = reference.name();
58
80786
            let value = previous_solution.get(&name)?;
59
80786
            let reference_domain = reference.resolved_domain().map(|x| x.as_ref().clone());
60
80786
            let value = literal_for_reference_domain(reference_domain, value)?;
61
80786
            Some(Expression::Atomic(Metadata::new(), Atom::Literal(value)))
62
        }
63
359752
        _ => Some(expr.clone()),
64
    }
65
440538
}
66

            
67
/// Replaces direct variable references with values from the candidate solution.
68
440538
fn substitute_current_solution_refs(
69
440538
    expr: &Expression,
70
440538
    candidate_solution: &BTreeMap<Name, Literal>,
71
440538
) -> Option<Expression> {
72
174748
    match expr {
73
80786
        Expression::Atomic(_, Atom::Reference(reference)) => {
74
80786
            let name = reference.name();
75
80786
            let value = candidate_solution.get(&name)?;
76
80786
            let reference_domain = reference.resolved_domain().map(|x| x.as_ref().clone());
77
80786
            let value = literal_for_reference_domain(reference_domain, value)?;
78
80786
            Some(Expression::Atomic(Metadata::new(), Atom::Literal(value)))
79
        }
80
359752
        _ => Some(expr.clone()),
81
    }
82
440538
}
83

            
84
/// Evaluates whether `candidate_solution` dominates `previous_solution`.
85
///
86
/// The dominance expression is instantiated with values from both solutions and
87
/// then constant-folded to a boolean result.
88
15500
fn does_solution_dominate(
89
15500
    dominance_expression: &Expression,
90
15500
    candidate_solution: &BTreeMap<Name, Literal>,
91
15500
    previous_solution: &BTreeMap<Name, Literal>,
92
15500
) -> bool {
93
15500
    let expr = dominance_expression
94
440538
        .rewrite(&|e| substitute_from_solution(&e, previous_solution))
95
440538
        .rewrite(&|e| substitute_current_solution_refs(&e, candidate_solution));
96

            
97
2521
    matches!(
98
15500
        conjure_cp::ast::eval::eval_constant(&expr),
99
        Some(Literal::Bool(true))
100
    )
101
15500
}
102

            
103
/// Removes solutions that are dominated by another solution in the result set.
104
1315
fn retroactively_prune_dominated(
105
1315
    solutions: Vec<BTreeMap<Name, Literal>>,
106
1315
    dominance_expression: &Expression,
107
1315
) -> Vec<BTreeMap<Name, Literal>> {
108
1315
    solutions
109
1315
        .iter()
110
1315
        .enumerate()
111
5528
        .filter_map(|(i, solution)| {
112
18580
            let dominated = solutions.iter().enumerate().any(|(j, candidate)| {
113
18580
                i != j && does_solution_dominate(dominance_expression, candidate, solution)
114
18580
            });
115

            
116
5528
            if dominated {
117
2521
                None
118
            } else {
119
3007
                Some(solution.clone())
120
            }
121
5528
        })
122
1315
        .collect()
123
1315
}
124

            
125
6783
pub fn get_solutions(
126
6783
    solver: Solver,
127
6783
    model: Model,
128
6783
    num_sols: i32,
129
6783
    solver_input_file: &Option<PathBuf>,
130
6783
    rule_trace_cdp: bool,
131
6783
) -> Result<Vec<BTreeMap<Name, Literal>>, anyhow::Error> {
132
6783
    set_rule_trace_enabled(rule_trace_cdp && configured_rule_trace_enabled());
133

            
134
6783
    let dominance_expression = model.dominance.as_ref().map(|expr| match expr {
135
1314
        Expression::DominanceRelation(_, inner) => inner.as_ref().clone(),
136
        _ => expr.clone(),
137
1314
    });
138

            
139
6783
    let adaptor_name = solver.get_name();
140

            
141
6783
    eprintln!("Building {adaptor_name} model...");
142

            
143
    // Create for later since we consume the model when loading it
144
6783
    let symbols_ptr = model.symbols_ptr_unchecked().clone();
145

            
146
6783
    let solver = solver.load_model(model)?;
147

            
148
6783
    if let Some(solver_input_file) = solver_input_file {
149
        eprintln!(
150
            "Writing solver input file to {}",
151
            solver_input_file.display()
152
        );
153
        let file = Box::new(std::fs::File::create(solver_input_file)?);
154
        solver.write_solver_input_file(&mut (file as Box<dyn std::io::Write>))?;
155
6783
    }
156

            
157
6783
    eprintln!("Running {adaptor_name}...");
158

            
159
    // Create two arcs, one to pass into the solver callback, one to get solutions out later
160
6783
    let all_solutions_ref = Arc::new(Mutex::<Vec<BTreeMap<Name, Literal>>>::new(vec![]));
161
6783
    let all_solutions_ref_2 = all_solutions_ref.clone();
162

            
163
6783
    let solver = if num_sols > 0 {
164
        // Get num_sols solutions
165
18
        let sols_left = Mutex::new(num_sols);
166

            
167
18
        solver
168
18
            .solve(Box::new(move |sols| {
169
12
                let mut all_solutions = (*all_solutions_ref_2).lock().unwrap();
170
12
                (*all_solutions).push(sols.into_iter().collect());
171
12
                let mut sols_left = sols_left.lock().unwrap();
172
12
                *sols_left -= 1;
173

            
174
12
                *sols_left != 0
175
12
            }))
176
18
            .map_err(|err| anyhow::anyhow!("solver failed while collecting solutions: {err}"))?
177
    } else {
178
        // Get all solutions
179
6765
        solver
180
148153
            .solve(Box::new(move |sols| {
181
148153
                let mut all_solutions = (*all_solutions_ref_2).lock().unwrap();
182
148153
                (*all_solutions).push(sols.into_iter().collect());
183
148153
                true
184
148153
            }))
185
6765
            .map_err(|err| anyhow::anyhow!("solver failed while collecting solutions: {err}"))?
186
    };
187

            
188
6783
    solver.save_stats_to_context();
189

            
190
    // Get the collections of solutions and model symbols
191
    #[allow(clippy::unwrap_used)]
192
6783
    let mut sols_guard = (*all_solutions_ref).lock().unwrap();
193
6783
    let sols = &mut *sols_guard;
194
6783
    let symbols = symbols_ptr.read();
195

            
196
    // Get the representations for each variable by name, since some variables are
197
    // divided into multiple auxiliary variables(see crate::representation::Representation)
198
6783
    let names = symbols.clone().into_iter().map(|x| x.0).collect_vec();
199
6783
    let representations = names
200
6783
        .into_iter()
201
569145
        .filter_map(|x| symbols.representations_for(&x).map(|repr| (x, repr)))
202
568077
        .filter_map(|(name, reprs)| {
203
568077
            if reprs.is_empty() {
204
563979
                return None;
205
4098
            }
206
4098
            assert!(
207
4098
                reprs.len() <= 1,
208
                "multiple representations for a variable is not yet implemented"
209
            );
210

            
211
4098
            assert_eq!(
212
4098
                reprs[0].len(),
213
                1,
214
                "nested representations are not yet implemented"
215
            );
216
4098
            Some((name, reprs[0][0].clone()))
217
568077
        })
218
6783
        .collect_vec();
219

            
220
148165
    for sol in sols.iter_mut() {
221
        // Get the value of complex variables using their auxiliary variables
222
185800
        for (name, representation) in representations.iter() {
223
185782
            let value = representation.value_up(sol).map_err(|err| {
224
                anyhow::anyhow!(
225
                    "failed to reconstruct value for variable {name} from solver solution: {err}"
226
                )
227
            })?;
228
185782
            sol.insert(name.clone(), value);
229
        }
230

            
231
        // Remove auxiliary variables since we've found the value of the
232
        // variable they represent
233
148165
        *sol = sol
234
148165
            .clone()
235
148165
            .into_iter()
236
3232451
            .filter(|(name, _)| {
237
3232451
                !matches!(name, Name::Represented(_)) && !matches!(name, Name::Machine(_))
238
3232451
            })
239
148165
            .collect();
240
    }
241

            
242
148165
    sols.retain(|x| !x.is_empty());
243
6783
    if let Some(dominance_expression) = dominance_expression.as_ref() {
244
1314
        let pre_prune_len = sols.len();
245
1314
        let pruned = retroactively_prune_dominated(sols.clone(), dominance_expression);
246
1314
        let post_prune_len = pruned.len();
247
1314

            
248
1314
        eprintln!("Dominance pruning retained {post_prune_len} of {pre_prune_len} solutions.");
249
1314

            
250
1314
        *sols = pruned;
251
5469
    }
252

            
253
6783
    Ok(sols.clone())
254
6783
}
255

            
256
#[allow(clippy::unwrap_used)]
257
2
pub fn get_solutions_from_conjure(
258
2
    essence_file: &str,
259
2
    param_file: Option<&str>,
260
2
    context: Arc<RwLock<Context<'static>>>,
261
2
) -> Result<Vec<BTreeMap<Name, Literal>>, anyhow::Error> {
262
2
    let tmp_dir = tempdir()?;
263

            
264
2
    let mut cmd = std::process::Command::new("conjure");
265

            
266
2
    cmd.arg("solve")
267
2
        .arg("--number-of-solutions=all")
268
2
        .arg("--copy-solutions=no")
269
2
        .arg("-o")
270
2
        .arg(tmp_dir.path())
271
2
        .arg(essence_file);
272

            
273
2
    if let Some(file) = param_file {
274
2
        cmd.arg(file);
275
2
    }
276

            
277
2
    let output = cmd.output()?;
278

            
279
2
    if !output.status.success() {
280
        let stderr =
281
            String::from_utf8(output.stderr).unwrap_or_else(|e| e.utf8_error().to_string());
282
        return Err(anyhow::Error::msg(format!(
283
            "Error: `conjure solve` exited with code {}; stderr: {}",
284
            output.status, stderr
285
        )));
286
2
    }
287

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

            
291
2
    let solutions_set: Vec<_> = solutions_files
292
2
        .par_iter()
293
2
        .map(|solutions_file| {
294
2
            let solutions_file = solutions_file.as_ref().unwrap();
295
2
            let model = parse_essence_file(solutions_file.to_str().unwrap(), Arc::clone(&context))
296
2
                .expect("conjure solutions files to be parsable");
297

            
298
2
            let mut solutions = BTreeMap::new();
299
2
            for (name, decl) in model.symbols().clone().into_iter() {
300
2
                match &decl.kind() as &DeclarationKind {
301
2
                    conjure_cp::ast::DeclarationKind::ValueLetting(expression, _) => {
302
2
                        let literal = expression
303
2
                            .clone()
304
2
                            .into_literal()
305
2
                            .expect("lettings in a solution should only contain literals");
306
2
                        solutions.insert(name, literal);
307
2
                    }
308
                    _ => {
309
                        bug!("only expect value letting declarations in solutions")
310
                    }
311
                }
312
            }
313
2
            solutions
314
2
        })
315
2
        .collect();
316

            
317
2
    Ok(solutions_set
318
2
        .into_iter()
319
2
        .filter(|x| !x.is_empty())
320
2
        .collect())
321
2
}
322

            
323
13548
pub fn solutions_to_json(solutions: &Vec<BTreeMap<Name, Literal>>) -> JsonValue {
324
13548
    let mut json_solutions = Vec::new();
325
291272
    for solution in solutions {
326
291272
        let mut json_solution = Map::new();
327
987982
        for (var_name, constant) in solution {
328
987982
            let serialized_constant = serde_json::to_value(constant).unwrap();
329
987982
            json_solution.insert(var_name.to_string(), serialized_constant);
330
987982
        }
331
291272
        json_solutions.push(JsonValue::Object(json_solution));
332
    }
333
13548
    let ans = JsonValue::Array(json_solutions);
334
13548
    sort_json_object(&ans, true)
335
13548
}
336

            
337
#[cfg(test)]
338
mod tests {
339
    use super::*;
340
    use conjure_cp::ast::{DeclarationPtr, Domain, Moo, Reference};
341

            
342
    #[test]
343
1
    fn retroactive_pruning_removes_dominated_prior_solution() {
344
1
        let x = Name::User("x".into());
345
1
        let x_ref = Expression::Atomic(
346
1
            Metadata::new(),
347
1
            Atom::Reference(Reference::new(DeclarationPtr::new_find(
348
1
                x.clone(),
349
1
                Domain::bool(),
350
1
            ))),
351
1
        );
352
1
        let dominance_expression = Expression::Imply(
353
1
            Metadata::new(),
354
1
            Moo::new(x_ref.clone()),
355
1
            Moo::new(Expression::FromSolution(
356
1
                Metadata::new(),
357
1
                Moo::new(Atom::Reference(Reference::new(DeclarationPtr::new_find(
358
1
                    x.clone(),
359
1
                    Domain::bool(),
360
1
                )))),
361
1
            )),
362
1
        );
363

            
364
1
        let mut sol_true = BTreeMap::new();
365
1
        sol_true.insert(x.clone(), Literal::Int(1));
366
1
        let mut sol_false = BTreeMap::new();
367
1
        sol_false.insert(x.clone(), Literal::Int(0));
368

            
369
1
        let pruned =
370
1
            retroactively_prune_dominated(vec![sol_true, sol_false.clone()], &dominance_expression);
371

            
372
1
        assert_eq!(pruned, vec![sol_false]);
373
1
    }
374
}