diff --git a/src/rules/remove_generalized_iteration.rs b/src/rules/remove_generalized_iteration.rs index 09544fd4..b1741c3a 100644 --- a/src/rules/remove_generalized_iteration.rs +++ b/src/rules/remove_generalized_iteration.rs @@ -9,13 +9,14 @@ use crate::rules::{Context, RuleConfiguration, RuleConfigurationError, RulePrope use super::runtime_variable::RuntimeVariableBuilder; use super::{Rule, RuleProcessResult}; -const METATABLE_VARIABLE_NAME: &str = "_m"; +const METATABLE_VARIABLE_NAME: &str = "m"; struct Processor { iterator_variable_name: String, invariant_variable_name: String, control_variable_name: String, skip_block_once: bool, + getmetatable: String } fn get_type_condition(arg: Expression, type_name: &str) -> Box { @@ -73,7 +74,7 @@ impl Processor { let mt_identifier = mt_typed_identifier.get_identifier().clone(); let get_mt_call = FunctionCall::new( - Prefix::from_name("getmetatable"), + Prefix::from_name(self.getmetatable.as_str()), TupleArguments::new(vec![iterator_exp.clone()]).into(), None, ); @@ -169,6 +170,7 @@ pub const REMOVE_GENERALIZED_ITERATION_RULE_NAME: &str = "remove_generalized_ite #[derive(Debug, PartialEq, Eq)] pub struct RemoveGeneralizedIteration { runtime_variable_format: String, + getmetatable: String } impl Default for RemoveGeneralizedIteration { @@ -176,6 +178,7 @@ impl Default for RemoveGeneralizedIteration { Self { runtime_variable_format: "_DARKLUA_REMOVE_GENERALIZED_ITERATION_{name}{hash}" .to_string(), + getmetatable: "getmetatable".to_string() } } } @@ -192,6 +195,7 @@ impl Rule for RemoveGeneralizedIteration { invariant_variable_name: var_builder.build("invar")?, control_variable_name: var_builder.build("control")?, skip_block_once: false, + getmetatable: self.getmetatable.to_owned() }; DefaultVisitor::visit_block(block, &mut processor); Ok(()) @@ -205,6 +209,9 @@ impl RuleConfiguration for RemoveGeneralizedIteration { "runtime_variable_format" => { self.runtime_variable_format = value.expect_string(&key)?; } + "getmetatable" => { + self.getmetatable = value.expect_string(&key)?; + } _ => return Err(RuleConfigurationError::UnexpectedProperty(key)), } } diff --git a/tests/rule_tests/remove_generalized_iteration.rs b/tests/rule_tests/remove_generalized_iteration.rs index a97aee90..02d4646a 100644 --- a/tests/rule_tests/remove_generalized_iteration.rs +++ b/tests/rule_tests/remove_generalized_iteration.rs @@ -9,7 +9,20 @@ test_rule!( }"# ).unwrap(), generic_for("for i,v in {1,2,3} do end") - => "do local iter={1,2,3} local invar,control if type(iter)=='table' then local _m=getmetatable(iter) if type(_m)=='table' and type(_m.__iter)=='function' then iter,invar,control=_m.__iter() else iter,invar,control=pairs(iter) end end for i,v in iter,invar,control do end end" + => "do local iter={1,2,3} local invar,control if type(iter)=='table' then local m=getmetatable(iter) if type(m)=='table' and type(m.__iter)=='function' then iter,invar,control=m.__iter() else iter,invar,control=pairs(iter) end end for i,v in iter,invar,control do end end" +); + +test_rule!( + remove_generalized_iteration_with_custom_getmetatable, + json5::from_str::>( + r#"{ + rule: 'remove_generalized_iteration', + runtime_variable_format: '{name}', + getmetatable: '_custom_getmetatable' + }"# + ).unwrap(), + generic_for("for i,v in {1,2,3} do end") + => "do local iter={1,2,3} local invar,control if type(iter)=='table' then local m=_custom_getmetatable(iter) if type(m)=='table' and type(m.__iter)=='function' then iter,invar,control=m.__iter() else iter,invar,control=pairs(iter) end end for i,v in iter,invar,control do end end" ); #[test]