darklua_core/process/
scope_visitor.rs1use std::collections::HashSet;
2use std::ops::DerefMut;
3
4use crate::nodes::*;
5use crate::process::utils::is_valid_identifier;
6use crate::process::{NodeProcessor, NodeVisitor};
7
8use super::utils::{identifier_permutator, Permutator};
9
10pub trait Scope {
13 fn push(&mut self);
15 fn pop(&mut self);
18 fn insert(&mut self, identifier: &mut String);
21 fn insert_self(&mut self);
23 fn insert_local(&mut self, identifier: &mut String, value: Option<&mut Expression>);
25 fn insert_local_function(&mut self, function: &mut LocalFunctionStatement);
27}
28
29pub struct ScopeVisitor;
31
32impl ScopeVisitor {
33 fn visit_block_without_push<T: NodeProcessor + Scope>(block: &mut Block, scope: &mut T) {
34 scope.process_block(block);
35
36 block
37 .iter_mut_statements()
38 .for_each(|statement| Self::visit_statement(statement, scope));
39
40 if let Some(last_statement) = block.mutate_last_statement() {
41 Self::visit_last_statement(last_statement, scope);
42 };
43 }
44}
45
46impl<T: NodeProcessor + Scope> NodeVisitor<T> for ScopeVisitor {
47 fn visit_block(block: &mut Block, scope: &mut T) {
48 scope.push();
49 Self::visit_block_without_push(block, scope);
50 scope.pop();
51 }
52
53 fn visit_local_assign(statement: &mut LocalAssignStatement, scope: &mut T) {
54 scope.process_local_assign_statement(statement);
55
56 statement
57 .iter_mut_values()
58 .for_each(|value| Self::visit_expression(value, scope));
59
60 for r#type in statement
61 .iter_mut_variables()
62 .filter_map(TypedIdentifier::mutate_type)
63 {
64 Self::visit_type(r#type, scope);
65 }
66
67 statement.for_each_assignment(|variable, expression| {
68 scope.insert_local(variable.mutate_name(), expression)
69 });
70 }
71
72 fn visit_function_expression(function: &mut FunctionExpression, scope: &mut T) {
73 scope.process_function_expression(function);
74
75 for r#type in function
76 .iter_mut_parameters()
77 .filter_map(TypedIdentifier::mutate_type)
78 {
79 Self::visit_type(r#type, scope);
80 }
81
82 if let Some(variadic_type) = function.mutate_variadic_type() {
83 Self::visit_function_variadic_type(variadic_type, scope);
84 }
85
86 if let Some(return_type) = function.mutate_return_type() {
87 Self::visit_function_return_type(return_type, scope);
88 }
89
90 scope.push();
91 function
92 .mutate_parameters()
93 .iter_mut()
94 .for_each(|parameter| scope.insert(parameter.mutate_name()));
95
96 scope.process_scope(function.mutate_block(), None);
97
98 Self::visit_block(function.mutate_block(), scope);
99 scope.pop();
100 }
101
102 fn visit_function_statement(statement: &mut FunctionStatement, scope: &mut T) {
103 scope.process_function_statement(statement);
104 scope.process_variable_expression(statement.mutate_function_name().mutate_identifier());
105
106 for r#type in statement
107 .iter_mut_parameters()
108 .filter_map(TypedIdentifier::mutate_type)
109 {
110 Self::visit_type(r#type, scope);
111 }
112
113 if let Some(variadic_type) = statement.mutate_variadic_type() {
114 Self::visit_function_variadic_type(variadic_type, scope);
115 }
116
117 if let Some(return_type) = statement.mutate_return_type() {
118 Self::visit_function_return_type(return_type, scope);
119 }
120
121 scope.push();
122 if statement.get_name().has_method() {
123 scope.insert_self();
124 }
125 statement
126 .mutate_parameters()
127 .iter_mut()
128 .for_each(|parameter| scope.insert(parameter.mutate_name()));
129
130 scope.process_scope(statement.mutate_block(), None);
131
132 Self::visit_block(statement.mutate_block(), scope);
133 scope.pop();
134 }
135
136 fn visit_local_function(statement: &mut LocalFunctionStatement, scope: &mut T) {
137 scope.process_local_function_statement(statement);
138
139 scope.insert_local_function(statement);
140
141 for r#type in statement
142 .iter_mut_parameters()
143 .filter_map(TypedIdentifier::mutate_type)
144 {
145 Self::visit_type(r#type, scope);
146 }
147
148 if let Some(variadic_type) = statement.mutate_variadic_type() {
149 Self::visit_function_variadic_type(variadic_type, scope);
150 }
151
152 if let Some(return_type) = statement.mutate_return_type() {
153 Self::visit_function_return_type(return_type, scope);
154 }
155
156 scope.push();
157 statement
158 .mutate_parameters()
159 .iter_mut()
160 .for_each(|parameter| scope.insert(parameter.mutate_name()));
161
162 scope.process_scope(statement.mutate_block(), None);
163
164 Self::visit_block(statement.mutate_block(), scope);
165 scope.pop();
166 }
167
168 fn visit_generic_for(statement: &mut GenericForStatement, scope: &mut T) {
169 scope.process_generic_for_statement(statement);
170
171 statement
172 .iter_mut_expressions()
173 .for_each(|expression| Self::visit_expression(expression, scope));
174
175 statement
176 .iter_mut_identifiers()
177 .for_each(|identifier| scope.insert(identifier.mutate_name()));
178
179 for r#type in statement
180 .iter_mut_identifiers()
181 .filter_map(TypedIdentifier::mutate_type)
182 {
183 Self::visit_type(r#type, scope);
184 }
185
186 scope.process_scope(statement.mutate_block(), None);
187
188 Self::visit_block(statement.mutate_block(), scope);
189 }
190
191 fn visit_numeric_for(statement: &mut NumericForStatement, scope: &mut T) {
192 scope.process_numeric_for_statement(statement);
193
194 Self::visit_expression(statement.mutate_start(), scope);
195 Self::visit_expression(statement.mutate_end(), scope);
196
197 if let Some(step) = statement.mutate_step() {
198 Self::visit_expression(step, scope);
199 };
200
201 if let Some(r#type) = statement.mutate_identifier().mutate_type() {
202 Self::visit_type(r#type, scope);
203 }
204
205 scope.push();
206 scope.insert(statement.mutate_identifier().mutate_name());
207
208 scope.process_scope(statement.mutate_block(), None);
209
210 Self::visit_block(statement.mutate_block(), scope);
211 scope.pop();
212 }
213
214 fn visit_repeat_statement(statement: &mut RepeatStatement, scope: &mut T) {
215 scope.process_repeat_statement(statement);
216
217 scope.push();
218
219 let (block, condition) = statement.mutate_block_and_condition();
220 scope.process_scope(block, Some(condition));
221
222 Self::visit_block_without_push(statement.mutate_block(), scope);
223 Self::visit_expression(statement.mutate_condition(), scope);
224
225 scope.pop();
226 }
227}
228
229#[derive(Debug, Clone, Default)]
230pub(crate) struct IdentifierTracker {
231 identifiers: Vec<HashSet<String>>,
232}
233
234impl IdentifierTracker {
235 fn insert_identifier(&mut self, identifier: &str) {
236 if let Some(set) = self.identifiers.last_mut() {
237 set.insert(identifier.to_string());
238 } else {
239 let mut set = HashSet::new();
240 set.insert(identifier.to_string());
241 self.identifiers.push(set);
242 }
243 }
244
245 pub fn new() -> IdentifierTracker {
246 Self {
247 identifiers: Vec::new(),
248 }
249 }
250
251 pub fn is_identifier_used(&self, identifier: &str) -> bool {
252 self.identifiers.iter().any(|set| set.contains(identifier))
253 }
254
255 pub fn generate_identifier(&mut self) -> String {
256 let mut permutator = identifier_permutator();
257
258 let identifier = permutator
259 .find(|identifier| {
260 is_valid_identifier(identifier) && !self.is_identifier_used(identifier)
261 })
262 .expect("the permutator should always ultimately return a valid identifier");
263 self.insert_identifier(&identifier);
264 identifier
265 }
266
267 pub fn generate_identifier_with_prefix(&mut self, prefix: impl Into<String>) -> String {
268 let mut identifier = prefix.into();
269 if identifier.is_empty() {
270 return self.generate_identifier();
271 }
272 let initial_length = identifier.len();
273 let mut permutator = Permutator::new("012345689".chars());
274
275 while self.is_identifier_used(&identifier) {
276 identifier.truncate(initial_length);
277 let next_suffix = permutator.next().unwrap_or_else(|| "_".to_owned());
278 identifier.push_str(&next_suffix);
279 }
280 self.insert_identifier(&identifier);
281 identifier
282 }
283}
284
285impl Scope for IdentifierTracker {
286 fn push(&mut self) {
287 self.identifiers.push(HashSet::new())
288 }
289
290 fn pop(&mut self) {
291 self.identifiers.pop();
292 }
293
294 fn insert(&mut self, identifier: &mut String) {
295 self.insert_identifier(identifier);
296 }
297
298 fn insert_self(&mut self) {
299 self.insert_identifier("self");
300 }
301
302 fn insert_local(&mut self, identifier: &mut String, _value: Option<&mut Expression>) {
303 self.insert_identifier(identifier);
304 }
305
306 fn insert_local_function(&mut self, function: &mut LocalFunctionStatement) {
307 self.insert_identifier(function.mutate_identifier().get_name());
308 }
309}
310
311impl<T, U> Scope for T
313where
314 T: DerefMut<Target = U>,
315 U: Scope,
316{
317 #[inline]
318 fn push(&mut self) {
319 self.deref_mut().push()
320 }
321
322 #[inline]
323 fn pop(&mut self) {
324 self.deref_mut().pop()
325 }
326
327 #[inline]
328 fn insert(&mut self, identifier: &mut String) {
329 self.deref_mut().insert(identifier);
330 }
331
332 #[inline]
333 fn insert_self(&mut self) {
334 self.deref_mut().insert_self();
335 }
336
337 #[inline]
338 fn insert_local(&mut self, identifier: &mut String, value: Option<&mut Expression>) {
339 self.deref_mut().insert_local(identifier, value)
340 }
341
342 #[inline]
343 fn insert_local_function(&mut self, function: &mut LocalFunctionStatement) {
344 self.deref_mut().insert_local_function(function)
345 }
346}