darklua_core/nodes/statements/
local_function.rs

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

darklua_core/nodes/statements/
local_function.rs
