1
//! Common utilities and types for rewriters.
2
use super::{
3
    Reduction,
4
    resolve_rules::{ResolveRulesError, RuleData},
5
    submodel_zipper::expression_ctx,
6
};
7
use crate::ast::{
8
    DeclarationPtr, Expression, Model, Name, SymbolTable,
9
    pretty::{pretty_variable_declaration, pretty_vec},
10
};
11
use crate::settings::{
12
    default_rule_trace_enabled, rule_trace_aggregates_enabled, rule_trace_enabled,
13
};
14

            
15
use itertools::Itertools;
16
use serde_json::json;
17
use std::collections::BTreeMap;
18
use std::fmt::Debug;
19
use std::sync::Arc;
20
use thiserror::Error;
21
use tracing::{info, trace};
22

            
23
#[derive(Debug, Clone)]
24
pub struct RuleResult<'a> {
25
    pub rule_data: RuleData<'a>,
26
    pub reduction: Reduction,
27
}
28

            
29
pub type VariableDeclarationSnapshot = BTreeMap<Name, String>;
30

            
31
14103492
pub fn snapshot_variable_declarations(symbols: &SymbolTable) -> VariableDeclarationSnapshot {
32
14103492
    symbols
33
14103492
        .clone()
34
14103492
        .into_iter_local()
35
1181514608
        .filter_map(|(name, _)| {
36
1181514608
            pretty_variable_declaration(symbols, &name).map(|declaration| (name, declaration))
37
1181514608
        })
38
14103492
        .collect()
39
14103492
}
40

            
41
/// Logs, to the main log, and the human readable traces used by the integration tester, that the
42
/// rule has been applied to the expression
43
634852
pub fn log_rule_application(
44
634852
    result: &RuleResult,
45
634852
    initial_expression: &Expression,
46
634852
    initial_symbols: &SymbolTable,
47
634852
    variable_declaration_snapshots: Option<(
48
634852
        &VariableDeclarationSnapshot,
49
634852
        &VariableDeclarationSnapshot,
50
634852
    )>,
51
634852
) {
52
634852
    let red = &result.reduction;
53
634852
    let rule = result.rule_data.rule;
54

            
55
    // A reduction can only modify either constraints or clauses, not both. So the the same
56
    // variable is used to hold changes in both (or empty if neither are changed).
57
634852
    let new_top_string = if !red.new_top.is_empty() {
58
56592
        pretty_vec(&red.new_top)
59
    } else {
60
578260
        pretty_vec(&red.new_clauses)
61
    };
62

            
63
634852
    info!(
64
        %new_top_string,
65
        "Applying rule: {} ({:?}), to expression: {}, resulting in: {}",
66
        rule.name,
67
        rule.rule_sets,
68
        initial_expression,
69
        red.new_expression
70
    );
71

            
72
634852
    if rule_trace_enabled() && default_rule_trace_enabled() {
73
321208
        let new_constraints_str = if !red.new_top.is_empty() {
74
26784
            let mut exprs: Vec<String> = vec![];
75
32328
            for expr in &red.new_top {
76
32328
                exprs.push(format!("  {expr}"));
77
32328
            }
78
26784
            let exprs = exprs.iter().join("\n");
79
26784
            format!("new constraints:\n{exprs}\n")
80
294424
        } else if !red.new_clauses.is_empty() {
81
87680
            let mut exprs: Vec<String> = vec![];
82
4679820
            for clause in &red.new_clauses {
83
4679820
                exprs.push(format!("  {clause}"));
84
4679820
            }
85
87680
            let exprs = exprs.iter().join("\n");
86
87680
            format!("new clauses:\n{exprs}\n")
87
        } else {
88
206744
            String::new()
89
        };
90

            
91
321208
        let (new_variables_str, updated_variables_str) =
92
321208
            if let Some((before, after)) = variable_declaration_snapshots {
93
94336
                let mut new_variables = Vec::new();
94
94336
                let mut updated_variables = Vec::new();
95

            
96
10198416
                for (name, declaration_after) in after {
97
10198416
                    match before.get(name) {
98
32616
                        None => new_variables.push(format!("  {declaration_after}")),
99
10165800
                        Some(declaration_before) if declaration_before != declaration_after => {
100
320
                            updated_variables
101
320
                                .push(format!("  {declaration_before} ~~> {declaration_after}"));
102
320
                        }
103
10165480
                        _ => {}
104
                    }
105
                }
106

            
107
94336
                let new_variables_str = if new_variables.is_empty() {
108
89732
                    String::new()
109
                } else {
110
4604
                    format!("new variables:\n{}\n", new_variables.join("\n"))
111
                };
112

            
113
94336
                let updated_variables_str = if updated_variables.is_empty() {
114
94056
                    String::new()
115
                } else {
116
280
                    format!("\nupdated variables:\n{}\n", updated_variables.join("\n"))
117
                };
118

            
119
94336
                (new_variables_str, updated_variables_str)
120
            } else {
121
                // empty if no new variables
122
226872
                let mut vars: Vec<String> = vec![];
123
1895032
                for var_name in red.added_symbols(initial_symbols) {
124
1894648
                    #[allow(clippy::unwrap_used)]
125
1894648
                    vars.push(format!(
126
1894648
                        "  {}",
127
1894648
                        pretty_variable_declaration(&red.symbols, &var_name).unwrap()
128
1894648
                    ));
129
1894648
                }
130
226872
                let new_variables_str = if vars.is_empty() {
131
151208
                    String::new()
132
                } else {
133
75664
                    format!("new variables:\n{}\n", vars.join("\n"))
134
                };
135
226872
                (new_variables_str, String::new())
136
            };
137

            
138
321208
        trace!(
139
            target: "rule_engine_rule_trace",
140
            "{}, \n   ~~> {} ({:?})\n{}\n{}{}{}\n--\n",
141
            initial_expression,
142
            rule.name,
143
            rule.rule_sets,
144
            red.new_expression,
145
            new_variables_str,
146
            updated_variables_str,
147
            new_constraints_str
148
        );
149
313644
    }
150

            
151
634852
    if rule_trace_enabled() && rule_trace_aggregates_enabled() {
152
3200
        trace!(
153
            target: "rule_engine_rule_trace_aggregates",
154
            rule_name = rule.name,
155
            "Applied rule"
156
        );
157
631652
    }
158

            
159
634852
    trace!(
160
        target: "rule_engine",
161
        "{}",
162
4428
    json!({
163
4428
        "rule_name": result.rule_data.rule.name,
164
4428
        "rule_priority": result.rule_data.priority,
165
4428
        "rule_set": {
166
4428
            "name": result.rule_data.rule_set.name,
167
        },
168
4428
        "initial_expression": serde_json::to_value(initial_expression).unwrap(),
169
4428
        "transformed_expression": serde_json::to_value(&red.new_expression).unwrap()
170
    })
171

            
172
    )
173
634852
}
174

            
175
type LettingCtxFn = Arc<dyn Fn(Expression) -> Expression>;
176
type ApplicableLettingRule<'a> = (
177
    RuleResult<'a>,
178
    u16,
179
    Expression,
180
    DeclarationPtr,
181
    LettingCtxFn,
182
);
183

            
184
457465
pub(crate) fn try_rewrite_value_letting_once(
185
457465
    model: &mut Model,
186
457465
    rules_grouped: &Vec<(u16, Vec<RuleData<'_>>)>,
187
457465
    prop_multiple_equally_applicable: bool,
188
457465
) -> Option<()> {
189
457465
    let symbols = model.symbols().clone();
190
457465
    let mut results: Vec<ApplicableLettingRule<'_>> = vec![];
191

            
192
8319361
    'top: for (priority, rules) in rules_grouped.iter() {
193
630480756
        for (_, decl) in symbols.clone().into_iter_local() {
194
630480756
            let Some(letting_expr) = decl.as_value_letting().map(|expr| expr.clone()) else {
195
628538048
                continue;
196
            };
197

            
198
3012280
            for (expr, ctx) in expression_ctx(letting_expr) {
199
3012280
                let expr = expr.clone();
200
3012280
                let ctx = ctx.clone();
201

            
202
12646776
                for rd in rules {
203
12646776
                    let Ok(reduction) = (rd.rule.application)(&expr, &symbols) else {
204
12645900
                        continue;
205
                    };
206

            
207
876
                    results.push((
208
876
                        RuleResult {
209
876
                            rule_data: rd.clone(),
210
876
                            reduction,
211
876
                        },
212
876
                        *priority,
213
876
                        expr.clone(),
214
876
                        decl.clone(),
215
876
                        ctx.clone(),
216
876
                    ));
217
                }
218

            
219
3012280
                if !results.is_empty() {
220
876
                    break 'top;
221
3011404
                }
222
            }
223
        }
224
    }
225

            
226
457465
    let (result, _, expr, decl, ctx) = match results.as_slice() {
227
457465
        [] => return None,
228
876
        [single, ..] => single,
229
    };
230

            
231
876
    if prop_multiple_equally_applicable && results.len() > 1 {
232
        let names: Vec<_> = results
233
            .iter()
234
            .map(|(result, _, _, _, _)| result.rule_data.rule.name)
235
            .collect();
236
        panic!("Multiple equally applicable rules for value letting expression {expr}: {names:?}");
237
876
    }
238

            
239
876
    log_rule_application(result, expr, &symbols, None);
240

            
241
876
    let rewritten_expr = ctx(result.reduction.new_expression.clone());
242
876
    result.reduction.clone().apply(model);
243

            
244
876
    let mut decl = decl.clone();
245
876
    *decl
246
876
        .as_value_letting_mut()
247
876
        .expect("declaration should still be a value letting") = rewritten_expr;
248

            
249
876
    Some(())
250
457465
}
251

            
252
/// Represents errors that can occur during the model rewriting process.
253
#[derive(Debug, Error)]
254
pub enum RewriteError {
255
    #[error("Error resolving rules {0}")]
256
    ResolveRulesError(ResolveRulesError),
257
}
258

            
259
impl From<ResolveRulesError> for RewriteError {
260
    fn from(error: ResolveRulesError) -> Self {
261
        RewriteError::ResolveRulesError(error)
262
    }
263
}