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

dalbit_core/modifiers/
convert_bit32.rs
