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

            
7
use conjure_cp::ast::{DeclarationKind, Literal, Name};
8
use conjure_cp::bug;
9
use conjure_cp::context::Context;
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

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

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

            
35
    // Create for later since we consume the model when loading it
36
    let symbols_rc = Rc::clone(model.as_submodel().symbols_ptr_unchecked());
37

            
38
    let solver = solver.load_model(model)?;
39

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

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

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

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

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

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

            
82
    solver.save_stats_to_context();
83

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

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

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

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

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

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

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

            
143
    let mut cmd = std::process::Command::new("conjure");
144
    let output = cmd
145
        .arg("solve")
146
        .arg("--number-of-solutions=all")
147
        .arg("--copy-solutions=no")
148
        .arg("-o")
149
        .arg(tmp_dir.path())
150
        .arg(essence_file)
151
        .output()?;
152

            
153
    if !output.status.success() {
154
        let stderr =
155
            String::from_utf8(output.stderr).unwrap_or_else(|e| e.utf8_error().to_string());
156
        return Err(anyhow::Error::msg(format!(
157
            "Error: `conjure solve` exited with code {}; stderr: {}",
158
            output.status, stderr
159
        )));
160
    }
161

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

            
165
    let solutions_set: Vec<_> = solutions_files
166
        .par_iter()
167
        .map(|solutions_file| {
168
            let solutions_file = solutions_file.as_ref().unwrap();
169
            let model = parse_essence_file(solutions_file.to_str().unwrap(), Arc::clone(&context))
170
                .expect("conjure solutions files to be parsable");
171

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

            
191
    Ok(solutions_set
192
        .into_iter()
193
        .filter(|x| !x.is_empty())
194
        .collect())
195
}
196

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