1
//! conjure_oxide solve sub-command
2
#![allow(clippy::unwrap_used)]
3
use std::time::Duration;
4
use std::{
5
    fs::File,
6
    io::Write as _,
7
    path::PathBuf,
8
    process::exit,
9
    sync::{Arc, RwLock},
10
};
11

            
12
use anyhow::anyhow;
13
use clap::ValueHint;
14
use conjure_cp::defaults::DEFAULT_RULE_SETS;
15
use conjure_cp::instantiate::instantiate_model;
16
use conjure_cp::{
17
    Model,
18
    context::Context,
19
    rule_engine::{resolve_rule_sets, rewrite_morph, rewrite_naive},
20
    settings::{
21
        Rewriter, set_comprehension_expander, set_current_parser, set_current_rewriter,
22
        set_current_solver_family, set_minion_discrete_threshold,
23
    },
24
    solver::Solver,
25
};
26
use conjure_cp::{
27
    parse::conjure_json::model_from_json, rule_engine::get_rules, settings::SolverFamily,
28
};
29
use conjure_cp::{parse::tree_sitter::parse_essence_file_native, solver::adaptors::*};
30
use conjure_cp_cli::find_conjure::conjure_executable;
31
use conjure_cp_cli::utils::conjure::{get_solutions, solutions_to_json};
32
use serde_json::to_string_pretty;
33

            
34
use crate::cli::{GlobalArgs, LOGGING_HELP_HEADING};
35

            
36
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
37
pub enum NumberOfSolutions {
38
    All,
39
    Limit(i32),
40
}
41

            
42
impl NumberOfSolutions {
43
12
    fn as_solver_limit(self) -> i32 {
44
12
        match self {
45
10
            NumberOfSolutions::All => 0,
46
2
            NumberOfSolutions::Limit(limit) => limit,
47
        }
48
12
    }
49
}
50

            
51
26
fn parse_number_of_solutions(input: &str) -> Result<NumberOfSolutions, String> {
52
26
    if input.eq_ignore_ascii_case("all") {
53
10
        return Ok(NumberOfSolutions::All);
54
16
    }
55

            
56
16
    let limit = input
57
16
        .parse::<i32>()
58
16
        .map_err(|_| "expected a positive integer or 'all'".to_string())?;
59

            
60
16
    if limit <= 0 {
61
        return Err("expected a positive integer or 'all'".to_string());
62
16
    }
63

            
64
16
    Ok(NumberOfSolutions::Limit(limit))
65
26
}
66

            
67
#[derive(Clone, Debug, clap::Args)]
68
pub struct Args {
69
    /// The input Essence problem file
70
    #[arg(value_name = "INPUT_ESSENCE", value_hint = ValueHint::FilePath)]
71
    pub essence_file: PathBuf,
72

            
73
    /// The input Essence parameter file
74
    #[arg(value_name = "PARAM_ESSENCE", value_hint = ValueHint::FilePath)]
75
    pub param_file: Option<PathBuf>,
76

            
77
    /// Save execution info as JSON to the given filepath.
78
    #[arg(long ,value_hint=ValueHint::FilePath,help_heading=LOGGING_HELP_HEADING)]
79
    pub info_json_path: Option<PathBuf>,
80

            
81
    /// Do not run the solver.
82
    ///
83
    /// The rewritten model is printed to stdout in an Essence-style syntax
84
    /// (but is not necessarily valid Essence).
85
    #[arg(long, default_value_t = false)]
86
    pub no_run_solver: bool,
87

            
88
    /// Number of solutions to return. Use a positive integer, or `all`.
89
    #[arg(
90
        long,
91
        short = 'n',
92
        default_value = "1",
93
        value_name = "N|all",
94
        value_parser = parse_number_of_solutions
95
    )]
96
    pub number_of_solutions: NumberOfSolutions,
97

            
98
    /// Save solutions to the given JSON file
99
    #[arg(long, short = 'o', value_hint = ValueHint::FilePath,help_heading=LOGGING_HELP_HEADING)]
100
    pub output: Option<PathBuf>,
101
}
102

            
103
26
pub fn run_solve_command(global_args: GlobalArgs, solve_args: Args) -> anyhow::Result<()> {
104
26
    let essence_file = solve_args.essence_file.clone();
105
26
    let param_file = solve_args.param_file.clone();
106

            
107
    // each step is in its own method so that similar commands
108
    // (e.g. testsolve) can reuse some of these steps.
109

            
110
26
    let context = init_context(&global_args, essence_file, param_file)?;
111

            
112
26
    let ctx_lock = context.read().unwrap();
113
26
    let essence_file_name = ctx_lock
114
26
        .essence_file_name
115
26
        .as_ref()
116
26
        .expect("context should contain the problem input file");
117
26
    let param_file_name = ctx_lock.param_file_name.as_ref();
118

            
119
    // parse models
120
26
    let problem_model = parse(&global_args, Arc::clone(&context), essence_file_name)?;
121

            
122
    // unify models
123
22
    let unified_model = match param_file_name {
124
8
        Some(param_file_name) => {
125
8
            let param_model = parse(&global_args, Arc::clone(&context), param_file_name)?;
126
8
            instantiate_model(problem_model, param_model)?
127
        }
128
14
        None => problem_model,
129
    };
130
16
    drop(ctx_lock);
131

            
132
16
    let rewritten_model = rewrite(unified_model, &global_args, Arc::clone(&context))?;
133

            
134
16
    let solver = init_solver(&global_args);
135

            
136
16
    if solve_args.no_run_solver {
137
4
        println!("{}", &rewritten_model);
138

            
139
4
        if let Some(path) = global_args.save_solver_input_file {
140
            let solver = solver.load_model(rewritten_model)?;
141
            eprintln!("Writing solver input file to {}", path.display());
142
            let mut file: Box<dyn std::io::Write> = Box::new(File::create(path)?);
143
            solver.write_solver_input_file(&mut file)?;
144
4
        }
145
    } else {
146
12
        run_solver(solver, &global_args, &solve_args, rewritten_model)?
147
    }
148

            
149
    // still do postamble even if we didn't run the solver
150
16
    if let Some(ref path) = solve_args.info_json_path {
151
        let context_obj = context.read().unwrap().clone();
152
        let generated_json = &serde_json::to_value(context_obj)?;
153
        let pretty_json = serde_json::to_string_pretty(&generated_json)?;
154
        File::create(path)?.write_all(pretty_json.as_bytes())?;
155
16
    }
156
16
    Ok(())
157
26
}
158

            
159
/// Returns a new Context and Solver for solving.
160
44
pub(crate) fn init_context(
161
44
    global_args: &GlobalArgs,
162
44
    essence_file: PathBuf,
163
44
    param_file: Option<PathBuf>,
164
44
) -> anyhow::Result<Arc<RwLock<Context<'static>>>> {
165
44
    set_current_parser(global_args.parser);
166
44
    set_current_rewriter(global_args.rewriter);
167
44
    set_comprehension_expander(global_args.comprehension_expander);
168
44
    set_current_solver_family(global_args.solver);
169
44
    set_minion_discrete_threshold(global_args.minion_discrete_threshold);
170

            
171
44
    let target_family = global_args.solver;
172
44
    let mut extra_rule_sets: Vec<&str> = DEFAULT_RULE_SETS.to_vec();
173
44
    for rs in &global_args.extra_rule_sets {
174
        extra_rule_sets.push(rs.as_str());
175
    }
176

            
177
44
    if let SolverFamily::Sat(sat_encoding) = target_family {
178
        extra_rule_sets.push(sat_encoding.as_rule_set());
179
44
    }
180

            
181
44
    let rule_sets = match resolve_rule_sets(target_family, &extra_rule_sets) {
182
44
        Ok(rs) => rs,
183
        Err(e) => {
184
            tracing::error!("Error resolving rule sets: {}", e);
185
            exit(1);
186
        }
187
    };
188

            
189
44
    let pretty_rule_sets = rule_sets
190
44
        .iter()
191
44
        .map(|rule_set| rule_set.name)
192
44
        .collect::<Vec<_>>()
193
44
        .join(", ");
194

            
195
44
    tracing::info!("Enabled rule sets: [{}]", pretty_rule_sets);
196
44
    tracing::info!(
197
        target: "file",
198
        "Rule sets: {}",
199
        pretty_rule_sets
200
    );
201

            
202
44
    let rules = get_rules(&rule_sets)?.into_iter().collect::<Vec<_>>();
203
44
    tracing::info!(
204
        target: "file",
205
        "Rules: {}",
206
950
        rules.iter().map(|rd| format!("{rd}")).collect::<Vec<_>>().join("\n")
207
    );
208
44
    let context = Context::new_ptr(
209
44
        target_family,
210
132
        extra_rule_sets.iter().map(|rs| rs.to_string()).collect(),
211
44
        rules,
212
44
        rule_sets.clone(),
213
    );
214

            
215
44
    context.write().unwrap().essence_file_name = Some(essence_file.to_str().expect("").into());
216
44
    if let Some(param_file) = param_file {
217
12
        context.write().unwrap().param_file_name = Some(param_file.to_str().expect("").into());
218
32
    }
219

            
220
44
    Ok(context)
221
44
}
222

            
223
18
pub(crate) fn init_solver(global_args: &GlobalArgs) -> Solver {
224
18
    let family = global_args.solver;
225
18
    let timeout_ms = global_args
226
18
        .solver_timeout
227
18
        .map(|dur| Duration::from(dur).as_millis())
228
18
        .map(|timeout_ms| u64::try_from(timeout_ms).expect("Timeout too large"));
229

            
230
18
    match family {
231
18
        SolverFamily::Minion => Solver::new(Minion::default()),
232
        SolverFamily::Sat(_) => Solver::new(Sat::default()),
233
        SolverFamily::Smt(theory_cfg) => Solver::new(Smt::new(timeout_ms, theory_cfg)),
234
    }
235
18
}
236

            
237
56
pub(crate) fn parse(
238
56
    global_args: &GlobalArgs,
239
56
    context: Arc<RwLock<Context<'static>>>,
240
56
    file_path: &str,
241
56
) -> anyhow::Result<Model> {
242
56
    tracing::info!(target: "file", "Input file: {}", file_path);
243

            
244
56
    match global_args.parser {
245
        conjure_cp::settings::Parser::TreeSitter => {
246
14
            parse_essence_file_native(file_path, context.clone()).map_err(|e| e.into())
247
        }
248
42
        conjure_cp::settings::Parser::ViaConjure => parse_with_conjure(file_path, context.clone()),
249
    }
250
56
}
251

            
252
42
pub(crate) fn parse_with_conjure(
253
42
    input_file: &str,
254
42
    context: Arc<RwLock<Context<'static>>>,
255
42
) -> anyhow::Result<Model> {
256
42
    conjure_executable().map_err(|e| anyhow!("Could not find correct conjure executable: {e}"))?;
257

            
258
42
    let mut cmd = std::process::Command::new("conjure");
259
42
    let output = cmd
260
42
        .arg("pretty")
261
42
        .arg("--output-format=astjson")
262
42
        .arg(input_file)
263
42
        .output()?;
264

            
265
42
    if !output.status.success() {
266
        println!("Parsing error: {}", String::from_utf8(output.stderr)?);
267
42
    }
268

            
269
42
    let astjson = String::from_utf8(output.stdout)?;
270

            
271
42
    if cfg!(feature = "extra-rule-checks") {
272
42
        tracing::info!("extra-rule-checks: enabled");
273
    } else {
274
        tracing::info!("extra-rule-checks: disabled");
275
    }
276

            
277
42
    model_from_json(&astjson, context.clone()).map_err(|e| anyhow!(e))
278
42
}
279

            
280
18
pub(crate) fn rewrite(
281
18
    model: Model,
282
18
    global_args: &GlobalArgs,
283
18
    context: Arc<RwLock<Context<'static>>>,
284
18
) -> anyhow::Result<Model> {
285
18
    tracing::info!("Initial model: \n{}\n", model);
286

            
287
18
    set_current_rewriter(global_args.rewriter);
288

            
289
18
    let comprehension_expander = global_args.comprehension_expander;
290
18
    set_comprehension_expander(comprehension_expander);
291
18
    tracing::info!("Comprehension expander: {}", comprehension_expander);
292

            
293
18
    let rule_sets = context.read().unwrap().rule_sets.clone();
294

            
295
18
    let new_model = match global_args.rewriter {
296
        Rewriter::Morph => {
297
            tracing::info!("Rewriting the model using the morph rewriter");
298
            rewrite_morph(
299
                model,
300
                &rule_sets,
301
                global_args.check_equally_applicable_rules,
302
            )
303
        }
304
        Rewriter::Naive => {
305
18
            tracing::info!("Rewriting the model using the default / naive rewriter");
306
18
            rewrite_naive(
307
18
                &model,
308
18
                &rule_sets,
309
18
                global_args.check_equally_applicable_rules,
310
            )?
311
        }
312
    };
313

            
314
18
    tracing::info!("Rewritten model: \n{}\n", new_model);
315
18
    Ok(new_model)
316
18
}
317

            
318
12
fn run_solver(
319
12
    solver: Solver,
320
12
    global_args: &GlobalArgs,
321
12
    cmd_args: &Args,
322
12
    model: Model,
323
12
) -> anyhow::Result<()> {
324
12
    let out_file: Option<File> = match &cmd_args.output {
325
12
        None => None,
326
        Some(pth) => Some(
327
            File::options()
328
                .create(true)
329
                .truncate(true)
330
                .write(true)
331
                .open(pth)?,
332
        ),
333
    };
334

            
335
12
    let solutions = get_solutions(
336
12
        solver,
337
12
        model,
338
12
        cmd_args.number_of_solutions.as_solver_limit(),
339
12
        &global_args.save_solver_input_file,
340
    )?;
341
12
    tracing::info!(target: "file", "Solutions: {}", solutions_to_json(&solutions));
342

            
343
12
    let solutions_json = solutions_to_json(&solutions);
344
12
    let solutions_str = to_string_pretty(&solutions_json)?;
345
12
    match out_file {
346
12
        None => {
347
12
            println!("Solutions:");
348
12
            println!("{solutions_str}");
349
12
        }
350
        Some(mut outf) => {
351
            outf.write_all(solutions_str.as_bytes())?;
352
            println!(
353
                "Solutions saved to {:?}",
354
                &cmd_args.output.clone().unwrap().canonicalize()?
355
            )
356
        }
357
    }
358
12
    Ok(())
359
12
}