Repository: seaofvoices/darklua (rust)

Issue:
`remove_compound_assignment` copying comments/newlines in between statements
### How to replicate:

1. Go to https://darklua.com/try-it/
2. Use the following rule below:
```json5
{
  rules: [
    'remove_compound_assignment',
  ],
}
```
### Results:

<table>
<tr>
<th>Input</th>
<th>Output</th>
<th>Expected</th>
</tr>
<tr>
<td valign="top">
  
```lua
local x = 0
x += 1-- this is ok

local y = 0
print(y) 
y += 1-- this is ok


local y = 0

y += 1-- one newline appears

local z = 0
--duplicate
z += 1--it copies the comments

local a = 0
--copies
--all
--comments

--in between
a += 1
```
  
</td>
<td valign="top">
  
```lua
local x = 0
x =x +1-- this is ok

local y = 0
print(y) 
y =y +1-- this is ok


local y = 0

y =
y +1-- one newline appears

local z = 0
--duplicate
z =--duplicate
z +1--it copies the comments

local a = 0
--copies
--all
--comments

--in between
a =--copies
--all
--comments

--in between
a +1
```
  
</td>
<td valign="top">

```lua
local x = 0
x =x +1-- this is ok

local y = 0
print(y) 
y =y +1-- this is ok


local y = 0

y =y +1-- one newline appears

local z = 0
--duplicate
z =z +1--it copies the comments

local a = 0
--copies
--all
--comments

--in between
a =a +1
```

</td>
</tr>
</table>

Pull request description:
Fix `remove_compound_assignment` rule to avoid copying variable tokens
Closes #174 

- [x] add entry to the changelog


Patch:
diff --git a/.github/workflows/site.yml b/.github/workflows/site.yml
index 0840041..d54109c 100644
--- a/.github/workflows/site.yml
+++ b/.github/workflows/site.yml
@@ -11,7 +11,7 @@ on:
 jobs:
   build-site:
     name: Build darklua site
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-latest
     defaults:
       run:
         shell: bash
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 77920b6..09bdee3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+* fix `remove_compound_assignment` rule to avoid copying variable tokens ([#176](https://github.com/seaofvoices/darklua/pull/176))
 * add get link button to [try-it](https://darklua.com/try-it/) page ([#175](https://github.com/seaofvoices/darklua/pull/175))
 * add rule to remove unused variables (`remove_unused_variable`). Fix issue with `rename_variables` where `self` variables and some cases of variable shadowing were not correctly renamed ([#172](https://github.com/seaofvoices/darklua/pull/172))
 
diff --git a/src/rules/remove_comments.rs b/src/rules/remove_comments.rs
index d163e6d..568b805 100644
--- a/src/rules/remove_comments.rs
+++ b/src/rules/remove_comments.rs
@@ -7,7 +7,7 @@ use crate::rules::{
 use super::verify_no_rule_properties;
 
 #[derive(Debug, Default)]
-struct RemoveCommentProcessor {}
+pub(crate) struct RemoveCommentProcessor {}
 
 impl NodeProcessor for RemoveCommentProcessor {
     fn process_block(&mut self, block: &mut Block) {
diff --git a/src/rules/remove_compound_assign.rs b/src/rules/remove_compound_assign.rs
index d2524e8..4061aa0 100644
--- a/src/rules/remove_compound_assign.rs
+++ b/src/rules/remove_compound_assign.rs
@@ -4,15 +4,17 @@ use crate::nodes::{
     AssignStatement, BinaryExpression, Block, CompoundAssignStatement, DoStatement, Expression,
     FieldExpression, IndexExpression, LocalAssignStatement, Prefix, Statement, Variable,
 };
-use crate::process::{IdentifierTracker, NodeProcessor, NodeVisitor, ScopeVisitor};
+use crate::process::{DefaultVisitor, IdentifierTracker, NodeProcessor, NodeVisitor, ScopeVisitor};
 use crate::rules::{
     Context, FlawlessRule, RuleConfiguration, RuleConfigurationError, RuleProperties,
 };
 
-use super::verify_no_rule_properties;
+use super::{verify_no_rule_properties, RemoveCommentProcessor, RemoveWhitespacesProcessor};
 
 struct Processor {
     identifier_tracker: IdentifierTracker,
+    remove_comments: RemoveCommentProcessor,
+    remove_spaces: RemoveWhitespacesProcessor,
 }
 
 impl Processor {
@@ -111,17 +113,11 @@ impl Processor {
                             self.remove_parentheses(index.get_index().clone()),
                         );
 
-                        Some(
-                            AssignStatement::from_variable(
-                                variable.clone(),
-                                BinaryExpression::new(
-                                    assignment.get_operator().to_binary_operator(),
-                                    variable,
-                                    assignment.get_value().clone(),
-                                ),
-                            )
-                            .into(),
-                        )
+                        Some(self.create_new_assignment_with_variable(
+                            assignment,
+                            variable.clone().into(),
+                            Some(variable.into()),
+                        ))
                     }
                     (None, Some(index_variable)) => {
                         let assign = LocalAssignStatement::from_variable(index_variable.clone())
@@ -131,7 +127,7 @@ impl Processor {
                                 .unwrap_or_else(|| index.get_prefix().clone()),
                             Expression::identifier(index_variable),
                         );
-                        Some(into_do_assignment(assignment, assign, variable))
+                        Some(self.create_do_assignment(assignment, assign, variable))
                     }
                     (Some(prefix_variable), None) => {
                         let assign = LocalAssignStatement::from_variable(prefix_variable.clone())
@@ -141,7 +137,7 @@ impl Processor {
                             index.get_index().clone(),
                         );
 
-                        Some(into_do_assignment(assignment, assign, variable))
+                        Some(self.create_do_assignment(assignment, assign, variable))
                     }
                     (Some(prefix_variable), Some(index_variable)) => {
                         let assign = LocalAssignStatement::from_variable(prefix_variable.clone())
@@ -152,7 +148,7 @@ impl Processor {
                             Prefix::from_name(prefix_variable),
                             Expression::identifier(index_variable),
                         );
-                        Some(into_do_assignment(assignment, assign, variable))
+                        Some(self.create_do_assignment(assignment, assign, variable))
                     }
                 }
             }
@@ -178,17 +174,11 @@ impl Processor {
                         };
                     let new_variable = FieldExpression::new(new_prefix, field.get_field().clone());
 
-                    Some(
-                        AssignStatement::from_variable(
-                            new_variable.clone(),
-                            BinaryExpression::new(
-                                assignment.get_operator().to_binary_operator(),
-                                new_variable,
-                                assignment.get_value().clone(),
-                            ),
-                        )
-                        .into(),
-                    )
+                    Some(self.create_new_assignment_with_variable(
+                        assignment,
+                        new_variable.clone().into(),
+                        Some(new_variable.into()),
+                    ))
                 }
                 Prefix::Call(_) | Prefix::Field(_) | Prefix::Index(_) | Prefix::Parenthese(_) => {
                     let identifier = self.generate_variable();
@@ -196,41 +186,82 @@ impl Processor {
                         Prefix::from_name(&identifier),
                         field.get_field().clone(),
                     );
-                    Some(
-                        DoStatement::new(
-                            Block::default()
-                                .with_statement(
-                                    LocalAssignStatement::from_variable(identifier).with_value(
-                                        match field.get_prefix().clone() {
-                                            Prefix::Parenthese(parenthese) => {
-                                                parenthese.into_inner_expression()
-                                            }
-                                            prefix => prefix.into(),
-                                        },
-                                    ),
-                                )
-                                .with_statement(AssignStatement::from_variable(
-                                    new_variable.clone(),
-                                    BinaryExpression::new(
-                                        assignment.get_operator().to_binary_operator(),
-                                        new_variable,
-                                        assignment.get_value().clone(),
-                                    ),
-                                )),
-                        )
-                        .into(),
-                    )
+
+                    let assign = LocalAssignStatement::from_variable(identifier).with_value(
+                        match field.get_prefix().clone() {
+                            Prefix::Parenthese(parenthese) => parenthese.into_inner_expression(),
+                            prefix => prefix.into(),
+                        },
+                    );
+
+                    Some(self.create_do_assignment(assignment, assign, new_variable))
                 }
             },
             Variable::Identifier(_) => None,
         }
     }
+
+    fn create_do_assignment(
+        &mut self,
+        compound_assignment: &CompoundAssignStatement,
+        assign: impl Into<Statement>,
+        variable: impl Into<Variable>,
+    ) -> Statement {
+        let variable = variable.into();
+        DoStatement::new(
+            Block::default()
+                .with_statement(assign.into())
+                .with_statement(self.create_new_assignment_with_variable(
+                    compound_assignment,
+                    variable.clone().into(),
+                    Some(variable),
+                )),
+        )
+        .into()
+    }
+
+    fn create_new_assignment(
+        &mut self,
+        assignment: &CompoundAssignStatement,
+        variable: impl Into<Expression>,
+    ) -> Statement {
+        self.create_new_assignment_with_variable(assignment, variable.into(), None)
+    }
+
+    fn create_new_assignment_with_variable(
+        &mut self,
+        assignment: &CompoundAssignStatement,
+        mut value: Expression,
+        variable: Option<Variable>,
+    ) -> Statement {
+        let operator = assignment.get_operator().to_binary_operator();
+
+        DefaultVisitor::visit_expression(&mut value, &mut self.remove_comments);
+        DefaultVisitor::visit_expression(&mut value, &mut self.remove_spaces);
+
+        let mut expression = BinaryExpression::new(operator, value, assignment.get_value().clone());
+        if let Some(token) = assignment.get_tokens().map(|tokens| {
+            let mut new_token = tokens.operator.clone();
+            new_token.replace_with_content(operator.to_str());
+            new_token
+        }) {
+            expression.set_token(token);
+        }
+
+        AssignStatement::from_variable(
+            variable.unwrap_or_else(|| assignment.get_variable().clone()),
+            expression,
+        )
+        .into()
+    }
 }
 
 impl Default for Processor {
     fn default() -> Self {
         Self {
             identifier_tracker: IdentifierTracker::new(),
+            remove_comments: RemoveCommentProcessor::default(),
+            remove_spaces: RemoveWhitespacesProcessor::default(),
         }
     }
 }
@@ -253,42 +284,13 @@ impl NodeProcessor for Processor {
     fn process_statement(&mut self, statement: &mut Statement) {
         if let Statement::CompoundAssign(assignment) = statement {
             let variable = assignment.get_variable();
-            *statement = self.replace_with(assignment).unwrap_or_else(|| {
-                AssignStatement::from_variable(
-                    variable.clone(),
-                    BinaryExpression::new(
-                        assignment.get_operator().to_binary_operator(),
-                        variable.clone(),
-                        assignment.get_value().clone(),
-                    ),
-                )
-                .into()
-            });
+            *statement = self
+                .replace_with(assignment)
+                .unwrap_or_else(|| self.create_new_assignment(assignment, variable.clone()));
         }
     }
 }
 
-fn into_do_assignment(
-    compound_assignment: &CompoundAssignStatement,
-    assign: impl Into<Statement>,
-    variable: impl Into<Variable>,
-) -> Statement {
-    let variable = variable.into();
-    DoStatement::new(
-        Block::default()
-            .with_statement(assign.into())
-            .with_statement(AssignStatement::from_variable(
-                variable.clone(),
-                BinaryExpression::new(
-                    compound_assignment.get_operator().to_binary_operator(),
-                    variable,
-                    compound_assignment.get_value().clone(),
-                ),
-            )),
-    )
-    .into()
-}
-
 pub const REMOVE_COMPOUND_ASSIGNMENT_RULE_NAME: &str = "remove_compound_assignment";
 
 /// A rule that converts compound assignment (like `+=`) into regular assignments.
diff --git a/src/rules/remove_spaces.rs b/src/rules/remove_spaces.rs
index af3e671..200322a 100644
--- a/src/rules/remove_spaces.rs
+++ b/src/rules/remove_spaces.rs
@@ -7,7 +7,7 @@ use crate::rules::{
 use super::verify_no_rule_properties;
 
 #[derive(Debug, Default)]
-struct RemoveWhitespacesProcessor {}
+pub(crate) struct RemoveWhitespacesProcessor {}
 
 impl NodeProcessor for RemoveWhitespacesProcessor {
     fn process_block(&mut self, block: &mut Block) {


Test patch:
diff --git a/tests/rule_tests/remove_compound_assignment.rs b/tests/rule_tests/remove_compound_assignment.rs
index b7d844e..c6246d2 100644
--- a/tests/rule_tests/remove_compound_assignment.rs
+++ b/tests/rule_tests/remove_compound_assignment.rs
@@ -35,6 +35,15 @@ test_rule!(
         => "do local __DARKLUA_VAR = a.object __DARKLUA_VAR.counter = __DARKLUA_VAR.counter + 1 end do local __DARKLUA_VAR0 = b.object __DARKLUA_VAR0.counter = __DARKLUA_VAR0.counter - 1 end",
 );
 
+test_rule_with_tokens!(
+    remove_compound_assignment_with_tokens,
+    RemoveCompoundAssignment::default(),
+    trailing_comment("i += 1 -- comment") => "i =i+ 1 -- comment",
+    trailing_comment_on_second_line("\ni += 1 -- comment") => "\ni =i+ 1 -- comment",
+    comment_after_operator("i += --[[ comment ]] 1") => "i =i+ --[[ comment ]] 1",
+    comment_after_variable("i --[[ comment ]] += 1") => "i --[[ comment ]] =i+ 1",
+);
+
 #[test]
 fn deserialize_from_object_notation() {
     json5::from_str::<Box<dyn Rule>>(