darklua_core/process/
scope_visitor.rs

1use 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
10/// Defines methods to interact with the concept of lexical scoping. The struct implementing this
11/// trait should be able to keep track of identifiers when used along the ScopeVisitor.
12pub trait Scope {
13    /// This method is called when a new block is entered.
14    fn push(&mut self);
15    /// When a block is left, this method should should free all identifiers inserted in the
16    /// previous block.
17    fn pop(&mut self);
18    /// Called when entering a function block (with each parameters of the function), with the
19    /// identifiers from a generic for statement or the identifier from a numeric for loop.
20    fn insert(&mut self, identifier: &mut String);
21    /// Called when entering a function defined with a method
22    fn insert_self(&mut self);
23    /// Called when a new local variable is initialized.
24    fn insert_local(&mut self, identifier: &mut String, value: Option<&mut Expression>);
25    /// Called when a new local function is initialized.
26    fn insert_local_function(&mut self, function: &mut LocalFunctionStatement);
27}
28
29/// A visitor that can be used only with a NodeProcessor that also implements the Scope trait.
30pub 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
311// implement Scope on anything that can deref into a Scope
312impl<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}