1use darklua_core::nodes::{
2 AssignStatement, BinaryExpression, BinaryOperator, Block, DoStatement, Expression,
3 FieldExpression, FunctionCall, Identifier, IfBranch, IfStatement, LocalAssignStatement, Prefix,
4 Statement, StringExpression, TupleArguments, TypedIdentifier, Variable,
5};
6use darklua_core::process::{DefaultVisitor, NodeProcessor, NodeVisitor};
7use darklua_core::rules::{Context, RuleConfiguration, RuleConfigurationError, RuleProperties};
8
9use super::runtime_identifier::RuntimeIdentifierBuilder;
10use darklua_core::rules::{Rule, RuleProcessResult};
11
12const METATABLE_VARIABLE_NAME: &str = "m";
13const GETMETATABLE_IDENTIFIER: &str = "__DALBIT_getmetatable_iter";
14
15struct Processor {
16 iterator_identifier: String,
17 invariant_identifier: String,
18 control_identifier: String,
19}
20
21fn get_type_condition(arg: Expression, type_name: &str) -> Box<BinaryExpression> {
22 let type_call = Box::new(FunctionCall::new(
23 Prefix::from_name("type"),
24 TupleArguments::new(vec![arg]).into(),
25 None,
26 ));
27 Box::new(BinaryExpression::new(
28 BinaryOperator::Equal,
29 Expression::Call(type_call),
30 Expression::String(StringExpression::from_value(type_name)),
31 ))
32}
33
34impl NodeProcessor for Processor {
35 fn process_block(&mut self, block: &mut Block) {
36 let mut result: Vec<(usize, Statement)> = Vec::new();
38 for (i, stmt) in block.iter_mut_statements().enumerate() {
39 if let Statement::GenericFor(generic_for) = stmt {
40 let exps = generic_for.mutate_expressions();
41 if exps.len() == 1 {
42 let mut stmts: Vec<Statement> = Vec::new();
43 let iterator_typed_identifier =
44 TypedIdentifier::new(self.iterator_identifier.as_str());
45 let iterator_identifier = iterator_typed_identifier.get_identifier().clone();
46
47 let invariant_typed_identifier =
48 TypedIdentifier::new(self.invariant_identifier.as_str());
49 let invariant_identifier = invariant_typed_identifier.get_identifier().clone();
50
51 let control_typed_identifier =
52 TypedIdentifier::new(self.control_identifier.as_str());
53 let control_identifier = control_typed_identifier.get_identifier().clone();
54
55 let iter_invar_control_local_assign = LocalAssignStatement::new(
56 vec![
57 iterator_typed_identifier,
58 invariant_typed_identifier,
59 control_typed_identifier,
60 ],
61 vec![exps[0].to_owned()],
62 );
63
64 let iterator_exp = Expression::Identifier(iterator_identifier.clone());
65 exps[0] = iterator_exp.clone();
66 let invariant_exp = Expression::Identifier(invariant_identifier.clone());
67 exps.push(invariant_exp);
68 let control_exp = Expression::Identifier(control_identifier.clone());
69 exps.push(control_exp);
70
71 let if_table_condition = get_type_condition(iterator_exp.clone(), "table");
72
73 let mt_typed_identifier = TypedIdentifier::new(METATABLE_VARIABLE_NAME);
74 let mt_identifier = mt_typed_identifier.get_identifier().clone();
75
76 let get_mt_call = FunctionCall::new(
77 Prefix::from_name(GETMETATABLE_IDENTIFIER),
78 TupleArguments::new(vec![iterator_exp.clone()]).into(),
79 None,
80 );
81 let mt_local_assign = LocalAssignStatement::new(
82 vec![mt_typed_identifier],
83 vec![get_mt_call.into()],
84 );
85
86 let if_mt_table_condition =
87 get_type_condition(mt_identifier.clone().into(), "table");
88 let mt_iter = FieldExpression::new(
89 Prefix::Identifier(mt_identifier.clone()),
90 Identifier::new("__iter"),
91 );
92 let if_mt_iter_function_condition =
93 get_type_condition(mt_iter.clone().into(), "function");
94
95 let mut mt_iter_call = FunctionCall::from_prefix(Box::new(mt_iter));
96 mt_iter_call = mt_iter_call
97 .with_argument(Expression::identifier(iterator_identifier.clone()));
98
99 let assign_from_iter = AssignStatement::new(
100 vec![
101 Variable::Identifier(iterator_identifier.clone()),
102 Variable::Identifier(invariant_identifier.clone()),
103 Variable::Identifier(control_identifier.clone()),
104 ],
105 vec![mt_iter_call.into()],
106 );
107
108 let assign_from_pairs = AssignStatement::new(
109 vec![
110 Variable::Identifier(iterator_identifier.clone()),
111 Variable::Identifier(invariant_identifier),
112 Variable::Identifier(control_identifier),
113 ],
114 vec![Identifier::new("next").into(), iterator_identifier.into()],
115 );
116
117 let if_mt_table_block = Block::new(vec![assign_from_iter.into()], None);
118 let if_not_mt_table_block = Block::new(vec![assign_from_pairs.into()], None);
119 let if_mt_table_branch = IfBranch::new(
120 Expression::Binary(Box::new(BinaryExpression::new(
121 BinaryOperator::And,
122 Expression::Binary(if_mt_table_condition),
123 Expression::Binary(if_mt_iter_function_condition),
124 ))),
125 if_mt_table_block,
126 );
127 let if_mt_table_stmt =
128 IfStatement::new(vec![if_mt_table_branch], Some(if_not_mt_table_block));
129
130 let if_table_block =
131 Block::new(vec![mt_local_assign.into(), if_mt_table_stmt.into()], None);
132 let if_table_branch =
133 IfBranch::new(Expression::Binary(if_table_condition), if_table_block);
134 let if_table_stmt = IfStatement::new(vec![if_table_branch], None);
135
136 stmts.push(iter_invar_control_local_assign.into());
137 stmts.push(if_table_stmt.into());
138 stmts.push(generic_for.clone().into());
139
140 result.push((i, DoStatement::new(Block::new(stmts, None)).into()))
141 }
142 }
143 }
144
145 for (i, stmt) in result {
146 block.remove_statement(i);
147 block.insert_statement(i, stmt);
148 }
149 }
150}
151
152pub const REMOVE_GENERALIZED_ITERATION_MODIFIER_NAME: &str = "remove_generalized_iteration";
153
154#[derive(Debug, PartialEq, Eq)]
156pub struct RemoveGeneralizedIteration {
157 runtime_identifier_format: String,
158}
159
160impl Default for RemoveGeneralizedIteration {
161 fn default() -> Self {
162 Self {
163 runtime_identifier_format: "_DALBIT_REMOVE_GENERALIZED_ITERATION_{name}{hash}"
164 .to_string(),
165 }
166 }
167}
168
169impl Rule for RemoveGeneralizedIteration {
170 fn process(&self, block: &mut Block, _: &Context) -> RuleProcessResult {
171 let var_builder = RuntimeIdentifierBuilder::new(
172 self.runtime_identifier_format.as_str(),
173 format!("{block:?}").as_bytes(),
174 Some(vec![METATABLE_VARIABLE_NAME.to_string()]),
175 )?;
176 let mut processor = Processor {
177 iterator_identifier: var_builder.build("iter")?,
178 invariant_identifier: var_builder.build("invar")?,
179 control_identifier: var_builder.build("control")?,
180 };
181 DefaultVisitor::visit_block(block, &mut processor);
182 Ok(())
183 }
184}
185
186impl RuleConfiguration for RemoveGeneralizedIteration {
187 fn configure(&mut self, _: RuleProperties) -> Result<(), RuleConfigurationError> {
188 Ok(())
189 }
190
191 fn get_name(&self) -> &'static str {
192 REMOVE_GENERALIZED_ITERATION_MODIFIER_NAME
193 }
194
195 fn serialize_to_properties(&self) -> RuleProperties {
196 RuleProperties::new()
197 }
198}