1
use std::collections::HashMap;
2
use std::env;
3

            
4
use crate::ast::ReturnType;
5
use crate::bug;
6
use crate::stats::RewriterStats;
7
use uniplate::Uniplate;
8

            
9
use crate::rule_engine::{Reduction, Rule, RuleSet};
10
use crate::{
11
    ast::Expression,
12
    rule_engine::resolve_rules::{get_rule_priorities, get_rules_vec},
13
    Model,
14
};
15

            
16
use super::rewriter_common::{log_rule_application, RewriteError, RuleResult};
17

            
18
/// Checks if the OPTIMIZATIONS environment variable is set to "1".
19
///
20
/// # Returns
21
/// - true if the environment variable is set to "1".
22
/// - false if the environment variable is not set or set to any other value.
23
68
fn optimizations_enabled() -> bool {
24
68
    match env::var("OPTIMIZATIONS") {
25
        Ok(val) => val == "1",
26
68
        Err(_) => false, // Assume optimizations are disabled if the environment variable is not set
27
    }
28
68
}
29

            
30
/// Rewrites the given model by applying a set of rules to all its constraints.
31
///
32
/// This function iteratively applies transformations to the model's constraints using the specified rule sets.
33
/// It returns a modified version of the model with all applicable rules applied, ensuring that any side-effects
34
/// such as updates to the symbol table and top-level constraints are properly reflected in the returned model.
35
///
36
/// # Parameters
37
/// - `model`: A reference to the [`Model`] to be rewritten. The function will clone this model to produce a modified version.
38
/// - `rule_sets`: A vector of references to [`RuleSet`]s that define the rules to be applied to the model's constraints.
39
///   Each `RuleSet` is expected to contain a collection of rules that can transform one or more constraints
40
///   within the model. The lifetime parameter `'a` ensures that the rules' references are valid for the
41
///   duration of the function execution.
42
///
43
/// # Returns
44
/// - `Ok(Model)`: If successful, it returns a modified copy of the [`Model`] after all applicable rules have been
45
///   applied. This new model includes any side-effects such as updates to the symbol table or modifications
46
///   to the constraints.
47
/// - `Err(RewriteError)`: If an error occurs during rule application (e.g., invalid rules or failed constraints),
48
///   it returns a [`RewriteError`] with details about the failure.
49
///
50
/// # Side-Effects
51
/// - When the model is rewritten, related data structures such as the symbol table (which tracks variable names and types)
52
///   or other top-level constraints may also be updated to reflect these changes. These updates are applied to the returned model,
53
///   ensuring that all related components stay consistent and aligned with the changes made during the rewrite.
54
/// - The function collects statistics about the rewriting process, including the number of rule applications
55
///   and the total runtime of the rewriter. These statistics are then stored in the model's context for
56
///   performance monitoring and analysis.
57
///
58
/// # Example
59
/// - Using `rewrite_model` with the Expression `a + min(x, y)`
60
///
61
///   Initial expression: a + min(x, y)
62
///   A model containing the expression is created. The variables of the model are represented by a SymbolTable and contain a,x,y.
63
///   The contraints of the initail model is the expression itself.
64
///
65
///   After getting the rules by their priorities and getting additional statistics the while loop of single interations is executed.
66
///   Details for this process can be found in [`rewrite_iteration`] documentation.
67
///
68
///   The loop is exited only when no more rules can be applied, when rewrite_iteration returns None and [`while let Some(step) = None`] occurs
69
///
70
///
71
///   Will result in side effects ((d<=x ^ d<=y) being the [`new_top`] and the model will now be a conjuction of that and (a+d)
72
///   Rewritten expression: ((a + d) ^ (d<=x ^ d<=y))
73
///
74
/// # Performance Considerations
75
/// - The function checks if optimizations are enabled before applying rules, which may affect the performance
76
///   of the rewriting process.
77
/// - Depending on the size of the model and the number of rules, the rewriting process might take a significant
78
///   amount of time. Use the statistics collected (`rewriter_run_time` and `rewriter_rule_application_attempts`)
79
///   to monitor and optimize performance.
80
///
81
/// # Panics
82
/// - This function may panic if the model's context is unavailable or if there is an issue with locking the context.
83
///
84
/// # See Also
85
/// - [`get_rule_priorities`]: Retrieves the priorities for the given rules.
86
/// - [`rewrite_iteration`]: Executes a single iteration of rewriting the model using the specified rules.
87
34
pub fn rewrite_model<'a>(
88
34
    model: &Model,
89
34
    rule_sets: &Vec<&'a RuleSet<'a>>,
90
34
) -> Result<Model, RewriteError> {
91
34
    let rule_priorities = get_rule_priorities(rule_sets)?;
92
34
    let rules = get_rules_vec(&rule_priorities);
93
34
    let mut new_model = model.clone();
94
34
    let mut stats = RewriterStats {
95
34
        is_optimization_enabled: Some(optimizations_enabled()),
96
34
        rewriter_run_time: None,
97
34
        rewriter_rule_application_attempts: Some(0),
98
34
        rewriter_rule_applications: Some(0),
99
34
    };
100
34

            
101
34
    // Check if optimizations are enabled
102
34
    let apply_optimizations = optimizations_enabled();
103
34

            
104
34
    let start = std::time::Instant::now();
105
34

            
106
34
    //the while loop is exited when None is returned implying the sub-expression is clean
107
34
    let mut i: usize = 0;
108
68
    while i < new_model.constraints.len() {
109
102
        while let Some(step) = rewrite_iteration(
110
102
            &new_model.constraints[i],
111
102
            &new_model,
112
102
            &rules,
113
102
            apply_optimizations,
114
102
            &mut stats,
115
102
        ) {
116
68
            debug_assert!(is_vec_bool(&step.new_top)); // All new_top expressions should be boolean
117
68
            new_model.constraints[i] = step.new_expression.clone();
118
68
            step.apply(&mut new_model); // Apply side-effects (e.g., symbol table updates)
119
        }
120

            
121
        // If new constraints are added, continue processing them in the next iterations.
122
34
        i += 1;
123
    }
124

            
125
34
    stats.rewriter_run_time = Some(start.elapsed());
126
34
    model.context.write().unwrap().stats.add_rewriter_run(stats);
127
34
    Ok(new_model)
128
34
}
129

            
130
/// Checks if all expressions in `Vec<Expr>` are booleans. This needs to be true so
131
/// the vector could be conjuncted to the model.
132
/// # Returns
133
///
134
/// - true: all expressions are booleans (so can be conjuncted).
135
///
136
/// - false: not all expressions are booleans.
137
68
fn is_vec_bool(exprs: &[Expression]) -> bool {
138
68
    exprs
139
68
        .iter()
140
68
        .all(|expr| expr.return_type() == Some(ReturnType::Bool))
141
68
}
142

            
143
/// Attempts to apply a set of rules to the given expression and its sub-expressions in the model.
144
///
145
/// This function recursively traverses the provided expression, applying any applicable rules from the given set.
146
/// If a rule is successfully applied to the expression or any of its sub-expressions, it returns a `Reduction`
147
/// containing the new expression, modified top-level constraints, and any changes to symbols. If no rules can be
148
/// applied at any level, it returns `None`.
149
///
150
/// # Parameters
151
/// - `expression`: A reference to the [`Expression`] to be rewritten. This is the main expression that the function
152
///   attempts to modify using the given rules.
153
/// - `model`: A reference to the [`Model`] that provides context and additional constraints for evaluating the rules.
154
/// - `rules`: A vector of references to [`Rule`]s that define the transformations to apply to the expression.
155
/// - `apply_optimizations`: A boolean flag that indicates whether optimization checks should be applied during the rewriting process.
156
///   If `true`, the function skips already "clean" (fully optimized or processed) expressions and marks them accordingly
157
///   to avoid redundant work.
158
/// - `stats`: A mutable reference to [`RewriterStats`] to collect statistics about the rule application process, such as
159
///   the number of rules applied and the time taken for each iteration.
160
///
161
/// # Returns
162
/// - `Some(<Reduction>)`: A [`Reduction`] containing the new expression and any associated modifications if a rule was applied
163
///   to `expr` or one of its sub-expressions.
164
/// - `None`: If no rule is applicable to the expression or any of its sub-expressions.
165
///
166
/// # Side-Effects
167
/// - If `apply_optimizations` is enabled, the function will skip "clean" expressions and mark successfully rewritten
168
///   expressions as "dirty". This is done to avoid unnecessary recomputation of expressions that have already been
169
///   optimized or processed.
170
///
171
/// # Example
172
/// - Recursively applying [`rewrite_iteration`]  to [`a + min(x, y)`]
173
///
174
///   Initially [`if apply_optimizations && expression.is_clean()`] is not true yet since intially our expression is dirty.
175
///
176
///   [`apply_results`] returns a null vector since no rules can be applied at the top level.
177
///   After calling function [`children`] on the expression a vector of sub-expression [`[a, min(x, y)]`] is returned.
178
///
179
///   The function iterates through the vector of the children from the top expression and calls itself.
180
///
181
///   [rewrite_iteration] on on the child [`a`] returns None, but on [`min(x, y)`] returns a [`Reduction`] object [`red`].
182
///   In this case, a rule (min simplification) can apply:
183
///   - d is added to the SymbolTable and the variables field is updated in the model. new_top is the side effects: (d<=x ^ d<=y)
184
///   - [`red = Reduction::new(new_expression = d, new_top, symbols)`];
185
///   - [`sub[1] = red.new_expression`] - Updates the second element in the vector of sub-expressions from [`min(x, y)`] to [`d`]
186
///
187
///   Since a child expression [`min(x, y)`] was rewritten to d, the parent expression [`a + min(x, y)`] is updated with the new child [`a+d`].
188
///   New [`Reduction`] is returned containing the modifications
189
///
190
///   The condition [`Some(step) = Some(new reduction)`] in the while loop in [`rewrite_model`] is met -> side effects are applied.
191
///
192
///   No more rules in our example can apply to the modified model -> mark all the children as clean and return a pure [`Reduction`].
193
///   [`return Some(Reduction::pure(expression))`]
194
///
195
///   On the last execution of rewrite_iteration condition [`apply_optimizations && expression.is_clean()`] is met, [`None`] is returned.
196
///
197
///
198
/// # Notes
199
/// - This function works recursively, meaning it traverses all sub-expressions within the given `expression` to find the
200
///   first rule that can be applied. If a rule is applied, it immediately returns the modified expression and stops
201
///   further traversal for that branch.
202
306
fn rewrite_iteration<'a>(
203
306
    expression: &'a Expression,
204
306
    model: &'a Model,
205
306
    rules: &'a Vec<&'a Rule<'a>>,
206
306
    apply_optimizations: bool,
207
306
    stats: &mut RewriterStats,
208
306
) -> Option<Reduction> {
209
306
    if apply_optimizations && expression.is_clean() {
210
        // Skip processing this expression if it's clean
211
        return None;
212
306
    }
213
306

            
214
306
    // Mark the expression as clean - will be marked dirty if any rule is applied
215
306
    let mut expression = expression.clone();
216
306

            
217
306
    let rule_results = apply_all_rules(&expression, model, rules, stats);
218
306
    if let Some(result) = choose_rewrite(&rule_results, &expression) {
219
        // If a rule is applied, mark the expression as dirty
220
68
        log_rule_application(&result, &expression, model);
221
68
        return Some(result.reduction);
222
238
    }
223
238

            
224
238
    let mut sub = expression.children();
225
238
    for i in 0..sub.len() {
226
204
        if let Some(red) = rewrite_iteration(&sub[i], model, rules, apply_optimizations, stats) {
227
51
            sub[i] = red.new_expression;
228
51
            let res = expression.with_children(sub.clone());
229
51
            return Some(Reduction::new(res, red.new_top, red.symbols));
230
153
        }
231
    }
232
    // If all children are clean, mark this expression as clean
233
187
    if apply_optimizations {
234
        assert!(expression.children().iter().all(|c| c.is_clean()));
235
        expression.set_clean(true);
236
        return Some(Reduction::pure(expression));
237
187
    }
238
187
    None
239
306
}
240

            
241
/// Applies all the given rules to a specific expression within the model.
242
///
243
/// This function iterates through the provided rules and attempts to apply each rule to the given `expression`.
244
/// If a rule is successfully applied, it creates a [`RuleResult`] containing the original rule and the resulting
245
/// [`Reduction`]. The statistics (`stats`) are updated to reflect the number of rule application attempts and successful
246
/// applications.
247
///
248
/// The function does not modify the provided `expression` directly. Instead, it collects all applicable rule results
249
/// into a vector, which can then be used for further processing or selection (e.g., with [`choose_rewrite`]).
250
///
251
/// # Parameters
252
/// - `expression`: A reference to the [`Expression`] that will be evaluated against the given rules. This is the main
253
///   target for rule transformations and is expected to remain unchanged during the function execution.
254
/// - `model`: A reference to the [`Model`] that provides context for rule evaluation, such as constraints and symbols.
255
///   Rules may depend on information in the model to determine if they can be applied.
256
/// - `rules`: A vector of references to [`Rule`]s that define the transformations to be applied to the expression.
257
///   Each rule is applied independently, and all applicable rules are collected.
258
/// - `stats`: A mutable reference to [`RewriterStats`] used to track statistics about rule application, such as
259
///   the number of attempts and successful applications.
260
///
261
/// # Returns
262
/// - A `Vec<RuleResult>` containing all rule applications that were successful. Each element in the vector represents
263
///   a rule that was applied to the given `expression` along with the resulting transformation.
264
/// - An empty vector if no rules were applicable to the expression.
265
///
266
/// # Side-Effects
267
/// - The function updates the provided `stats` with the number of rule application attempts and successful applications.
268
/// - Debug or trace logging may be performed to track which rules were applicable or not for a given expression.
269
///
270
/// # Example
271
///
272
/// let applicable_rules = apply_all_rules(&expr, &model, &rules, &mut stats);
273
/// if !applicable_rules.is_empty() {
274
///     for result in applicable_rules {
275
///         println!("Rule applied: {:?}", result.rule);
276
///     }
277
/// }
278
///
279
///
280
/// # Notes
281
/// - This function does not modify the input `expression` or `model` directly. The returned `RuleResult` vector
282
///   provides information about successful transformations, allowing the caller to decide how to process them.
283
/// - The function performs independent rule applications. If rules have dependencies or should be applied in a
284
///   specific order, consider handling that logic outside of this function.
285
///
286
/// # See Also
287
/// - [`choose_rewrite`]: Chooses a single reduction from the rule results provided by `apply_all_rules`.
288
306
fn apply_all_rules<'a>(
289
306
    expression: &'a Expression,
290
306
    model: &'a Model,
291
306
    rules: &'a Vec<&'a Rule<'a>>,
292
306
    stats: &mut RewriterStats,
293
306
) -> Vec<RuleResult<'a>> {
294
306
    let mut results = Vec::new();
295
14688
    for rule in rules {
296
14382
        match rule.apply(expression, model) {
297
85
            Ok(red) => {
298
85
                stats.rewriter_rule_application_attempts =
299
85
                    Some(stats.rewriter_rule_application_attempts.unwrap() + 1);
300
85
                stats.rewriter_rule_applications =
301
85
                    Some(stats.rewriter_rule_applications.unwrap() + 1);
302
85
                // Assert no clean children
303
85
                // assert!(!red.new_expression.children().iter().any(|c| c.is_clean()), "Rule that caused assertion to fail: {:?}", rule.name);
304
85
                // assert!(!red.new_expression.children().iter().any(|c| c.children().iter().any(|c| c.is_clean())));
305
85
                results.push(RuleResult {
306
85
                    rule,
307
85
                    reduction: red,
308
85
                });
309
85
            }
310
            Err(_) => {
311
14297
                log::trace!(
312
                    "Rule attempted but not applied: {} ({:?}), to expression: {}",
313
                    rule.name,
314
                    rule.rule_sets,
315
                    expression
316
                );
317
14297
                stats.rewriter_rule_application_attempts =
318
14297
                    Some(stats.rewriter_rule_application_attempts.unwrap() + 1);
319
14297
                continue;
320
            }
321
        }
322
    }
323
306
    results
324
306
}
325

            
326
/// Chooses the first applicable rule result from a list of rule applications.
327
///
328
/// This function selects a reduction from the provided `RuleResult` list, prioritizing the first rule
329
/// that successfully transforms the expression. This strategy can be modified in the future to incorporate
330
/// more complex selection criteria, such as prioritizing rules based on cost, complexity, or other heuristic metrics.
331
///
332
/// The function also checks the priorities of all the applicable rules and detects if there are multiple rules of the same proirity
333
///
334
/// # Parameters
335
/// - `results`: A slice of [`RuleResult`] containing potential rule applications to be considered. Each element
336
///   represents a rule that was successfully applied to the expression, along with the resulting transformation.
337
/// -  `initial_expression`: [`Expression`] before the rule tranformation.
338
///
339
/// # Returns
340
/// - `Some(<Reduction>)`: Returns a [`Reduction`] representing the first rule's application if there is at least one
341
///   rule that produced a successful transformation.
342
/// - `None`: If no rule applications are available in the `results` slice (i.e., it is empty), it returns `None`.
343
///
344
/// # Example
345
///
346
/// let rule_results = vec![rule1_result, rule2_result];
347
/// if let Some(reduction) = choose_rewrite(&rule_results) {
348
/// Process the chosen reduction
349
/// }
350
///
351
306
fn choose_rewrite<'a>(
352
306
    results: &[RuleResult<'a>],
353
306
    initial_expression: &Expression,
354
306
) -> Option<RuleResult<'a>> {
355
306
    //in the case where multiple rules are applicable
356
306
    if results.len() > 1 {
357
17
        let expr = results[0].reduction.new_expression.clone();
358
34
        let rules: Vec<_> = results.iter().map(|result| &result.rule).collect();
359
17

            
360
17
        check_priority(rules.clone(), initial_expression, &expr);
361
289
    }
362

            
363
306
    if results.is_empty() {
364
238
        return None;
365
68
    }
366
68

            
367
68
    // Return the first result for now
368
68
    Some(results[0].clone())
369
306
}
370

            
371
/// Function filters all the applicable rules based on their priority.
372
/// In the case where there are multiple rules of the same prioriy, a bug! is thrown listing all those duplicates.
373
/// Otherwise, if there are multiple rules applicable but they all have different priorities, a warning message is dispalyed.
374
///
375
/// # Parameters
376
/// - `rules`: a vector of [`Rule`] containing all the applicable rules and their metadata for a specific expression.
377
/// - `initial_expression`: [`Expression`] before rule the tranformation.
378
/// - `new_expr`: [`Expression`] after the rule transformation.
379
///
380
17
fn check_priority<'a>(
381
17
    rules: Vec<&&Rule<'_>>,
382
17
    initial_expr: &'a Expression,
383
17
    new_expr: &'a Expression,
384
17
) {
385
17
    //getting the rule sets from the applicable rules
386
34
    let rule_sets: Vec<_> = rules.iter().map(|rule| &rule.rule_sets).collect();
387
17

            
388
17
    //a map with keys being rule priorities and their values neing all the rules of that priority found in the rule_sets
389
17
    let mut rules_by_priorities: HashMap<u16, Vec<&str>> = HashMap::new();
390

            
391
    //iterates over each rule_set and groups by the rule priority
392
51
    for rule_set in &rule_sets {
393
34
        if let Some((name, priority)) = rule_set.first() {
394
34
            rules_by_priorities
395
34
                .entry(*priority)
396
34
                .or_default()
397
34
                .push(*name);
398
34
        }
399
    }
400

            
401
    //filters the map, retaining only entries where there is more than 1 rule of the same priority
402
17
    let duplicate_rules: HashMap<u16, Vec<&str>> = rules_by_priorities
403
17
        .into_iter()
404
34
        .filter(|(_, group)| group.len() > 1)
405
17
        .collect();
406
17

            
407
17
    if !duplicate_rules.is_empty() {
408
        //accumulates all duplicates into a formatted message
409
        let mut message = format!("Found multiple rules of the same priority applicable to to expression: {} \n resulting in expression: {}", initial_expr, new_expr);
410
        for (priority, rules) in &duplicate_rules {
411
            message.push_str(&format!("Priority {:?} \n Rules: {:?}", priority, rules));
412
        }
413
        bug!("{}", message);
414

            
415
    //no duplicate rules of the same priorities were found in the set of applicable rules
416
    } else {
417
17
        log::warn!("Multiple rules of different priorities are applicable to expression {} \n resulting in expression: {}
418
        \n Rules{:?}", initial_expr, new_expr, rules)
419
    }
420
17
}