Add getmetatable config field

CavefulGames · Sep 30, 2024 · dddb313 · dddb313
1 parent 42cecdd
commit dddb313
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git 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<BinaryExpression> {
@@ -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,13 +170,15 @@ 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 {
     fn default() -> Self {
         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
@@ -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::<Box<dyn Rule>>(
+        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]