Repository: seaofvoices/darklua (rust)

Issue:
`remove_assertions` changes return value when used in a call expression
For example, applying `remove_assertions` with the default parameter for the rule, the following input:
```lua
local head = assert(character:FindFirstChild("Head"), "head is missing!")
```
becomes:
```lua
local head = character:FindFirstChild("Head") and nil
```
This changes runtime behavior in an unexpected way, because `head` is always `nil` even when the character does have a child named Head. Expected output is:
```lua
local head = character:FindFirstChild("Head")
```
<hr>

When `preserve_aguments_side_effects` is set to `false`, the output is:
```lua
local head = nil
```
Not sure if this is intended either. I think whether the return value is used (so, whether the assert is an expression or a statement?) should probably be factored in here.

Tested on version 0.13.0.



Pull request description:
Fix `remove_assertions` rule
Closes #189 

Fix the `remove_assertions` rule to make sure that the `assert` calls return their arguments.

- [x] add entry to the changelog


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9242081..31da1c3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+* fix `remove_assertions` rule to make the `assert` calls return their arguments ([#229](https://github.com/seaofvoices/darklua/pull/229))
 * add rule to remove continue statements (`remove_continue`) ([#227](https://github.com/seaofvoices/darklua/pull/227))
 * fix negative zero sign erasure ([#222](https://github.com/seaofvoices/darklua/pull/222))
 * add `remove_if_expression` rule ([#221](https://github.com/seaofvoices/darklua/pull/221))
diff --git a/src/rules/remove_assertions.rs b/src/rules/remove_assertions.rs
index 6613c34..2836b7c 100644
--- a/src/rules/remove_assertions.rs
+++ b/src/rules/remove_assertions.rs
@@ -1,10 +1,13 @@
-use crate::nodes::{Block, Prefix};
+use std::collections::HashMap;
+use std::iter::{self, FromIterator};
+
+use crate::nodes::{Block, Expression, FunctionCall, Prefix, TupleArguments};
 use crate::process::{IdentifierTracker, NodeVisitor, ScopeVisitor};
 use crate::rules::{
     Context, FlawlessRule, RuleConfiguration, RuleConfigurationError, RuleProperties,
 };
 
-use super::remove_call_match::RemoveFunctionCallProcessor;
+use super::remove_call_match::{CallMatch, RemoveFunctionCallProcessor};
 
 const ASSERT_FUNCTION_NAME: &str = "assert";
 
@@ -24,22 +27,60 @@ impl Default for RemoveAssertions {
     }
 }
 
-fn should_remove_call(identifiers: &IdentifierTracker, prefix: &Prefix) -> bool {
-    if identifiers.is_identifier_used(ASSERT_FUNCTION_NAME) {
-        return false;
+struct AssertMatcher;
+
+impl CallMatch<()> for AssertMatcher {
+    fn matches(&self, identifiers: &IdentifierTracker, prefix: &Prefix) -> bool {
+        if identifiers.is_identifier_used(ASSERT_FUNCTION_NAME) {
+            return false;
+        }
+
+        match prefix {
+            Prefix::Identifier(identifier) => identifier.get_name() == ASSERT_FUNCTION_NAME,
+            _ => false,
+        }
     }
 
-    match prefix {
-        Prefix::Identifier(identifier) => identifier.get_name() == ASSERT_FUNCTION_NAME,
-        _ => false,
+    fn compute_result(
+        &self,
+        call: &FunctionCall,
+        mappings: &HashMap<&'static str, String>,
+    ) -> Option<Expression> {
+        let expressions = call.get_arguments().clone().to_expressions();
+
+        Some(match expressions.len() {
+            0 => Expression::nil(),
+            1 => expressions
+                .into_iter()
+                .next()
+                .expect("at least one expression is expected"),
+            _ => FunctionCall::from_name(
+                mappings
+                    .get("select")
+                    .cloned()
+                    .unwrap_or_else(|| "select".to_owned()),
+            )
+            .with_arguments(TupleArguments::from_iter(
+                iter::once(Expression::from(1)).chain(expressions),
+            ))
+            .into(),
+        })
+    }
+
+    fn reserve_globals(&self) -> impl Iterator<Item = &'static str> {
+        iter::once("select")
     }
 }
 
 impl FlawlessRule for RemoveAssertions {
     fn flawless_process(&self, block: &mut Block, _: &Context) {
         let mut processor =
-            RemoveFunctionCallProcessor::new(self.preserve_args_side_effects, should_remove_call);
+            RemoveFunctionCallProcessor::new(self.preserve_args_side_effects, AssertMatcher);
         ScopeVisitor::visit_block(block, &mut processor);
+
+        if let Some(statement) = processor.extract_reserved_globals() {
+            block.insert_statement(0, statement);
+        }
     }
 }
 
diff --git a/src/rules/remove_call_match.rs b/src/rules/remove_call_match.rs
index 0bcea24..5d70945 100644
--- a/src/rules/remove_call_match.rs
+++ b/src/rules/remove_call_match.rs
@@ -1,16 +1,34 @@
-use std::ops;
+use std::collections::HashMap;
+use std::{iter, ops};
 
-use crate::nodes::{Arguments, DoStatement, Expression, Prefix, Statement, TableEntry};
+use crate::nodes::{
+    Arguments, DoStatement, Expression, FunctionCall, Identifier, LocalAssignStatement, Prefix,
+    Statement, TableEntry, TypedIdentifier,
+};
 use crate::process::{Evaluator, IdentifierTracker, NodeProcessor};
 use crate::utils::{expressions_as_expression, expressions_as_statement};
 
 pub(crate) trait CallMatch<T> {
     fn matches(&self, identifiers: &IdentifierTracker, prefix: &Prefix) -> bool;
+
+    fn compute_result(
+        &self,
+        _call: &FunctionCall,
+        _mappings: &HashMap<&'static str, String>,
+    ) -> Option<Expression> {
+        None
+    }
+
+    fn reserve_globals(&self) -> impl Iterator<Item = &'static str> {
+        iter::empty()
+    }
 }
 
 #[derive(Default)]
 pub(crate) struct RemoveFunctionCallProcessor<Args, T: CallMatch<Args>> {
     identifier_tracker: IdentifierTracker,
+    global_mappings: HashMap<&'static str, String>,
+    global_counter: u32,
     evaluator: Evaluator,
     preserve_args_side_effects: bool,
     matcher: T,
@@ -39,6 +57,8 @@ impl<Args, T: CallMatch<Args>> RemoveFunctionCallProcessor<Args, T> {
     pub(crate) fn new(preserve_args_side_effects: bool, matcher: T) -> Self {
         Self {
             identifier_tracker: Default::default(),
+            global_mappings: Default::default(),
+            global_counter: 0,
             evaluator: Default::default(),
             preserve_args_side_effects,
             matcher,
@@ -46,6 +66,23 @@ impl<Args, T: CallMatch<Args>> RemoveFunctionCallProcessor<Args, T> {
         }
     }
 
+    pub(crate) fn extract_reserved_globals(&mut self) -> Option<Statement> {
+        let (variables, values) = self.global_mappings.drain().fold(
+            (Vec::new(), Vec::new()),
+            |(mut variables, mut values), (global, reserved_name)| {
+                variables.push(TypedIdentifier::new(reserved_name));
+                values.push(Identifier::new(global).into());
+                (variables, values)
+            },
+        );
+
+        if variables.is_empty() {
+            None
+        } else {
+            Some(LocalAssignStatement::new(variables, values).into())
+        }
+    }
+
     fn preserve_side_effects(&self, arguments: &Arguments) -> Vec<Expression> {
         match arguments {
             Arguments::Tuple(tuple) => tuple
@@ -88,6 +125,11 @@ impl<Args, T: CallMatch<Args>> RemoveFunctionCallProcessor<Args, T> {
             Arguments::String(_) => Vec::new(),
         }
     }
+
+    fn get_reserved_global(&mut self) -> String {
+        self.global_counter += 1;
+        format!("__DARKLUA_REMOVE_CALL_RESERVED_{}", self.global_counter)
+    }
 }
 
 impl<Args, T: CallMatch<Args>> ops::Deref for RemoveFunctionCallProcessor<Args, T> {
@@ -107,9 +149,10 @@ impl<Args, T: CallMatch<Args>> ops::DerefMut for RemoveFunctionCallProcessor<Arg
 impl<Args, T: CallMatch<Args>> NodeProcessor for RemoveFunctionCallProcessor<Args, T> {
     fn process_statement(&mut self, statement: &mut Statement) {
         if let Statement::Call(call) = statement {
-            if self
-                .matcher
-                .matches(&self.identifier_tracker, call.get_prefix())
+            if call.get_method().is_none()
+                && self
+                    .matcher
+                    .matches(&self.identifier_tracker, call.get_prefix())
             {
                 *statement = if self.preserve_args_side_effects {
                     expressions_as_statement(self.preserve_side_effects(call.get_arguments()))
@@ -122,15 +165,34 @@ impl<Args, T: CallMatch<Args>> NodeProcessor for RemoveFunctionCallProcessor<Arg
 
     fn process_expression(&mut self, expression: &mut Expression) {
         if let Expression::Call(call) = expression {
-            if self
-                .matcher
-                .matches(&self.identifier_tracker, call.get_prefix())
+            if call.get_method().is_none()
+                && self
+                    .matcher
+                    .matches(&self.identifier_tracker, call.get_prefix())
             {
-                *expression = if self.preserve_args_side_effects {
-                    expressions_as_expression(self.preserve_side_effects(call.get_arguments()))
+                let insert_globals = self
+                    .matcher
+                    .reserve_globals()
+                    .filter(|global| {
+                        self.is_identifier_used(global)
+                            && !self.global_mappings.contains_key(global)
+                    })
+                    .collect::<Vec<_>>();
+
+                for global in insert_globals {
+                    let new_reserved_name = self.get_reserved_global();
+                    self.global_mappings.insert(global, new_reserved_name);
+                }
+
+                if let Some(result) = self.matcher.compute_result(call, &self.global_mappings) {
+                    *expression = result;
                 } else {
-                    Expression::nil()
-                };
+                    *expression = if self.preserve_args_side_effects {
+                        expressions_as_expression(self.preserve_side_effects(call.get_arguments()))
+                    } else {
+                        Expression::nil()
+                    };
+                }
             }
         }
     }
diff --git a/src/rules/remove_debug_profiling.rs b/src/rules/remove_debug_profiling.rs
index 5821d4e..8afe18f 100644
--- a/src/rules/remove_debug_profiling.rs
+++ b/src/rules/remove_debug_profiling.rs
@@ -47,6 +47,10 @@ impl FlawlessRule for RemoveDebugProfiling {
         let mut processor =
             RemoveFunctionCallProcessor::new(self.preserve_args_side_effects, should_remove_call);
         ScopeVisitor::visit_block(block, &mut processor);
+
+        if let Some(statement) = processor.extract_reserved_globals() {
+            block.insert_statement(0, statement);
+        }
     }
 }
 


Test patch:
diff --git a/tests/rule_tests/remove_assertions.rs b/tests/rule_tests/remove_assertions.rs
index 19b1398..b35aaa6 100644
--- a/tests/rule_tests/remove_assertions.rs
+++ b/tests/rule_tests/remove_assertions.rs
@@ -1,17 +1,24 @@
 use darklua_core::rules::{RemoveAssertions, Rule};
 
 test_rule!(
-    remove_debug_profiling,
+    remove_assertions,
     RemoveAssertions::default(),
     remove_variable_condition("assert(condition)") => "do end",
     remove_variable_condition_with_message("assert(condition, 'message')") => "do end",
     remove_function_call_condition("assert(validate(value))") => "validate(value)",
     remove_variable_condition_with_function_call_message("assert(condition, formatter(condition))") => "formatter(condition)",
     remove_function_call_condition_and_function_call_message("assert(validate(value), formatter(value))") => "do validate(value) formatter(value) end",
+    as_expression_remove_variable_condition("return assert(condition)") => "return condition",
+    as_expression_remove_variable_condition_with_message("return { assert(condition, 'message') }") => "return { select(1, condition, 'message') }",
+    as_expression_remove_variable_condition_with_message_but_select_is_used("local select = true\nreturn { assert(condition, 'message') }")
+        => "local __DARKLUA_REMOVE_CALL_RESERVED_1 = select local select = true\nreturn { __DARKLUA_REMOVE_CALL_RESERVED_1(1, condition, 'message') }",
+    as_expression_remove_function_call_condition("call(assert(validate(value)))") => "call(validate(value))",
+    as_expression_remove_variable_condition_with_function_call_message("return (assert(condition, formatter(condition)))") => "return (select(1, condition, formatter(condition)))",
+    as_expression_remove_function_call_condition_and_function_call_message("return assert(validate(value), formatter(value))") => "return select(1, validate(value), formatter(value))",
 );
 
 test_rule!(
-    remove_debug_profiling_without_side_effects,
+    remove_assertions_without_side_effects,
     json5::from_str::<Box<dyn Rule>>(
         r#"{
         rule: 'remove_assertions',
@@ -28,6 +35,7 @@ test_rule!(
 test_rule_without_effects!(
     RemoveAssertions::default(),
     assert_function_used("local function assert() end assert('label')"),
+    assert_with_method_call("assert:oops(condition)"),
 );
 
 #[test]