1
use super::{RewriteError, RuleSet, resolve_rules::RuleData};
2
use crate::{
3
    Model,
4
    ast::{Expression as Expr, SubModel, comprehension::Comprehension},
5
    bug,
6
    rule_engine::{
7
        get_rules_grouped,
8
        rewriter_common::{RuleResult, log_rule_application},
9
        submodel_zipper::submodel_ctx,
10
    },
11
    stats::RewriterStats,
12
};
13

            
14
use itertools::Itertools;
15
use std::{process::exit, sync::Arc, time::Instant};
16
use tracing::{Level, span, trace};
17
use uniplate::{Biplate, Uniplate};
18

            
19
/// A naive, exhaustive rewriter for development purposes. Applies rules in priority order,
20
/// favouring expressions found earlier during preorder traversal of the tree.
21
pub fn rewrite_naive<'a>(
22
    model: &Model,
23
    rule_sets: &Vec<&'a RuleSet<'a>>,
24
    prop_multiple_equally_applicable: bool,
25
    exit_after_unrolling: bool,
26
) -> Result<Model, RewriteError> {
27
    let rules_grouped = get_rules_grouped(rule_sets)
28
        .unwrap_or_else(|_| bug!("get_rule_priorities() failed!"))
29
        .into_iter()
30
        .collect_vec();
31

            
32
    let mut model = model.clone();
33
    let mut done_something = true;
34

            
35
    let mut rewriter_stats = RewriterStats::new();
36
    rewriter_stats.is_optimization_enabled = Some(false);
37
    let run_start = Instant::now();
38

            
39
    trace!(
40
        target: "rule_engine_human",
41
        "Model before rewriting:\n\n{}\n--\n",
42
        model
43
    );
44

            
45
    // Rewrite until there are no more rules left to apply.
46
    while done_something {
47
        let mut new_model = None;
48
        done_something = false;
49

            
50
        // Rewrite each sub-model in the tree, largest first.
51
        for (mut submodel, ctx) in <_ as Biplate<SubModel>>::contexts_bi(&model) {
52
            if try_rewrite_model(
53
                &mut submodel,
54
                &rules_grouped,
55
                prop_multiple_equally_applicable,
56
                &mut rewriter_stats,
57
            )
58
            .is_some()
59
            {
60
                new_model = Some(ctx(submodel));
61
                done_something = true;
62
                break;
63
            }
64
        }
65
        if let Some(new_model) = new_model {
66
            model = new_model;
67
        }
68

            
69
        if Biplate::<Comprehension>::universe_bi(model.as_submodel()).is_empty()
70
            && exit_after_unrolling
71
        {
72
            println!("{}", model.as_submodel().root().universe().len());
73
            exit(0);
74
        }
75
    }
76

            
77
    let run_end = Instant::now();
78
    rewriter_stats.rewriter_run_time = Some(run_end - run_start);
79

            
80
    model
81
        .context
82
        .write()
83
        .unwrap()
84
        .stats
85
        .add_rewriter_run(rewriter_stats);
86

            
87
    trace!(
88
        target: "rule_engine_human",
89
        "Final model:\n\n{}",
90
        model
91
    );
92
    Ok(model)
93
}
94

            
95
// Tries to do a single rewrite on the model.
96
//
97
// Returns None if no change was made.
98
fn try_rewrite_model(
99
    submodel: &mut SubModel,
100
    rules_grouped: &Vec<(u16, Vec<RuleData<'_>>)>,
101
    prop_multiple_equally_applicable: bool,
102
    stats: &mut RewriterStats,
103
) -> Option<()> {
104
    type CtxFn = Arc<dyn Fn(Expr) -> SubModel>;
105
    let mut results: Vec<(RuleResult<'_>, u16, Expr, CtxFn)> = vec![];
106

            
107
    // Iterate over rules by priority in descending order.
108
    'top: for (priority, rules) in rules_grouped.iter() {
109
        // Using Biplate, rewrite both the expression tree, and any value lettings in the symbol
110
        // table.
111
        for (expr, ctx) in submodel_ctx(submodel.clone()) {
112
            // Clone expr and ctx so they can be reused
113
            let expr = expr.clone();
114
            let ctx = ctx.clone();
115
            for rd in rules {
116
                // Count rule application attempts
117
                stats.rewriter_rule_application_attempts =
118
                    Some(stats.rewriter_rule_application_attempts.unwrap_or(0) + 1);
119

            
120
                #[cfg(debug_assertions)]
121
                let span = span!(Level::TRACE,"trying_rule_application",rule_name=rd.rule.name,rule_target_expression=%expr);
122

            
123
                #[cfg(debug_assertions)]
124
                let _guard = span.enter();
125

            
126
                #[cfg(debug_assertions)]
127
                tracing::trace!(rule_name = rd.rule.name, "Trying rule");
128

            
129
                match (rd.rule.application)(&expr, &submodel.symbols()) {
130
                    Ok(red) => {
131
                        // Count successful rule applications
132
                        stats.rewriter_rule_applications =
133
                            Some(stats.rewriter_rule_applications.unwrap_or(0) + 1);
134

            
135
                        // Collect applicable rules
136
                        results.push((
137
                            RuleResult {
138
                                rule_data: rd.clone(),
139
                                reduction: red,
140
                            },
141
                            *priority,
142
                            expr.clone(),
143
                            ctx.clone(),
144
                        ));
145
                    }
146
                    Err(_) => {
147
                        // when called a lot, this becomes very expensive!
148
                        #[cfg(debug_assertions)]
149
                        tracing::trace!(
150
                            "Rule attempted but not applied: {} (priority {}, rule set {}), to expression: {}",
151
                            rd.rule.name,
152
                            priority,
153
                            rd.rule_set.name,
154
                            expr
155
                        );
156
                    }
157
                }
158
            }
159
            // This expression has the highest rule priority so far, so this is what we want to
160
            // rewrite.
161
            if !results.is_empty() {
162
                break 'top;
163
            }
164
        }
165
    }
166

            
167
    match results.as_slice() {
168
        [] => {
169
            return None;
170
        } // no rules are applicable.
171
        [(result, _priority, expr, ctx), ..] => {
172
            if prop_multiple_equally_applicable {
173
                assert_no_multiple_equally_applicable_rules(&results, rules_grouped);
174
            }
175

            
176
            // Extract the single applicable rule and apply it
177
            log_rule_application(result, expr, submodel);
178

            
179
            // Replace expr with new_expression
180
            *submodel = ctx(result.reduction.new_expression.clone());
181

            
182
            // Apply new symbols and top level
183
            result.reduction.clone().apply(submodel);
184
        }
185
    }
186

            
187
    Some(())
188
}
189

            
190
// Exits with a bug if there are multiple equally applicable rules for an expression.
191
fn assert_no_multiple_equally_applicable_rules<CtxFnType>(
192
    results: &Vec<(RuleResult<'_>, u16, Expr, CtxFnType)>,
193
    rules_grouped: &Vec<(u16, Vec<RuleData<'_>>)>,
194
) {
195
    if results.len() <= 1 {
196
        return;
197
    }
198

            
199
    let names: Vec<_> = results
200
        .iter()
201
        .map(|(result, _, _, _)| result.rule_data.rule.name)
202
        .collect();
203

            
204
    // Extract the expression from the first result
205
    let expr = results[0].2.clone();
206

            
207
    // Construct a single string to display the names of the rules grouped by priority
208
    let mut rules_by_priority_string = String::new();
209
    rules_by_priority_string.push_str("Rules grouped by priority:\n");
210
    for (priority, rules) in rules_grouped.iter() {
211
        rules_by_priority_string.push_str(&format!("Priority {priority}:\n"));
212
        for rd in rules {
213
            rules_by_priority_string.push_str(&format!(
214
                "  - {} (from {})\n",
215
                rd.rule.name, rd.rule_set.name
216
            ));
217
        }
218
    }
219
    bug!("Multiple equally applicable rules for {expr}: {names:#?}\n\n{rules_by_priority_string}");
220
}