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

            
10
use crate::exit;
11
use anyhow::anyhow;
12
use clap::ValueHint;
13
use conjure_cp::instantiate::{instantiate_model, validate_instantiation_conditions};
14
use conjure_cp::{
15
    Model,
16
    context::Context,
17
    defaults::DEFAULT_RULE_SETS,
18
    rule_engine::{resolve_rule_sets, rewrite_model},
19
    settings::{
20
        Rewriter, set_channelling, set_comprehension_expander, set_current_parser,
21
        set_current_rewriter, set_current_solver_family, set_default_rule_trace_enabled,
22
        set_heuristic, set_heuristic_responses, set_heuristic_seed, set_minion_discrete_threshold,
23
        set_rule_attempt_trace_enabled, set_rule_trace_aggregates_enabled, set_rule_trace_enabled,
24
    },
25
    solver::Solver,
26
};
27
use conjure_cp::{
28
    parse::conjure_json::model_from_json, rule_engine::get_rules, settings::SolverFamily,
29
};
30
use conjure_cp::{parse::tree_sitter::parse_essence_file_native, solver::adaptors::*};
31
use conjure_cp_cli::find_conjure::conjure_executable;
32
use conjure_cp_cli::utils::conjure::{get_solutions, solutions_to_essence, solutions_to_json};
33
use conjure_cp_cli::utils::simplified_json::{
34
    domains_from_model, param_model_from_assignments, params_from_simplified_json_str,
35
    solution_to_simplified_json, solutions_to_simplified_json_string,
36
};
37

            
38
use crate::cli::{GlobalArgs, LOGGING_HELP_HEADING};
39

            
40
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, clap::ValueEnum)]
41
pub enum OutputFormat {
42
    /// Conjure-compatible Essence multi-solution / per-solution files
43
    #[default]
44
    Essence,
45
    /// Conjure `--output-format=json` simplified JSON
46
    Json,
47
}
48

            
49
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
50
pub enum NumberOfSolutions {
51
    All,
52
    Limit(i32),
53
}
54

            
55
impl NumberOfSolutions {
56
104
    fn as_solver_limit(self) -> i32 {
57
104
        match self {
58
46
            NumberOfSolutions::All => 0,
59
58
            NumberOfSolutions::Limit(limit) => limit,
60
        }
61
104
    }
62
}
63

            
64
197
fn parse_number_of_solutions(input: &str) -> Result<NumberOfSolutions, String> {
65
197
    if input.eq_ignore_ascii_case("all") {
66
46
        return Ok(NumberOfSolutions::All);
67
151
    }
68

            
69
151
    let limit = input
70
151
        .parse::<i32>()
71
        .map_err(|_| "expected a positive integer or 'all'".to_string())?;
72

            
73
151
    if limit <= 0 {
74
        return Err("expected a positive integer or 'all'".to_string());
75
151
    }
76

            
77
151
    Ok(NumberOfSolutions::Limit(limit))
78
197
}
79

            
80
#[derive(Clone, Debug, clap::Args)]
81
pub struct Args {
82
    /// The input Essence problem file
83
    #[arg(value_name = "INPUT_ESSENCE", value_hint = ValueHint::FilePath)]
84
    pub essence_file: PathBuf,
85

            
86
    /// The input Essence parameter file
87
    #[arg(value_name = "PARAM_ESSENCE", value_hint = ValueHint::FilePath)]
88
    pub param_file: Option<PathBuf>,
89

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

            
94
    /// Do not run the solver.
95
    ///
96
    /// The rewritten model is printed to stdout in an Essence-style syntax
97
    /// (but is not necessarily valid Essence).
98
    #[arg(long, default_value_t = false)]
99
    pub no_run_solver: bool,
100

            
101
    /// Number of solutions to return. Use a positive integer, or `all`.
102
    #[arg(
103
        long,
104
        short = 'n',
105
        default_value = "1",
106
        value_name = "N|all",
107
        value_parser = parse_number_of_solutions
108
    )]
109
    pub number_of_solutions: NumberOfSolutions,
110

            
111
    /// Format for printed / saved solutions
112
    #[arg(long, value_enum, default_value_t = OutputFormat::Essence, help_heading=LOGGING_HELP_HEADING)]
113
    pub output_format: OutputFormat,
114

            
115
    /// Write all solutions into a single file (default: true).
116
    ///
117
    /// Pass `--solutions-in-one-file=false` to emit one file (or stdout record) per solution.
118
    #[arg(
119
        long,
120
        default_value_t = true,
121
        action = clap::ArgAction::Set,
122
        value_name = "BOOL",
123
        help_heading = LOGGING_HELP_HEADING
124
    )]
125
    pub solutions_in_one_file: bool,
126

            
127
    /// Save solutions to the given path (Essence `.solutions` / JSON `.solutions.json`, or a
128
    /// filename stem when `--solutions-in-one-file=false`)
129
    #[arg(long, short = 'o', value_hint = ValueHint::FilePath,help_heading=LOGGING_HELP_HEADING)]
130
    pub output: Option<PathBuf>,
131

            
132
    /// When optimising, retain every improving solution Minion reports instead of only the last
133
    /// one.
134
    #[arg(long, default_value_t = false)]
135
    pub keep_intermediate_solutions: bool,
136
}
137

            
138
182
pub fn run_solve_command(global_args: GlobalArgs, solve_args: Args) -> anyhow::Result<()> {
139
182
    let essence_file = solve_args.essence_file.clone();
140
182
    let param_file = solve_args.param_file.clone();
141

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

            
145
182
    let context = init_context(&global_args, essence_file, param_file)?;
146

            
147
182
    let ctx_lock = context.read().unwrap();
148
182
    let essence_file_name = ctx_lock
149
182
        .essence_file_name
150
182
        .as_ref()
151
182
        .expect("context should contain the problem input file");
152
182
    let param_file_name = ctx_lock.param_file_name.as_ref();
153

            
154
    // parse models
155
182
    let problem_model = parse(&global_args, Arc::clone(&context), essence_file_name)?;
156

            
157
    // unify models
158
154
    let unified_model = match param_file_name {
159
        Some(param_file_name) => {
160
34
            let param_model = parse_param(
161
34
                &global_args,
162
34
                Arc::clone(&context),
163
34
                param_file_name,
164
34
                &problem_model,
165
34
            )?;
166
34
            instantiate_model(problem_model, param_model)?
167
        }
168
        None => {
169
120
            let mut problem_model = problem_model;
170
120
            validate_instantiation_conditions(&mut problem_model)?;
171
120
            problem_model
172
        }
173
    };
174
148
    drop(ctx_lock);
175

            
176
148
    let rewritten_model = rewrite(unified_model, &global_args, Arc::clone(&context))?;
177

            
178
148
    let solver = init_solver(&global_args)?;
179

            
180
148
    if solve_args.no_run_solver {
181
44
        println!("{}", rewritten_model);
182

            
183
44
        if let Some(path) = global_args.save_solver_input_file {
184
            let solver = solver.load_model(rewritten_model)?;
185
            tracing::info!(
186
                target: "conjure::stage",
187
                path = %path.display(),
188
                "writing solver input file"
189
            );
190
            let mut file: Box<dyn std::io::Write> = Box::new(File::create(path)?);
191
            solver.write_solver_input_file(&mut file)?;
192
44
        }
193
    } else {
194
104
        run_solver(solver, &global_args, &solve_args, rewritten_model)?
195
    }
196

            
197
    // still do postamble even if we didn't run the solver
198
148
    if let Some(ref path) = solve_args.info_json_path {
199
        let context_obj = context.read().unwrap().clone();
200
        let generated_json = &serde_json::to_value(context_obj)?;
201
        let pretty_json = serde_json::to_string_pretty(&generated_json)?;
202
        std::fs::write(path, format!("{pretty_json}\n"))?;
203
148
    }
204
148
    Ok(())
205
182
}
206

            
207
/// Returns a new Context and Solver for solving.
208
314
pub(crate) fn init_context(
209
314
    global_args: &GlobalArgs,
210
314
    essence_file: PathBuf,
211
314
    param_file: Option<PathBuf>,
212
314
) -> anyhow::Result<Arc<RwLock<Context<'static>>>> {
213
314
    let default_rule_trace_enabled = global_args.rule_trace.is_some();
214
314
    let rule_attempt_trace_enabled = global_args.rule_attempt_trace.is_some();
215
314
    let rule_trace_aggregates_enabled = global_args.rule_trace_aggregates.is_some();
216
    let rule_trace_enabled =
217
314
        default_rule_trace_enabled || rule_attempt_trace_enabled || rule_trace_aggregates_enabled;
218

            
219
314
    set_current_parser(global_args.parser);
220
314
    set_current_rewriter(global_args.rewriter);
221
314
    set_comprehension_expander(global_args.comprehension_expander);
222
314
    set_current_solver_family(global_args.solver);
223
314
    set_minion_discrete_threshold(global_args.minion_discrete_threshold);
224
314
    set_rule_trace_enabled(rule_trace_enabled);
225
314
    set_default_rule_trace_enabled(default_rule_trace_enabled);
226
314
    set_rule_attempt_trace_enabled(rule_attempt_trace_enabled);
227
314
    set_rule_trace_aggregates_enabled(rule_trace_aggregates_enabled);
228

            
229
314
    let target_family = global_args.solver;
230
314
    let mut extra_rule_sets: Vec<&str> = DEFAULT_RULE_SETS.to_vec();
231
314
    for rs in &global_args.extra_rule_sets {
232
        extra_rule_sets.push(rs.as_str());
233
    }
234

            
235
314
    let rule_sets = match resolve_rule_sets(target_family, &extra_rule_sets) {
236
314
        Ok(rs) => rs,
237
        Err(e) => {
238
            tracing::error!("Error resolving rule sets: {}", e);
239
            exit(1);
240
        }
241
    };
242

            
243
314
    tracing::info!(
244
        target: "conjure::stage",
245
314
        count = rule_sets.len(),
246
314
        "resolved rule sets"
247
    );
248
    tracing::debug!(
249
314
        rule_sets = %rule_sets
250
314
            .iter()
251
3454
            .map(|rule_set| rule_set.name)
252
            .collect::<Vec<_>>()
253
314
            .join(", "),
254
314
        "enabled rule sets"
255
    );
256

            
257
314
    let rules = get_rules(&rule_sets)?.into_iter().collect::<Vec<_>>();
258
    tracing::debug!(
259
        "Rules: {}",
260
314
        rules
261
314
            .iter()
262
83824
            .map(|rd| format!("{rd}"))
263
            .collect::<Vec<_>>()
264
314
            .join("\n")
265
    );
266
314
    let context = Context::new_ptr(
267
314
        target_family,
268
628
        extra_rule_sets.iter().map(|rs| rs.to_string()).collect(),
269
314
        rules,
270
314
        rule_sets.clone(),
271
    );
272

            
273
314
    context.write().unwrap().essence_file_name = Some(essence_file.to_str().expect("").into());
274
314
    if let Some(param_file) = param_file {
275
38
        context.write().unwrap().param_file_name = Some(param_file.to_str().expect("").into());
276
276
    }
277

            
278
314
    Ok(context)
279
314
}
280

            
281
150
pub(crate) fn init_solver(global_args: &GlobalArgs) -> anyhow::Result<Solver> {
282
150
    let family = global_args.solver;
283
150
    let timeout = global_args.solver_timeout.map(Duration::from);
284

            
285
150
    Ok(match family {
286
92
        SolverFamily::Minion => Solver::new(
287
92
            Minion::with_search_orders(global_args.minion_varorder, global_args.minion_valorder)
288
92
                .with_solver_seed(global_args.solver_seed)
289
92
                .with_timeout(timeout),
290
92
        ),
291
42
        SolverFamily::Sat => Solver::new(
292
42
            Sat::new(global_args.sat_backend)
293
42
                .with_solver_seed(global_args.solver_seed)
294
42
                .with_timeout(timeout),
295
42
        ),
296
        #[cfg(feature = "z3")]
297
        SolverFamily::Z3 => {
298
16
            Solver::new(Smt::new(timeout).with_solver_seed(global_args.solver_seed))
299
        }
300
        #[cfg(not(feature = "z3"))]
301
        SolverFamily::Z3 => {
302
            anyhow::bail!("Z3 solver support was not compiled in (enable the `z3` feature).")
303
        }
304
    })
305
150
}
306

            
307
352
pub(crate) fn parse(
308
352
    global_args: &GlobalArgs,
309
352
    context: Arc<RwLock<Context<'static>>>,
310
352
    file_path: &str,
311
352
) -> anyhow::Result<Model> {
312
352
    tracing::info!(target: "conjure::stage", path = %file_path, "parsing input model");
313

            
314
352
    match global_args.parser {
315
        conjure_cp::settings::Parser::TreeSitter => {
316
184
            parse_essence_file_native(file_path, context.clone()).map_err(|e| e.into())
317
        }
318
168
        conjure_cp::settings::Parser::ViaConjure => parse_with_conjure(file_path, context.clone()),
319
    }
320
352
}
321

            
322
/// Parse an Essence or simplified-JSON parameter file.
323
34
pub(crate) fn parse_param(
324
34
    global_args: &GlobalArgs,
325
34
    context: Arc<RwLock<Context<'static>>>,
326
34
    file_path: &str,
327
34
    problem_model: &Model,
328
34
) -> anyhow::Result<Model> {
329
34
    if file_path.ends_with(".json") {
330
        tracing::info!(target: "conjure::stage", path = %file_path, "parsing parameter file");
331
        let text = std::fs::read_to_string(file_path)?;
332
        let given_domains = domains_from_model(problem_model);
333
        let params = params_from_simplified_json_str(&text, &given_domains)?;
334
        return Ok(param_model_from_assignments(
335
            params,
336
            &given_domains,
337
            context,
338
        ));
339
34
    }
340
34
    parse(global_args, context, file_path)
341
34
}
342

            
343
168
pub(crate) fn parse_with_conjure(
344
168
    input_file: &str,
345
168
    context: Arc<RwLock<Context<'static>>>,
346
168
) -> anyhow::Result<Model> {
347
168
    conjure_executable().map_err(|e| anyhow!("Could not find correct conjure executable: {e}"))?;
348

            
349
168
    let mut cmd = std::process::Command::new("conjure");
350
168
    let output = cmd
351
168
        .arg("pretty")
352
168
        .arg("--output-format=astjson")
353
168
        .arg(input_file)
354
168
        .output()?;
355

            
356
168
    if !output.status.success() {
357
        println!("Parsing error: {}", String::from_utf8(output.stderr)?);
358
168
    }
359

            
360
168
    let astjson = String::from_utf8(output.stdout)?;
361

            
362
168
    if cfg!(feature = "extra-rule-checks") {
363
        tracing::debug!("extra-rule-checks: enabled");
364
    } else {
365
        tracing::debug!("extra-rule-checks: disabled");
366
    }
367

            
368
168
    model_from_json(&astjson, context.clone()).map_err(|e| anyhow!(e))
369
168
}
370

            
371
150
pub(crate) fn rewrite(
372
150
    model: Model,
373
150
    global_args: &GlobalArgs,
374
150
    context: Arc<RwLock<Context<'static>>>,
375
150
) -> anyhow::Result<Model> {
376
    tracing::debug!(model = %model, "initial model");
377

            
378
150
    let rewriter = global_args.rewriter;
379
150
    set_current_rewriter(rewriter);
380

            
381
150
    let comprehension_expander = global_args.comprehension_expander;
382
150
    set_comprehension_expander(comprehension_expander);
383
    tracing::debug!(%comprehension_expander, "configured comprehension expander");
384

            
385
150
    set_heuristic(global_args.heuristic);
386
150
    set_heuristic_seed(global_args.seed);
387
150
    set_heuristic_responses(global_args.responses.clone());
388
150
    set_channelling(global_args.channelling);
389
150
    conjure_cp::settings::set_amo_encoding(global_args.amo_encoding);
390
150
    conjure_cp::settings::set_pb_encoding(global_args.pb_encoding);
391
150
    conjure_cp::settings::set_element_encoding(global_args.element_encoding);
392
150
    conjure_cp::settings::set_table_encoding(global_args.table_encoding);
393
150
    conjure_cp::settings::set_alldifferent_encoding(global_args.alldifferent_encoding);
394
150
    conjure_cp::settings::set_cardinality_encoding(global_args.cardinality_encoding);
395
    tracing::debug!(
396
        "Heuristic: {}, seed: {}, responses: {:?}, channelling: {}, solver seed: {}",
397
150
        global_args.heuristic,
398
150
        global_args.seed,
399
150
        global_args.responses,
400
150
        global_args.channelling,
401
150
        global_args.solver_seed
402
    );
403

            
404
150
    let rule_sets = context.read().unwrap().rule_sets.clone();
405

            
406
150
    let Rewriter::Rewrite(config) = rewriter;
407
150
    tracing::info!(target: "conjure::stage", %config, "rewriting model");
408
150
    let new_model = rewrite_model(&model, &rule_sets, config)?;
409

            
410
    tracing::debug!(model = %new_model, "rewritten model");
411
150
    Ok(new_model)
412
150
}
413

            
414
104
fn run_solver(
415
104
    solver: Solver,
416
104
    global_args: &GlobalArgs,
417
104
    cmd_args: &Args,
418
104
    model: Model,
419
104
) -> anyhow::Result<()> {
420
104
    let domains = domains_from_model(&model);
421
104
    let solutions = get_solutions(
422
104
        solver,
423
104
        model,
424
104
        cmd_args.number_of_solutions.as_solver_limit(),
425
104
        cmd_args.keep_intermediate_solutions,
426
104
        &global_args.save_solver_input_file,
427
104
        global_args.rule_trace_cdp,
428
104
    )?;
429
    tracing::debug!(solutions = %solutions_to_json(&solutions), "solver solutions");
430

            
431
104
    let solutions = coerce_bools_in_solutions(&solutions, &domains);
432
104
    write_solutions(&solutions, cmd_args)?;
433
104
    Ok(())
434
104
}
435

            
436
/// Turn solver `0`/`1` assignments back into booleans when the find domain is bool, so JSON/Essence
437
/// output matches Conjure's simplified JSON (`true`/`false`).
438
104
fn coerce_bools_in_solutions(
439
104
    solutions: &[std::collections::BTreeMap<conjure_cp::ast::Name, conjure_cp::ast::Literal>],
440
104
    domains: &std::collections::BTreeMap<conjure_cp::ast::Name, conjure_cp::ast::DomainPtr>,
441
104
) -> Vec<std::collections::BTreeMap<conjure_cp::ast::Name, conjure_cp::ast::Literal>> {
442
    use conjure_cp::ast::{GroundDomain, Literal, Name};
443
104
    solutions
444
104
        .iter()
445
144
        .map(|solution| {
446
144
            solution
447
144
                .iter()
448
430
                .map(|(name, value)| {
449
430
                    let value = match domains.get(name).and_then(|domain| domain.as_ground()) {
450
240
                        Some(GroundDomain::Bool) => match value {
451
112
                            Literal::Int(1) => Literal::Bool(true),
452
116
                            Literal::Int(0) => Literal::Bool(false),
453
12
                            other => other.clone(),
454
                        },
455
190
                        _ => value.clone(),
456
                    };
457
430
                    (name.clone(), value)
458
430
                })
459
                .collect::<std::collections::BTreeMap<Name, Literal>>()
460
144
        })
461
        .collect()
462
104
}
463

            
464
104
fn write_solutions(
465
104
    solutions: &[std::collections::BTreeMap<conjure_cp::ast::Name, conjure_cp::ast::Literal>],
466
104
    cmd_args: &Args,
467
104
) -> anyhow::Result<()> {
468
104
    if cmd_args.solutions_in_one_file {
469
104
        let body = match cmd_args.output_format {
470
104
            OutputFormat::Essence => solutions_to_essence(solutions),
471
            OutputFormat::Json => solutions_to_simplified_json_string(solutions)?,
472
        };
473
104
        match &cmd_args.output {
474
104
            None => print!("{body}"),
475
            Some(path) => {
476
                std::fs::write(path, body.as_bytes())?;
477
                println!("Solutions saved to {:?}", path.canonicalize()?);
478
            }
479
        }
480
104
        return Ok(());
481
    }
482

            
483
    // One file / record per solution.
484
    for (index, solution) in solutions.iter().enumerate() {
485
        let body = match cmd_args.output_format {
486
            OutputFormat::Essence => solutions_to_essence(std::slice::from_ref(solution)),
487
            OutputFormat::Json => {
488
                let value = solution_to_simplified_json(solution)?;
489
                format!("{}\n", serde_json::to_string_pretty(&value)?)
490
            }
491
        };
492
        match &cmd_args.output {
493
            None => print!("{body}"),
494
            Some(path) => {
495
                let file_path = per_solution_output_path(path, index + 1, cmd_args.output_format);
496
                std::fs::write(&file_path, body.as_bytes())?;
497
                println!("Solution saved to {:?}", file_path.canonicalize()?);
498
            }
499
        }
500
    }
501
    Ok(())
502
104
}
503

            
504
fn per_solution_output_path(base: &Path, index: usize, format: OutputFormat) -> PathBuf {
505
    let stem = base
506
        .file_stem()
507
        .and_then(|s| s.to_str())
508
        .unwrap_or("solution");
509
    let parent = base.parent().unwrap_or_else(|| std::path::Path::new("."));
510
    let extension = match format {
511
        OutputFormat::Essence => "solution",
512
        OutputFormat::Json => "solution.json",
513
    };
514
    parent.join(format!("{stem}-{index:06}.{extension}"))
515
}