darklua_core/nodes/statements/
local_function.rs

1use crate::nodes::{
2    Block, FunctionBodyTokens, FunctionReturnType, FunctionVariadicType, GenericParameters,
3    Identifier, Token, TypedIdentifier,
4};
5
6/// Tokens associated with a local function statement.
7#[derive(Clone, Debug, PartialEq, Eq)]
8pub struct LocalFunctionTokens {
9    pub local: Token,
10    pub function_body: FunctionBodyTokens,
11}
12
13impl LocalFunctionTokens {
14    super::impl_token_fns!(target = [local, function_body]);
15}
16
17impl std::ops::Deref for LocalFunctionTokens {
18    type Target = FunctionBodyTokens;
19
20    fn deref(&self) -> &Self::Target {
21        &self.function_body
22    }
23}
24
25impl std::ops::DerefMut for LocalFunctionTokens {
26    fn deref_mut(&mut self) -> &mut Self::Target {
27        &mut self.function_body
28    }
29}
30
31/// Represents a local function declaration statement.
32#[derive(Clone, Debug, PartialEq, Eq)]
33pub struct LocalFunctionStatement {
34    identifier: Identifier,
35    block: Block,
36    parameters: Vec<TypedIdentifier>,
37    is_variadic: bool,
38    variadic_type: Option<FunctionVariadicType>,
39    return_type: Option<FunctionReturnType>,
40    generic_parameters: Option<GenericParameters>,
41    tokens: Option<Box<LocalFunctionTokens>>,
42}
43
44impl LocalFunctionStatement {
45    /// Creates a new local function statement.
46    pub fn new(
47        identifier: impl Into<Identifier>,
48        block: Block,
49        parameters: Vec<TypedIdentifier>,
50        is_variadic: bool,
51    ) -> Self {
52        Self {
53            identifier: identifier.into(),
54            block,
55            parameters,
56            is_variadic,
57            variadic_type: None,
58            return_type: None,
59            generic_parameters: None,
60            tokens: None,
61        }
62    }
63
64    /// Creates a new local function statement with a given name and block.
65    pub fn from_name(identifier: impl Into<Identifier>, block: impl Into<Block>) -> Self {
66        Self {
67            identifier: identifier.into(),
68            block: block.into(),
69            parameters: Vec::new(),
70            is_variadic: false,
71            variadic_type: None,
72            return_type: None,
73            generic_parameters: None,
74            tokens: None,
75        }
76    }
77
78    /// Sets the tokens for this local function statement.
79    pub fn with_tokens(mut self, tokens: LocalFunctionTokens) -> Self {
80        self.tokens = Some(tokens.into());
81        self
82    }
83
84    /// Sets the tokens for this local function statement.
85    #[inline]
86    pub fn set_tokens(&mut self, tokens: LocalFunctionTokens) {
87        self.tokens = Some(tokens.into());
88    }
89
90    /// Returns the tokens for this local function statement, if any.
91    #[inline]
92    pub fn get_tokens(&self) -> Option<&LocalFunctionTokens> {
93        self.tokens.as_deref()
94    }
95
96    /// Returns a mutable reference to the tokens, if any.
97    #[inline]
98    pub fn mutate_tokens(&mut self) -> Option<&mut LocalFunctionTokens> {
99        self.tokens.as_deref_mut()
100    }
101
102    /// Adds a parameter to this function.
103    pub fn with_parameter(mut self, parameter: impl Into<TypedIdentifier>) -> Self {
104        self.parameters.push(parameter.into());
105        self
106    }
107
108    /// Marks this function as variadic.
109    pub fn variadic(mut self) -> Self {
110        self.is_variadic = true;
111        self
112    }
113
114    /// Sets the variadic type for this function.
115    ///
116    /// If the function is not already variadic, this will make it variadic.
117    pub fn with_variadic_type(mut self, r#type: impl Into<FunctionVariadicType>) -> Self {
118        self.is_variadic = true;
119        self.variadic_type = Some(r#type.into());
120        self
121    }
122
123    /// Sets the variadic type for this function.
124    ///
125    /// If the function is not already variadic, this will make it variadic.
126    pub fn set_variadic_type(&mut self, r#type: impl Into<FunctionVariadicType>) {
127        self.is_variadic = true;
128        self.variadic_type = Some(r#type.into());
129    }
130
131    /// Returns the variadic type, if any.
132    #[inline]
133    pub fn get_variadic_type(&self) -> Option<&FunctionVariadicType> {
134        self.variadic_type.as_ref()
135    }
136
137    /// Returns whether this function has a variadic type.
138    #[inline]
139    pub fn has_variadic_type(&self) -> bool {
140        self.variadic_type.is_some()
141    }
142
143    /// Returns a mutable reference to the variadic type, if any.
144    #[inline]
145    pub fn mutate_variadic_type(&mut self) -> Option<&mut FunctionVariadicType> {
146        self.variadic_type.as_mut()
147    }
148
149    /// Sets the return type for this function.
150    pub fn with_return_type(mut self, return_type: impl Into<FunctionReturnType>) -> Self {
151        self.return_type = Some(return_type.into());
152        self
153    }
154
155    /// Sets the return type for this function.
156    pub fn set_return_type(&mut self, return_type: impl Into<FunctionReturnType>) {
157        self.return_type = Some(return_type.into());
158    }
159
160    /// Returns the return type, if any.
161    #[inline]
162    pub fn get_return_type(&self) -> Option<&FunctionReturnType> {
163        self.return_type.as_ref()
164    }
165
166    /// Returns whether this function has a return type.
167    #[inline]
168    pub fn has_return_type(&self) -> bool {
169        self.return_type.is_some()
170    }
171
172    /// Returns a mutable reference to the return type, if any.
173    #[inline]
174    pub fn mutate_return_type(&mut self) -> Option<&mut FunctionReturnType> {
175        self.return_type.as_mut()
176    }
177
178    /// Sets the generic parameters for this function.
179    pub fn with_generic_parameters(mut self, generic_parameters: GenericParameters) -> Self {
180        self.generic_parameters = Some(generic_parameters);
181        self
182    }
183
184    /// Sets the generic parameters for this function.
185    #[inline]
186    pub fn set_generic_parameters(&mut self, generic_parameters: GenericParameters) {
187        self.generic_parameters = Some(generic_parameters);
188    }
189
190    /// Returns the generic parameters, if any.
191    #[inline]
192    pub fn get_generic_parameters(&self) -> Option<&GenericParameters> {
193        self.generic_parameters.as_ref()
194    }
195
196    /// Returns a mutable reference to the parameters.
197    #[inline]
198    pub fn mutate_parameters(&mut self) -> &mut Vec<TypedIdentifier> {
199        &mut self.parameters
200    }
201
202    /// Returns a mutable reference to the block.
203    #[inline]
204    pub fn mutate_block(&mut self) -> &mut Block {
205        &mut self.block
206    }
207
208    /// Returns a mutable reference to the identifier.
209    #[inline]
210    pub fn mutate_identifier(&mut self) -> &mut Identifier {
211        &mut self.identifier
212    }
213
214    /// Returns the function's block.
215    #[inline]
216    pub fn get_block(&self) -> &Block {
217        &self.block
218    }
219
220    /// Returns the function's parameters.
221    #[inline]
222    pub fn get_parameters(&self) -> &Vec<TypedIdentifier> {
223        &self.parameters
224    }
225
226    /// Returns an iterator over the parameters.
227    #[inline]
228    pub fn iter_parameters(&self) -> impl Iterator<Item = &TypedIdentifier> {
229        self.parameters.iter()
230    }
231
232    /// Returns a mutable iterator over the parameters.
233    #[inline]
234    pub fn iter_mut_parameters(&mut self) -> impl Iterator<Item = &mut TypedIdentifier> {
235        self.parameters.iter_mut()
236    }
237
238    /// Returns the function's identifier.
239    #[inline]
240    pub fn get_identifier(&self) -> &Identifier {
241        &self.identifier
242    }
243
244    /// Returns the function's name.
245    #[inline]
246    pub fn get_name(&self) -> &str {
247        self.identifier.get_name()
248    }
249
250    /// Returns whether this function has a parameter with the given name.
251    #[inline]
252    pub fn has_parameter(&self, name: &str) -> bool {
253        self.parameters
254            .iter()
255            .any(|parameter| parameter.get_name() == name)
256    }
257
258    /// Returns whether this function has parameters.
259    #[inline]
260    pub fn has_parameters(&self) -> bool {
261        !self.parameters.is_empty()
262    }
263
264    /// Returns whether this function is variadic.
265    #[inline]
266    pub fn is_variadic(&self) -> bool {
267        self.is_variadic
268    }
269
270    /// Returns the number of parameters.
271    #[inline]
272    pub fn parameters_count(&self) -> usize {
273        self.parameters.len()
274    }
275
276    /// Removes all type annotations from this function.
277    pub fn clear_types(&mut self) {
278        self.return_type.take();
279        self.variadic_type.take();
280        self.generic_parameters.take();
281        for parameter in &mut self.parameters {
282            parameter.remove_type();
283        }
284        if let Some(tokens) = &mut self.tokens {
285            tokens.variable_arguments_colon.take();
286        }
287    }
288
289    /// Returns a mutable reference to the first token for this statement, creating it if missing.
290    pub fn mutate_first_token(&mut self) -> &mut Token {
291        self.set_default_tokens();
292        &mut self.tokens.as_deref_mut().unwrap().local
293    }
294
295    /// Returns a mutable reference to the last token for this statement,
296    /// creating it if missing.
297    pub fn mutate_last_token(&mut self) -> &mut Token {
298        self.set_default_tokens();
299        &mut self.tokens.as_deref_mut().unwrap().end
300    }
301
302    fn set_default_tokens(&mut self) {
303        if self.tokens.is_none() {
304            self.tokens = Some(
305                LocalFunctionTokens {
306                    local: Token::from_content("local"),
307                    function_body: FunctionBodyTokens {
308                        function: Token::from_content("function"),
309                        opening_parenthese: Token::from_content("("),
310                        closing_parenthese: Token::from_content(")"),
311                        end: Token::from_content("end"),
312                        parameter_commas: Vec::new(),
313                        variable_arguments: None,
314                        variable_arguments_colon: None,
315                        return_type_colon: None,
316                    },
317                }
318                .into(),
319            );
320        }
321    }
322
323    super::impl_token_fns!(
324        target = [identifier]
325        iter = [parameters, generic_parameters, tokens]
326    );
327}
328
329#[cfg(test)]
330mod test {
331    use super::*;
332
333    #[test]
334    fn has_parameter_is_true_when_single_param_matches() {
335        let func = LocalFunctionStatement::from_name("foo", Block::default()).with_parameter("bar");
336
337        assert!(func.has_parameter("bar"));
338    }
339
340    #[test]
341    fn has_parameter_is_true_when_at_least_one_param_matches() {
342        let func = LocalFunctionStatement::from_name("foo", Block::default())
343            .with_parameter("bar")
344            .with_parameter("baz");
345
346        assert!(func.has_parameter("baz"));
347    }
348
349    #[test]
350    fn has_parameter_is_false_when_none_matches() {
351        let func = LocalFunctionStatement::from_name("foo", Block::default())
352            .with_parameter("bar")
353            .with_parameter("baz");
354
355        assert!(!func.has_parameter("foo"));
356    }
357}
darklua_core/nodes/statements/local_function.rs

darklua_core/nodes/statements/
local_function.rs
