luaur_compiler/records/
table_mutation_tracker.rs

1use luaur_ast::records::ast_expr::AstExpr;
2use luaur_ast::records::ast_expr_binary::AstExprBinary;
3use luaur_ast::records::ast_expr_call::AstExprCall;
4use luaur_ast::records::ast_expr_group::AstExprGroup;
5use luaur_ast::records::ast_expr_if_else::AstExprIfElse;
6use luaur_ast::records::ast_expr_index_expr::AstExprIndexExpr;
7use luaur_ast::records::ast_expr_index_name::AstExprIndexName;
8use luaur_ast::records::ast_expr_instantiate::AstExprInstantiate;
9use luaur_ast::records::ast_expr_local::AstExprLocal;
10use luaur_ast::records::ast_expr_table::AstExprTable;
11use luaur_ast::records::ast_expr_type_assertion::AstExprTypeAssertion;
12use luaur_ast::records::ast_stat_assign::AstStatAssign;
13use luaur_ast::records::ast_stat_compound_assign::AstStatCompoundAssign;
14use luaur_ast::records::ast_stat_for_in::AstStatForIn;
15use luaur_ast::records::ast_stat_function::AstStatFunction;
16use luaur_ast::records::ast_stat_local::AstStatLocal;
17use luaur_ast::records::ast_stat_return::AstStatReturn;
18use luaur_ast::records::ast_visitor::AstVisitor;
19
20use luaur_common::records::dense_hash_map::DenseHashMap;
21use luaur_common::records::dense_hash_set::DenseHashSet;
22
23use crate::records::variable::Variable;
24use luaur_ast::records::ast_local::AstLocal;
25
26#[derive(Debug)]
27pub struct TableMutationTracker<'a> {
28    pub(crate) variables: &'a DenseHashMap<*mut AstLocal, Variable>,
29    pub(crate) escaped: DenseHashSet<*mut AstLocal>,
30}
31
32impl<'a> TableMutationTracker<'a> {
33    pub fn table_mutation_tracker(variables: &'a DenseHashMap<*mut AstLocal, Variable>) -> Self {
34        Self {
35            variables,
36            escaped: DenseHashSet::new(core::ptr::null_mut()),
37        }
38    }
39
40    pub fn mark_escaped(&mut self, mut expr: *mut AstExpr) {
41        loop {
42            if expr.is_null() {
43                return;
44            }
45            unsafe {
46                let node_ptr = expr as *mut luaur_ast::records::ast_node::AstNode;
47
48                let local = luaur_ast::rtti::ast_node_as::<AstExprLocal>(node_ptr);
49                if !local.is_null() {
50                    self.escaped.insert((*local).local);
51                    return;
52                }
53
54                let group = luaur_ast::rtti::ast_node_as::<AstExprGroup>(node_ptr);
55                if !group.is_null() {
56                    expr = (*group).expr;
57                    continue;
58                }
59
60                let assertion = luaur_ast::rtti::ast_node_as::<AstExprTypeAssertion>(node_ptr);
61                if !assertion.is_null() {
62                    expr = (*assertion).expr;
63                    continue;
64                }
65
66                let inst = luaur_ast::rtti::ast_node_as::<AstExprInstantiate>(node_ptr);
67                if !inst.is_null() {
68                    expr = (*inst).expr;
69                    continue;
70                }
71
72                let if_else = luaur_ast::rtti::ast_node_as::<AstExprIfElse>(node_ptr);
73                if !if_else.is_null() {
74                    self.mark_escaped((*if_else).true_expr);
75                    expr = (*if_else).false_expr;
76                    continue;
77                }
78
79                let bin = luaur_ast::rtti::ast_node_as::<AstExprBinary>(node_ptr);
80                if !bin.is_null() {
81                    if (*bin).op == AstExprBinary::And || (*bin).op == AstExprBinary::Or {
82                        self.mark_escaped((*bin).left);
83                        expr = (*bin).right;
84                        continue;
85                    } else {
86                        return;
87                    }
88                }
89
90                return;
91            }
92        }
93    }
94
95    pub fn mark_escaped_table_index(&mut self, expr: *mut AstExpr, is_lvalue: bool) {
96        if expr.is_null() {
97            return;
98        }
99        unsafe {
100            let node_ptr = expr as *mut luaur_ast::records::ast_node::AstNode;
101
102            let idx_name = luaur_ast::rtti::ast_node_as::<AstExprIndexName>(node_ptr);
103            if !idx_name.is_null() {
104                self.mark_escaped((*idx_name).expr);
105                return;
106            }
107
108            let idx_expr = luaur_ast::rtti::ast_node_as::<AstExprIndexExpr>(node_ptr);
109            if !idx_expr.is_null() {
110                self.mark_escaped((*idx_expr).expr);
111                if is_lvalue {
112                    self.mark_escaped((*idx_expr).index);
113                }
114            }
115        }
116    }
117}
118
119impl<'a> AstVisitor for TableMutationTracker<'a> {
120    fn visit_expr_call(&mut self, node: *mut core::ffi::c_void) -> bool {
121        let node = node as *mut AstExprCall;
122        unsafe {
123            for arg in (*node).args.iter() {
124                self.mark_escaped(*arg);
125            }
126
127            if (*node).self_ {
128                self.mark_escaped_table_index((*node).func, false);
129            }
130        }
131
132        true
133    }
134
135    fn visit_expr_table(&mut self, node: *mut core::ffi::c_void) -> bool {
136        let node = node as *mut AstExprTable;
137        unsafe {
138            for item in (*node).items.iter() {
139                if !item.key.is_null() {
140                    self.mark_escaped(item.key);
141                }
142                self.mark_escaped(item.value);
143            }
144        }
145        true
146    }
147
148    fn visit_stat_local(&mut self, node: *mut core::ffi::c_void) -> bool {
149        let node = node as *mut AstStatLocal;
150        unsafe {
151            for (value, _var) in (*node).values.iter().zip((*node).vars.iter()) {
152                self.mark_escaped(*value);
153            }
154        }
155        true
156    }
157
158    fn visit_stat_assign(&mut self, node: *mut core::ffi::c_void) -> bool {
159        let node = node as *mut AstStatAssign;
160        unsafe {
161            for rhs in (*node).values.iter() {
162                self.mark_escaped(*rhs);
163            }
164            for lhs in (*node).vars.iter() {
165                self.mark_escaped_table_index(*lhs, true);
166            }
167        }
168        true
169    }
170
171    fn visit_stat_compound_assign(&mut self, node: *mut core::ffi::c_void) -> bool {
172        let node = node as *mut AstStatCompoundAssign;
173        unsafe {
174            self.mark_escaped_table_index((*node).var, true);
175        }
176        true
177    }
178
179    fn visit_stat_function(&mut self, node: *mut core::ffi::c_void) -> bool {
180        let node = node as *mut AstStatFunction;
181        unsafe {
182            self.mark_escaped_table_index((*node).name, true);
183        }
184        true
185    }
186
187    fn visit_stat_for_in(&mut self, node: *mut core::ffi::c_void) -> bool {
188        let node = node as *mut AstStatForIn;
189        unsafe {
190            for expr in (*node).values.iter() {
191                self.mark_escaped(*expr);
192            }
193        }
194        true
195    }
196
197    fn visit_stat_return(&mut self, node: *mut core::ffi::c_void) -> bool {
198        let node = node as *mut AstStatReturn;
199        unsafe {
200            for expr in (*node).list.iter() {
201                self.mark_escaped(*expr);
202            }
203        }
204        true
205    }
206}
luaur_compiler/records/table_mutation_tracker.rs

luaur_compiler/records/
table_mutation_tracker.rs
