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::{DeclarationKind, Literal, 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

            
24
1846
pub fn get_solutions(
25
1846
    solver: Solver,
26
1846
    model: Model,
27
1846
    num_sols: i32,
28
1846
    solver_input_file: &Option<PathBuf>,
29
1846
) -> Result<Vec<BTreeMap<Name, Literal>>, anyhow::Error> {
30
1846
    let adaptor_name = solver.get_name();
31

            
32
1846
    eprintln!("Building {adaptor_name} model...");
33

            
34
    // Create for later since we consume the model when loading it
35
1846
    let symbols_ptr = model.symbols_ptr_unchecked().clone();
36

            
37
1846
    let solver = solver.load_model(model)?;
38

            
39
1846
    if let Some(solver_input_file) = solver_input_file {
40
        eprintln!(
41
            "Writing solver input file to {}",
42
            solver_input_file.display()
43
        );
44
        let file = Box::new(std::fs::File::create(solver_input_file)?);
45
        solver.write_solver_input_file(&mut (file as Box<dyn std::io::Write>))?;
46
1846
    }
47

            
48
1846
    eprintln!("Running {adaptor_name}...");
49

            
50
    // Create two arcs, one to pass into the solver callback, one to get solutions out later
51
1846
    let all_solutions_ref = Arc::new(Mutex::<Vec<BTreeMap<Name, Literal>>>::new(vec![]));
52
1846
    let all_solutions_ref_2 = all_solutions_ref.clone();
53

            
54
1846
    let solver = if num_sols > 0 {
55
        // Get num_sols solutions
56
2
        let sols_left = Mutex::new(num_sols);
57

            
58
        #[allow(clippy::unwrap_used)]
59
2
        solver
60
2
            .solve(Box::new(move |sols| {
61
2
                let mut all_solutions = (*all_solutions_ref_2).lock().unwrap();
62
2
                (*all_solutions).push(sols.into_iter().collect());
63
2
                let mut sols_left = sols_left.lock().unwrap();
64
2
                *sols_left -= 1;
65

            
66
2
                *sols_left != 0
67
2
            }))
68
2
            .unwrap()
69
    } else {
70
        // Get all solutions
71
        #[allow(clippy::unwrap_used)]
72
1844
        solver
73
53874
            .solve(Box::new(move |sols| {
74
53874
                let mut all_solutions = (*all_solutions_ref_2).lock().unwrap();
75
53874
                (*all_solutions).push(sols.into_iter().collect());
76
53874
                true
77
53874
            }))
78
1844
            .unwrap()
79
    };
80

            
81
1846
    solver.save_stats_to_context();
82

            
83
    // Get the collections of solutions and model symbols
84
    #[allow(clippy::unwrap_used)]
85
1846
    let mut sols_guard = (*all_solutions_ref).lock().unwrap();
86
1846
    let sols = &mut *sols_guard;
87
1846
    let symbols = symbols_ptr.read();
88

            
89
    // Get the representations for each variable by name, since some variables are
90
    // divided into multiple auxiliary variables(see crate::representation::Representation)
91
1846
    let names = symbols.clone().into_iter().map(|x| x.0).collect_vec();
92
1846
    let representations = names
93
1846
        .into_iter()
94
83354
        .filter_map(|x| symbols.representations_for(&x).map(|repr| (x, repr)))
95
82894
        .filter_map(|(name, reprs)| {
96
82894
            if reprs.is_empty() {
97
81902
                return None;
98
992
            }
99
992
            assert!(
100
992
                reprs.len() <= 1,
101
                "multiple representations for a variable is not yet implemented"
102
            );
103

            
104
992
            assert_eq!(
105
992
                reprs[0].len(),
106
                1,
107
                "nested representations are not yet implemented"
108
            );
109
992
            Some((name, reprs[0][0].clone()))
110
82894
        })
111
1846
        .collect_vec();
112

            
113
53876
    for sol in sols.iter_mut() {
114
        // Get the value of complex variables using their auxiliary variables
115
67048
        for (name, representation) in representations.iter() {
116
67024
            let value = representation.value_up(sol).unwrap();
117
67024
            sol.insert(name.clone(), value);
118
67024
        }
119

            
120
        // Remove auxiliary variables since we've found the value of the
121
        // variable they represent
122
53876
        *sol = sol
123
53876
            .clone()
124
53876
            .into_iter()
125
13171212
            .filter(|(name, _)| {
126
13171212
                !matches!(name, Name::Represented(_)) && !matches!(name, Name::Machine(_))
127
13171212
            })
128
53876
            .collect();
129
    }
130

            
131
53876
    sols.retain(|x| !x.is_empty());
132
1846
    Ok(sols.clone())
133
1846
}
134

            
135
#[allow(clippy::unwrap_used)]
136
2
pub fn get_solutions_from_conjure(
137
2
    essence_file: &str,
138
2
    param_file: Option<&str>,
139
2
    context: Arc<RwLock<Context<'static>>>,
140
2
) -> Result<Vec<BTreeMap<Name, Literal>>, anyhow::Error> {
141
2
    let tmp_dir = tempdir()?;
142

            
143
2
    let mut cmd = std::process::Command::new("conjure");
144

            
145
2
    cmd.arg("solve")
146
2
        .arg("--number-of-solutions=all")
147
2
        .arg("--copy-solutions=no")
148
2
        .arg("-o")
149
2
        .arg(tmp_dir.path())
150
2
        .arg(essence_file);
151

            
152
2
    if let Some(file) = param_file {
153
2
        cmd.arg(file);
154
2
    }
155

            
156
2
    let output = cmd.output()?;
157

            
158
2
    if !output.status.success() {
159
        let stderr =
160
            String::from_utf8(output.stderr).unwrap_or_else(|e| e.utf8_error().to_string());
161
        return Err(anyhow::Error::msg(format!(
162
            "Error: `conjure solve` exited with code {}; stderr: {}",
163
            output.status, stderr
164
        )));
165
2
    }
166

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

            
170
2
    let solutions_set: Vec<_> = solutions_files
171
2
        .par_iter()
172
2
        .map(|solutions_file| {
173
2
            let solutions_file = solutions_file.as_ref().unwrap();
174
2
            let model = parse_essence_file(solutions_file.to_str().unwrap(), Arc::clone(&context))
175
2
                .expect("conjure solutions files to be parsable");
176

            
177
2
            let mut solutions = BTreeMap::new();
178
2
            for (name, decl) in model.symbols().clone().into_iter() {
179
2
                match &decl.kind() as &DeclarationKind {
180
2
                    conjure_cp::ast::DeclarationKind::ValueLetting(expression, _) => {
181
2
                        let literal = expression
182
2
                            .clone()
183
2
                            .into_literal()
184
2
                            .expect("lettings in a solution should only contain literals");
185
2
                        solutions.insert(name, literal);
186
2
                    }
187
                    _ => {
188
                        bug!("only expect value letting declarations in solutions")
189
                    }
190
                }
191
            }
192
2
            solutions
193
2
        })
194
2
        .collect();
195

            
196
2
    Ok(solutions_set
197
2
        .into_iter()
198
2
        .filter(|x| !x.is_empty())
199
2
        .collect())
200
2
}
201

            
202
3686
pub fn solutions_to_json(solutions: &Vec<BTreeMap<Name, Literal>>) -> JsonValue {
203
3686
    let mut json_solutions = Vec::new();
204
107742
    for solution in solutions {
205
107742
        let mut json_solution = Map::new();
206
356798
        for (var_name, constant) in solution {
207
356798
            let serialized_constant = serde_json::to_value(constant).unwrap();
208
356798
            json_solution.insert(var_name.to_string(), serialized_constant);
209
356798
        }
210
107742
        json_solutions.push(JsonValue::Object(json_solution));
211
    }
212
3686
    let ans = JsonValue::Array(json_solutions);
213
3686
    sort_json_object(&ans, true)
214
3686
}