darklua_core/rules/
remove_compound_assign.rs

1use std::ops::{Deref, DerefMut};
2
3use crate::nodes::{
4    AssignStatement, BinaryExpression, Block, CompoundAssignStatement, DoStatement, Expression,
5    FieldExpression, IndexExpression, LocalAssignStatement, Prefix, Statement, Variable,
6};
7use crate::process::{DefaultVisitor, IdentifierTracker, NodeProcessor, NodeVisitor, ScopeVisitor};
8use crate::rules::{
9    Context, FlawlessRule, RuleConfiguration, RuleConfigurationError, RuleProperties,
10};
11
12use super::{verify_no_rule_properties, RemoveCommentProcessor, RemoveWhitespacesProcessor};
13
14struct Processor {
15    identifier_tracker: IdentifierTracker,
16    remove_comments: RemoveCommentProcessor,
17    remove_spaces: RemoveWhitespacesProcessor,
18}
19
20impl Processor {
21    #[inline]
22    fn generate_variable(&mut self) -> String {
23        self.identifier_tracker
24            .generate_identifier_with_prefix("__DARKLUA_VAR")
25    }
26
27    fn simplify_prefix(&self, prefix: &Prefix) -> Option<Prefix> {
28        match prefix {
29            Prefix::Parenthese(parenthese) => {
30                if let Expression::Identifier(identifier) = parenthese.inner_expression() {
31                    Some(Prefix::from(identifier.clone()))
32                } else {
33                    None
34                }
35            }
36            Prefix::Identifier(_) | Prefix::Call(_) | Prefix::Field(_) | Prefix::Index(_) => None,
37        }
38    }
39
40    fn remove_parentheses(&self, expression: impl Into<Expression>) -> Expression {
41        let expression = expression.into();
42        if let Expression::Parenthese(parenthese) = expression {
43            parenthese.into_inner_expression()
44        } else {
45            expression
46        }
47    }
48
49    fn replace_with(&mut self, assignment: &CompoundAssignStatement) -> Option<Statement> {
50        match assignment.get_variable() {
51            Variable::Index(index) => {
52                let prefix_assignment = match index.get_prefix() {
53                    Prefix::Identifier(_) => None,
54                    Prefix::Parenthese(parenthese)
55                        if matches!(
56                            parenthese.inner_expression(),
57                            Expression::False(_)
58                                | Expression::Identifier(_)
59                                | Expression::Number(_)
60                                | Expression::Nil(_)
61                                | Expression::String(_)
62                                | Expression::True(_)
63                                | Expression::VariableArguments(_)
64                        ) =>
65                    {
66                        None
67                    }
68                    Prefix::Call(_)
69                    | Prefix::Field(_)
70                    | Prefix::Index(_)
71                    | Prefix::Parenthese(_) => Some(self.generate_variable()),
72                };
73                let index_assignment = match index.get_index() {
74                    Expression::False(_)
75                    | Expression::Identifier(_)
76                    | Expression::Number(_)
77                    | Expression::Nil(_)
78                    | Expression::InterpolatedString(_)
79                    | Expression::String(_)
80                    | Expression::True(_)
81                    | Expression::VariableArguments(_) => None,
82                    Expression::Parenthese(parenthese)
83                        if matches!(
84                            parenthese.inner_expression(),
85                            Expression::False(_)
86                                | Expression::Identifier(_)
87                                | Expression::Number(_)
88                                | Expression::Nil(_)
89                                | Expression::String(_)
90                                | Expression::True(_)
91                                | Expression::VariableArguments(_)
92                        ) =>
93                    {
94                        None
95                    }
96                    Expression::Binary(_)
97                    | Expression::Call(_)
98                    | Expression::Field(_)
99                    | Expression::Function(_)
100                    | Expression::If(_)
101                    | Expression::Index(_)
102                    | Expression::Parenthese(_)
103                    | Expression::Table(_)
104                    | Expression::TypeCast(_)
105                    | Expression::Unary(_) => Some(self.generate_variable()),
106                };
107
108                match (prefix_assignment, index_assignment) {
109                    (None, None) => {
110                        let variable = IndexExpression::new(
111                            self.simplify_prefix(index.get_prefix())
112                                .unwrap_or_else(|| index.get_prefix().clone()),
113                            self.remove_parentheses(index.get_index().clone()),
114                        );
115
116                        Some(self.create_new_assignment_with_variable(
117                            assignment,
118                            variable.clone().into(),
119                            Some(variable.into()),
120                        ))
121                    }
122                    (None, Some(index_variable)) => {
123                        let assign = LocalAssignStatement::from_variable(index_variable.clone())
124                            .with_value(self.remove_parentheses(index.get_index().clone()));
125                        let variable = IndexExpression::new(
126                            self.simplify_prefix(index.get_prefix())
127                                .unwrap_or_else(|| index.get_prefix().clone()),
128                            Expression::identifier(index_variable),
129                        );
130                        Some(self.create_do_assignment(assignment, assign, variable))
131                    }
132                    (Some(prefix_variable), None) => {
133                        let assign = LocalAssignStatement::from_variable(prefix_variable.clone())
134                            .with_value(self.remove_parentheses(index.get_prefix().clone()));
135                        let variable = IndexExpression::new(
136                            Prefix::from_name(prefix_variable),
137                            index.get_index().clone(),
138                        );
139
140                        Some(self.create_do_assignment(assignment, assign, variable))
141                    }
142                    (Some(prefix_variable), Some(index_variable)) => {
143                        let assign = LocalAssignStatement::from_variable(prefix_variable.clone())
144                            .with_value(self.remove_parentheses(index.get_prefix().clone()))
145                            .with_variable(index_variable.clone())
146                            .with_value(self.remove_parentheses(index.get_index().clone()));
147                        let variable = IndexExpression::new(
148                            Prefix::from_name(prefix_variable),
149                            Expression::identifier(index_variable),
150                        );
151                        Some(self.create_do_assignment(assignment, assign, variable))
152                    }
153                }
154            }
155            Variable::Field(field) => match field.get_prefix() {
156                Prefix::Identifier(_) => None,
157                Prefix::Parenthese(parenthese)
158                    if matches!(
159                        parenthese.inner_expression(),
160                        Expression::False(_)
161                            | Expression::Identifier(_)
162                            | Expression::Number(_)
163                            | Expression::Nil(_)
164                            | Expression::String(_)
165                            | Expression::True(_)
166                            | Expression::VariableArguments(_)
167                    ) =>
168                {
169                    let new_prefix =
170                        if let Expression::Identifier(identifier) = parenthese.inner_expression() {
171                            Prefix::from(identifier.clone())
172                        } else {
173                            parenthese.clone().into()
174                        };
175                    let new_variable = FieldExpression::new(new_prefix, field.get_field().clone());
176
177                    Some(self.create_new_assignment_with_variable(
178                        assignment,
179                        new_variable.clone().into(),
180                        Some(new_variable.into()),
181                    ))
182                }
183                Prefix::Call(_) | Prefix::Field(_) | Prefix::Index(_) | Prefix::Parenthese(_) => {
184                    let identifier = self.generate_variable();
185                    let new_variable = FieldExpression::new(
186                        Prefix::from_name(&identifier),
187                        field.get_field().clone(),
188                    );
189
190                    let assign = LocalAssignStatement::from_variable(identifier).with_value(
191                        match field.get_prefix().clone() {
192                            Prefix::Parenthese(parenthese) => parenthese.into_inner_expression(),
193                            prefix => prefix.into(),
194                        },
195                    );
196
197                    Some(self.create_do_assignment(assignment, assign, new_variable))
198                }
199            },
200            Variable::Identifier(_) => None,
201        }
202    }
203
204    fn create_do_assignment(
205        &mut self,
206        compound_assignment: &CompoundAssignStatement,
207        assign: impl Into<Statement>,
208        variable: impl Into<Variable>,
209    ) -> Statement {
210        let variable = variable.into();
211        DoStatement::new(
212            Block::default()
213                .with_statement(assign.into())
214                .with_statement(self.create_new_assignment_with_variable(
215                    compound_assignment,
216                    variable.clone().into(),
217                    Some(variable),
218                )),
219        )
220        .into()
221    }
222
223    fn create_new_assignment(
224        &mut self,
225        assignment: &CompoundAssignStatement,
226        variable: impl Into<Expression>,
227    ) -> Statement {
228        self.create_new_assignment_with_variable(assignment, variable.into(), None)
229    }
230
231    fn create_new_assignment_with_variable(
232        &mut self,
233        assignment: &CompoundAssignStatement,
234        mut value: Expression,
235        variable: Option<Variable>,
236    ) -> Statement {
237        let operator = assignment.get_operator().to_binary_operator();
238
239        DefaultVisitor::visit_expression(&mut value, &mut self.remove_comments);
240        DefaultVisitor::visit_expression(&mut value, &mut self.remove_spaces);
241
242        let mut expression = BinaryExpression::new(operator, value, assignment.get_value().clone());
243        if let Some(token) = assignment.get_tokens().map(|tokens| {
244            let mut new_token = tokens.operator.clone();
245            new_token.replace_with_content(operator.to_str());
246            new_token
247        }) {
248            expression.set_token(token);
249        }
250
251        AssignStatement::from_variable(
252            variable.unwrap_or_else(|| assignment.get_variable().clone()),
253            expression,
254        )
255        .into()
256    }
257}
258
259impl Default for Processor {
260    fn default() -> Self {
261        Self {
262            identifier_tracker: IdentifierTracker::new(),
263            remove_comments: RemoveCommentProcessor::default(),
264            remove_spaces: RemoveWhitespacesProcessor::default(),
265        }
266    }
267}
268
269impl Deref for Processor {
270    type Target = IdentifierTracker;
271
272    fn deref(&self) -> &Self::Target {
273        &self.identifier_tracker
274    }
275}
276
277impl DerefMut for Processor {
278    fn deref_mut(&mut self) -> &mut Self::Target {
279        &mut self.identifier_tracker
280    }
281}
282
283impl NodeProcessor for Processor {
284    fn process_statement(&mut self, statement: &mut Statement) {
285        if let Statement::CompoundAssign(assignment) = statement {
286            let variable = assignment.get_variable();
287            *statement = self
288                .replace_with(assignment)
289                .unwrap_or_else(|| self.create_new_assignment(assignment, variable.clone()));
290        }
291    }
292}
293
294pub const REMOVE_COMPOUND_ASSIGNMENT_RULE_NAME: &str = "remove_compound_assignment";
295
296/// A rule that converts compound assignment (like `+=`) into regular assignments.
297#[derive(Debug, Default, PartialEq, Eq)]
298pub struct RemoveCompoundAssignment {}
299
300impl FlawlessRule for RemoveCompoundAssignment {
301    fn flawless_process(&self, block: &mut Block, _: &Context) {
302        let mut processor = Processor::default();
303        ScopeVisitor::visit_block(block, &mut processor);
304    }
305}
306
307impl RuleConfiguration for RemoveCompoundAssignment {
308    fn configure(&mut self, properties: RuleProperties) -> Result<(), RuleConfigurationError> {
309        verify_no_rule_properties(&properties)?;
310
311        Ok(())
312    }
313
314    fn get_name(&self) -> &'static str {
315        REMOVE_COMPOUND_ASSIGNMENT_RULE_NAME
316    }
317
318    fn serialize_to_properties(&self) -> RuleProperties {
319        RuleProperties::new()
320    }
321}
322
323#[cfg(test)]
324mod test {
325    use super::*;
326    use crate::rules::Rule;
327
328    use insta::assert_json_snapshot;
329
330    fn new_rule() -> RemoveCompoundAssignment {
331        RemoveCompoundAssignment::default()
332    }
333
334    #[test]
335    fn serialize_default_rule() {
336        let rule: Box<dyn Rule> = Box::new(new_rule());
337
338        assert_json_snapshot!("default_remove_compound_assignment", rule);
339    }
340
341    #[test]
342    fn configure_with_extra_field_error() {
343        let result = json5::from_str::<Box<dyn Rule>>(
344            r#"{
345            rule: 'remove_compound_assignment',
346            prop: "something",
347        }"#,
348        );
349        pretty_assertions::assert_eq!(result.unwrap_err().to_string(), "unexpected field 'prop'");
350    }
351}