Skip to main content

darklua_core/process/processors/
collect_globals.rs

1use std::collections::HashSet;
2
3use crate::{
4    nodes::{Expression, FunctionAssignment, Identifier, TypeField},
5    process::{NodeProcessor, Scope},
6};
7
8#[derive(Debug, Default)]
9pub struct CollectGlobalsProcessor {
10    scopes: Vec<HashSet<String>>,
11    globals: HashSet<String>,
12}
13
14impl CollectGlobalsProcessor {
15    pub(crate) fn add(&mut self, identifier: &str) {
16        if self.scopes.is_empty() {
17            self.scopes.push(HashSet::new());
18        }
19
20        let current = self.scopes.last_mut().unwrap();
21        if !current.contains(identifier) {
22            current.insert(identifier.to_owned());
23        }
24    }
25
26    pub fn is_declared(&self, identifier: &str) -> bool {
27        self.scopes
28            .iter()
29            .rev()
30            .any(|scope| scope.contains(identifier))
31    }
32
33    pub fn iter_globals(&self) -> impl Iterator<Item = &str> {
34        self.globals.iter().map(String::as_str)
35    }
36
37    pub fn into_globals(self) -> impl Iterator<Item = String> {
38        self.globals.into_iter()
39    }
40}
41
42impl Scope for CollectGlobalsProcessor {
43    fn push(&mut self) {
44        self.scopes.push(HashSet::new())
45    }
46
47    fn pop(&mut self) {
48        self.scopes.pop();
49    }
50
51    fn insert(&mut self, identifier: &mut String) {
52        self.add(identifier);
53    }
54
55    fn insert_self(&mut self) {
56        self.add("self");
57    }
58
59    fn insert_local(&mut self, identifier: &mut String, _value: Option<&mut Expression>) {
60        self.add(identifier);
61    }
62
63    fn insert_local_function(&mut self, function: &mut FunctionAssignment) {
64        self.add(function.get_name());
65    }
66}
67
68impl NodeProcessor for CollectGlobalsProcessor {
69    fn process_variable_expression(&mut self, variable: &mut Identifier) {
70        if !self.is_declared(variable.get_name()) {
71            self.globals.insert(variable.get_name().clone());
72        }
73    }
74
75    fn process_type_field(&mut self, type_field: &mut TypeField) {
76        let namespace = type_field.get_namespace().get_name();
77        if !self.is_declared(namespace) {
78            self.globals.insert(namespace.clone());
79        }
80    }
81}
82
83#[cfg(test)]
84mod test {
85    use crate::{
86        process::{NodeVisitor, ScopeVisitor},
87        Parser,
88    };
89
90    use super::*;
91
92    fn extract_globals(code: &str) -> Vec<String> {
93        let mut processor = CollectGlobalsProcessor::default();
94
95        let mut block = Parser::default()
96            .parse(code)
97            .expect("expected code should parse");
98
99        ScopeVisitor::visit_block(&mut block, &mut processor);
100
101        let mut globals = processor.into_globals().collect::<Vec<_>>();
102        globals.sort();
103        globals
104    }
105
106    #[test]
107    fn catch_global_variable_in_module_scope() {
108        insta::assert_debug_snapshot!(extract_globals(r#"
109local g = game
110return g
111        "#), @r###"
112        [
113            "game",
114        ]
115        "###);
116    }
117
118    #[test]
119    fn catch_global_variable_within_function_scope() {
120        insta::assert_debug_snapshot!(extract_globals(r#"
121local function example()
122    return unpack(game:GetService("Players"):GetPlayers())
123end
124
125return { example = example }
126        "#), @r###"
127        [
128            "game",
129            "unpack",
130        ]
131        "###);
132    }
133}