dalbit_core/modifiers/
convert_bit32.rs

1use std::string::ToString;
2use std::{collections::HashMap, str::FromStr};
3
4use full_moon::ast::punctuated::Pair;
5use full_moon::ast::Do;
6use full_moon::node::Node;
7use full_moon::tokenizer::{Symbol, Token, TokenType};
8use full_moon::{
9    ast::{
10        punctuated::Punctuated, BinOp, Call, Expression, FunctionArgs, FunctionCall, Index, Stmt,
11        Suffix, UnOp, Var,
12    },
13    tokenizer::TokenReference,
14    visitors::VisitorMut,
15};
16use strum_macros::{Display, EnumString};
17
18use super::ast_util;
19
20pub const CONVERT_BIT32_MODIFIER_NAME: &str = "convert_bit32";
21const DEFAULT_BIT32_IDENTIFIER: &str = "bit32";
22const MASKING_NUMBER_TOKEN_SYMBOL: &str = "0xFFFFFFFF";
23
24#[inline]
25fn mask_32bit(exp: Expression) -> Expression {
26    ast_util::create_binary_operator(
27        exp,
28        BinOp::Ampersand(TokenReference::symbol("&").unwrap()),
29        ast_util::create_number(MASKING_NUMBER_TOKEN_SYMBOL),
30    )
31}
32
33fn index_to_string(index: &Index) -> Option<String> {
34    match index {
35        Index::Brackets {
36            brackets: _,
37            expression,
38        } => {
39            let mut string = expression.to_string();
40            string.remove(0);
41            string.pop();
42            Some(string)
43        }
44        Index::Dot { dot: _, name: _ } => {
45            let mut string = index.to_string();
46            string.remove(0);
47            Some(string)
48        }
49        _ => None,
50    }
51}
52
53#[derive(Debug, Display, EnumString)]
54enum Bit32Method {
55    #[strum(serialize = "rshift")]
56    RightShift,
57    #[strum(serialize = "lshift")]
58    LeftShift,
59    #[strum(serialize = "band")]
60    And,
61    #[strum(serialize = "bor")]
62    Or,
63    #[strum(serialize = "bxor")]
64    Xor,
65    #[strum(serialize = "bnot")]
66    Not,
67    #[strum(serialize = "btest")]
68    Test,
69}
70
71impl Bit32Method {
72    fn convert(&self, call: &Call) -> Option<Expression> {
73        if let Call::AnonymousCall(args) = call {
74            if let FunctionArgs::Parentheses {
75                parentheses,
76                arguments,
77            } = args
78            {
79                let mut iter = arguments.iter();
80                let first_arg = iter.next()?;
81
82                let binop = match self {
83                    Bit32Method::RightShift => {
84                        BinOp::DoubleGreaterThan(TokenReference::symbol(">>").unwrap())
85                    }
86                    Bit32Method::LeftShift => {
87                        BinOp::DoubleLessThan(TokenReference::symbol("<<").unwrap())
88                    }
89                    Bit32Method::And => BinOp::Ampersand(TokenReference::symbol("&").unwrap()),
90                    Bit32Method::Or => BinOp::Pipe(TokenReference::symbol("|").unwrap()),
91                    Bit32Method::Xor => BinOp::Tilde(TokenReference::symbol("~").unwrap()),
92                    Bit32Method::Not => {
93                        let masking_arg = mask_32bit(first_arg.clone());
94                        let parenthese = ast_util::create_parentheses(masking_arg, None);
95                        let bnot_exp = ast_util::create_unary_operator(
96                            UnOp::Tilde(TokenReference::symbol("~").unwrap()),
97                            parenthese,
98                        );
99                        let masking_unop = mask_32bit(bnot_exp);
100                        return Some(ast_util::create_parentheses(
101                            masking_unop,
102                            Some(parentheses.clone()),
103                        ));
104                    }
105                    Bit32Method::Test => {
106                        let second_arg = iter.next()?;
107                        let band_exp = ast_util::create_binary_operator(
108                            first_arg.clone(),
109                            BinOp::Ampersand(TokenReference::symbol("&").unwrap()),
110                            second_arg.clone(),
111                        );
112                        let parenthese = ast_util::create_parentheses(band_exp, None);
113                        let masking_bin_exp = mask_32bit(parenthese);
114                        let not_equal_exp = ast_util::create_binary_operator(
115                            masking_bin_exp,
116                            BinOp::TildeEqual(TokenReference::symbol("~=").unwrap()),
117                            ast_util::create_number("0"),
118                        );
119                        return Some(ast_util::create_parentheses(
120                            not_equal_exp,
121                            Some(parentheses.clone()),
122                        ));
123                    }
124                };
125
126                let second_arg = iter.next()?;
127                let bitop_exp =
128                    ast_util::create_binary_operator(first_arg.clone(), binop, second_arg.clone());
129                let parenthese = ast_util::create_parentheses(bitop_exp, None);
130                let masking_bin_exp = mask_32bit(parenthese);
131                return Some(ast_util::create_parentheses(
132                    masking_bin_exp,
133                    Some(parentheses.clone()),
134                ));
135            }
136        }
137        None
138    }
139}
140
141#[derive(Debug)]
142pub struct ConvertBit32 {
143    bit32_identifier: String,
144    bit32_methods: HashMap<String, Bit32Method>,
145}
146
147impl Default for ConvertBit32 {
148    fn default() -> Self {
149        Self {
150            bit32_identifier: DEFAULT_BIT32_IDENTIFIER.to_owned(),
151            bit32_methods: HashMap::new(),
152        }
153    }
154}
155
156impl VisitorMut for ConvertBit32 {
157    fn visit_stmt(&mut self, stmt: Stmt) -> Stmt {
158        match &stmt {
159            Stmt::FunctionCall(func_call) => {
160                if let Some(_) = self.convert(func_call) {
161                    if let Suffix::Call(Call::AnonymousCall(FunctionArgs::Parentheses {
162                        parentheses,
163                        arguments: _,
164                    })) = func_call.suffixes().last().unwrap()
165                    {
166                        return Stmt::Do(
167                            Do::new()
168                                .with_do_token(TokenReference::new(
169                                    func_call
170                                        .surrounding_trivia()
171                                        .0
172                                        .into_iter()
173                                        .cloned()
174                                        .collect(),
175                                    Token::new(TokenType::Symbol { symbol: Symbol::Do }),
176                                    vec![Token::new(TokenType::Whitespace {
177                                        characters: " ".into(),
178                                    })],
179                                ))
180                                .with_end_token(TokenReference::new(
181                                    Vec::new(),
182                                    Token::new(TokenType::Symbol {
183                                        symbol: Symbol::End,
184                                    }),
185                                    parentheses
186                                        .tokens()
187                                        .1
188                                        .trailing_trivia()
189                                        .cloned()
190                                        .chain(std::iter::once(Token::new(TokenType::Whitespace {
191                                            characters: " ".into(),
192                                        })))
193                                        .collect(),
194                                )),
195                        );
196                    }
197                }
198            }
199            Stmt::Assignment(assign) => {
200                if self.check_replaced(assign.variables(), assign.expressions()) {
201                    let mut do_trailing_trivia: Vec<Token> = Vec::new();
202                    for token_ref in assign.tokens() {
203                        for t in token_ref.leading_trivia() {
204                            do_trailing_trivia.push(t.to_owned());
205                        }
206                        for t in token_ref.trailing_trivia() {
207                            do_trailing_trivia.push(t.to_owned());
208                        }
209                    }
210                    return Stmt::Do(
211                        Do::new()
212                            .with_do_token(TokenReference::new(
213                                vec![Token::new(TokenType::Whitespace {
214                                    characters: " ".into(),
215                                })],
216                                Token::new(TokenType::Symbol { symbol: Symbol::Do }),
217                                do_trailing_trivia,
218                            ))
219                            .with_end_token(TokenReference::new(
220                                vec![Token::new(TokenType::Whitespace {
221                                    characters: " ".into(),
222                                })],
223                                Token::new(TokenType::Symbol {
224                                    symbol: Symbol::End,
225                                }),
226                                vec![Token::new(TokenType::Whitespace {
227                                    characters: " ".into(),
228                                })],
229                            )),
230                    );
231                }
232            }
233            Stmt::LocalAssignment(local_assign) => {
234                let mut variables: Punctuated<Var> = Punctuated::new();
235                for token in local_assign.names() {
236                    variables.push(Pair::new(Var::Name(token.clone()), None));
237                }
238                if self.check_replaced(&variables, local_assign.expressions()) {
239                    let mut do_trailing_trivia: Vec<Token> = Vec::new();
240                    for token_ref in local_assign.tokens() {
241                        for t in token_ref.leading_trivia() {
242                            do_trailing_trivia.push(t.to_owned());
243                        }
244                        for t in token_ref.trailing_trivia() {
245                            do_trailing_trivia.push(t.to_owned());
246                        }
247                    }
248                    return Stmt::Do(
249                        Do::new()
250                            .with_do_token(TokenReference::new(
251                                vec![Token::new(TokenType::Whitespace {
252                                    characters: " ".into(),
253                                })],
254                                Token::new(TokenType::Symbol { symbol: Symbol::Do }),
255                                do_trailing_trivia,
256                            ))
257                            .with_end_token(TokenReference::new(
258                                Vec::new(),
259                                Token::new(TokenType::Symbol {
260                                    symbol: Symbol::End,
261                                }),
262                                vec![Token::new(TokenType::Whitespace {
263                                    characters: " ".into(),
264                                })],
265                            )),
266                    );
267                }
268            }
269            _ => {}
270        }
271        stmt
272    }
273
274    /// To convert bit32 methods/calls and linked identifiers into bitwise operators
275    ///
276    /// Conversion Example: `local x = bit32.band; local y = x(1, 2)` -> `do then; local y = ((1&2)&0xFFFFFFFF)`
277    fn visit_expression(&mut self, exp: Expression) -> Expression {
278        if let Expression::FunctionCall(func_call) = &exp {
279            if let Some(exp) = self.convert(func_call) {
280                return exp;
281            }
282        }
283        exp
284    }
285}
286
287impl ConvertBit32 {
288    #[inline]
289    fn is_bit32_identifier(&self, string: impl Into<String>) -> bool {
290        string.into() == self.bit32_identifier
291    }
292
293    fn check_replaced(
294        &mut self,
295        variables: &Punctuated<Var>,
296        expressions: &Punctuated<Expression>,
297    ) -> bool {
298        for (var, exp) in variables.iter().zip(expressions.iter()) {
299            // local x = bit32.band
300            if let Expression::Var(exp) = exp {
301                match exp {
302                    Var::Expression(var_exp) => {
303                        if !self.is_bit32_identifier(var_exp.prefix().to_string()) {
304                            return false;
305                        }
306                        let mut iter = var_exp.suffixes();
307                        let first = iter.next();
308                        if let Some(first) = first {
309                            // there's an index(ex. `band`)
310                            if let Suffix::Index(index) = first {
311                                let index = index_to_string(index);
312                                if let Some(index) = index {
313                                    if let Ok(method) = Bit32Method::from_str(index.trim()) {
314                                        self.bit32_methods
315                                            .insert(var.to_string().trim().to_owned(), method);
316                                        return true;
317                                    }
318                                }
319                            }
320                        }
321                    }
322                    Var::Name(_) => {
323                        if self.is_bit32_identifier(exp.to_string().trim().to_owned()) {
324                            self.bit32_identifier = var.to_string().trim().to_owned();
325                            return true;
326                        }
327                    }
328                    _ => {}
329                }
330            }
331        }
332        false
333    }
334
335    fn convert(&mut self, func_call: &FunctionCall) -> Option<Expression> {
336        let mut iter = func_call.suffixes();
337        let first = iter.next();
338        let second = iter.next();
339        let prefix = func_call.prefix().to_string();
340        match (first, second) {
341            (Some(first), Some(second)) => {
342                if !self.is_bit32_identifier(prefix) {
343                    return None;
344                }
345                match (first, second) {
346                    // there's another index(ex. `band`) before a call(ex. `(1, 2)`)
347                    (Suffix::Index(index), Suffix::Call(call)) => {
348                        let index = index_to_string(index)?;
349                        if let Ok(method) = Bit32Method::from_str(index.trim()) {
350                            return method.convert(call);
351                        }
352                        None
353                    }
354                    _ => None,
355                }
356            }
357            (Some(first), None) => {
358                // there's only a call(ex. `(1, 2)`)
359                if let Suffix::Call(call) = first {
360                    if let Some(method) = self.bit32_methods.get(&prefix) {
361                        return method.convert(call);
362                    }
363                }
364                None
365            }
366            _ => None,
367        }
368    }
369}
dalbit_core/modifiers/convert_bit32.rs

dalbit_core/modifiers/
convert_bit32.rs
