dalbit_core/modifiers/
remove_generalized_iteration.rs

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    skip_block_once: bool,
20}
21
22fn get_type_condition(arg: Expression, type_name: &str) -> Box<BinaryExpression> {
23    let type_call = Box::new(FunctionCall::new(
24        Prefix::from_name("type"),
25        TupleArguments::new(vec![arg]).into(),
26        None,
27    ));
28    Box::new(BinaryExpression::new(
29        BinaryOperator::Equal,
30        Expression::Call(type_call),
31        Expression::String(StringExpression::from_value(type_name)),
32    ))
33}
34
35impl Processor {
36    fn process_into_do(&self, block: &mut Block) -> Option<(usize, Statement)> {
37        let mut result: Option<(usize, Statement)> = None;
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 pairs_call = FunctionCall::new(
109                    //     Prefix::from_name("pairs"),
110                    //     TupleArguments::new(vec![iterator_identifier.clone().into()]).into(),
111                    //     None,
112                    // );
113
114                    let assign_from_pairs = AssignStatement::new(
115                        vec![
116                            Variable::Identifier(iterator_identifier.clone()),
117                            Variable::Identifier(invariant_identifier),
118                            Variable::Identifier(control_identifier),
119                        ],
120                        vec![Identifier::new("next").into(), iterator_identifier.into()],
121                    );
122
123                    let if_mt_table_block = Block::new(vec![assign_from_iter.into()], None);
124                    let if_not_mt_table_block = Block::new(vec![assign_from_pairs.into()], None);
125                    let if_mt_table_branch = IfBranch::new(
126                        Expression::Binary(Box::new(BinaryExpression::new(
127                            BinaryOperator::And,
128                            Expression::Binary(if_mt_table_condition),
129                            Expression::Binary(if_mt_iter_function_condition),
130                        ))),
131                        if_mt_table_block,
132                    );
133                    let if_mt_table_stmt =
134                        IfStatement::new(vec![if_mt_table_branch], Some(if_not_mt_table_block));
135
136                    let if_table_block =
137                        Block::new(vec![mt_local_assign.into(), if_mt_table_stmt.into()], None);
138                    let if_table_branch =
139                        IfBranch::new(Expression::Binary(if_table_condition), if_table_block);
140                    let if_table_stmt = IfStatement::new(vec![if_table_branch], None);
141
142                    stmts.push(iter_invar_control_local_assign.into());
143                    stmts.push(if_table_stmt.into());
144                    stmts.push(generic_for.clone().into());
145
146                    result = Some((i, DoStatement::new(Block::new(stmts, None)).into()));
147                    // return Some((i, DoStatement::new(Block::new(stmts, None)).into()));
148                }
149            }
150        }
151
152        if let Some(result) = result {
153            block.remove_statement(result.0);
154            Some(result)
155        } else {
156            None
157        }
158    }
159}
160
161impl NodeProcessor for Processor {
162    fn process_block(&mut self, block: &mut Block) {
163        if self.skip_block_once {
164            self.skip_block_once = false;
165            return;
166        }
167        let do_stmt = self.process_into_do(block);
168        if let Some((i, stmt)) = do_stmt {
169            self.skip_block_once = true;
170            block.insert_statement(i, stmt);
171        }
172    }
173}
174
175pub const REMOVE_GENERALIZED_ITERATION_MODIFIER_NAME: &str = "remove_generalized_iteration";
176
177/// A rule that removes generalized iteration.
178#[derive(Debug, PartialEq, Eq)]
179pub struct RemoveGeneralizedIteration {
180    runtime_identifier_format: String,
181}
182
183impl Default for RemoveGeneralizedIteration {
184    fn default() -> Self {
185        Self {
186            runtime_identifier_format: "_DALBIT_REMOVE_GENERALIZED_ITERATION_{name}{hash}"
187                .to_string(),
188        }
189    }
190}
191
192impl Rule for RemoveGeneralizedIteration {
193    fn process(&self, block: &mut Block, _: &Context) -> RuleProcessResult {
194        let var_builder = RuntimeIdentifierBuilder::new(
195            self.runtime_identifier_format.as_str(),
196            format!("{block:?}").as_bytes(),
197            Some(vec![METATABLE_VARIABLE_NAME.to_string()]),
198        )?;
199        let mut processor = Processor {
200            iterator_identifier: var_builder.build("iter")?,
201            invariant_identifier: var_builder.build("invar")?,
202            control_identifier: var_builder.build("control")?,
203            skip_block_once: false,
204        };
205        DefaultVisitor::visit_block(block, &mut processor);
206        Ok(())
207    }
208}
209
210impl RuleConfiguration for RemoveGeneralizedIteration {
211    fn configure(&mut self, _: RuleProperties) -> Result<(), RuleConfigurationError> {
212        Ok(())
213    }
214
215    fn get_name(&self) -> &'static str {
216        REMOVE_GENERALIZED_ITERATION_MODIFIER_NAME
217    }
218
219    fn serialize_to_properties(&self) -> RuleProperties {
220        RuleProperties::new()
221    }
222}