darklua_core/nodes/statements/
function.rs

1use crate::nodes::{
2    Attribute, Attributes, Block, FunctionBodyTokens, FunctionReturnType, FunctionVariadicType,
3    GenericParameters, Identifier, Token, TypedIdentifier,
4};
5
6/// Tokens associated with a function name.
7#[derive(Clone, Debug, PartialEq, Eq)]
8pub struct FunctionNameTokens {
9    /// The tokens for the periods in the function name.
10    pub periods: Vec<Token>,
11    /// The token for the colon in the function name when a method is present.
12    pub colon: Option<Token>,
13}
14
15impl FunctionNameTokens {
16    super::impl_token_fns!(iter = [periods, colon]);
17}
18
19/// Represents the name portion of a function statement.
20///
21/// Function names can include table fields and methods
22/// ([e.g., `module.table:method`]).
23#[derive(Clone, Debug, PartialEq, Eq)]
24pub struct FunctionName {
25    name: Identifier,
26    field_names: Vec<Identifier>,
27    method: Option<Identifier>,
28    tokens: Option<FunctionNameTokens>,
29}
30
31impl FunctionName {
32    /// Creates a new function name with fields and an optional method.
33    pub fn new(name: Identifier, field_names: Vec<Identifier>, method: Option<Identifier>) -> Self {
34        Self {
35            name,
36            field_names,
37            method,
38            tokens: None,
39        }
40    }
41
42    /// Creates a new function name from a single identifier.
43    pub fn from_name<S: Into<Identifier>>(name: S) -> Self {
44        Self {
45            name: name.into(),
46            field_names: Vec::new(),
47            method: None,
48            tokens: None,
49        }
50    }
51
52    /// Sets the tokens for this function name.
53    pub fn with_tokens(mut self, tokens: FunctionNameTokens) -> Self {
54        self.tokens = Some(tokens);
55        self
56    }
57
58    /// Sets the tokens for this function name.
59    #[inline]
60    pub fn set_tokens(&mut self, tokens: FunctionNameTokens) {
61        self.tokens = Some(tokens);
62    }
63
64    /// Returns the tokens for this function name, if any.
65    #[inline]
66    pub fn get_tokens(&self) -> Option<&FunctionNameTokens> {
67        self.tokens.as_ref()
68    }
69
70    /// Adds a field to this function name.
71    pub fn with_field<S: Into<Identifier>>(mut self, field: S) -> Self {
72        self.field_names.push(field.into());
73        self
74    }
75
76    /// Sets the field names for this function name.
77    pub fn with_fields(mut self, field_names: Vec<Identifier>) -> Self {
78        self.field_names = field_names;
79        self
80    }
81
82    /// Sets a method for this function name.
83    pub fn with_method<S: Into<Identifier>>(mut self, method: S) -> Self {
84        self.method.replace(method.into());
85        self
86    }
87
88    /// Adds a field to this function name.
89    pub fn push_field<S: Into<Identifier>>(&mut self, field: S) {
90        self.field_names.push(field.into());
91    }
92
93    /// Removes and returns the method, if any.
94    #[inline]
95    pub fn remove_method(&mut self) -> Option<Identifier> {
96        self.method.take()
97    }
98
99    /// Returns the method, if any.
100    #[inline]
101    pub fn get_method(&self) -> Option<&Identifier> {
102        self.method.as_ref()
103    }
104
105    /// Returns whether this function name has a method component.
106    #[inline]
107    pub fn has_method(&self) -> bool {
108        self.method.is_some()
109    }
110
111    /// Returns the base name.
112    #[inline]
113    pub fn get_name(&self) -> &Identifier {
114        &self.name
115    }
116
117    /// Sets the base name.
118    #[inline]
119    pub fn set_name(&mut self, name: Identifier) {
120        self.name = name;
121    }
122
123    /// Returns the field names.
124    #[inline]
125    pub fn get_field_names(&self) -> &Vec<Identifier> {
126        &self.field_names
127    }
128
129    /// Returns a mutable reference to the base identifier.
130    #[inline]
131    pub fn mutate_identifier(&mut self) -> &mut Identifier {
132        &mut self.name
133    }
134
135    super::impl_token_fns!(iter = [tokens, field_names, method]);
136}
137
138/// Represents a function declaration statement.
139#[derive(Clone, Debug, PartialEq, Eq)]
140pub struct FunctionStatement {
141    name: FunctionName,
142    block: Block,
143    parameters: Vec<TypedIdentifier>,
144    is_variadic: bool,
145    variadic_type: Option<FunctionVariadicType>,
146    return_type: Option<FunctionReturnType>,
147    generic_parameters: Option<GenericParameters>,
148    attributes: Attributes,
149    tokens: Option<Box<FunctionBodyTokens>>,
150}
151
152impl FunctionStatement {
153    /// Creates a new function statement with the given name, block, parameters, and variadic flag.
154    pub fn new(
155        name: FunctionName,
156        block: Block,
157        parameters: Vec<TypedIdentifier>,
158        is_variadic: bool,
159    ) -> Self {
160        Self {
161            name,
162            block,
163            parameters,
164            is_variadic,
165            variadic_type: None,
166            return_type: None,
167            generic_parameters: None,
168            attributes: Attributes::new(),
169            tokens: None,
170        }
171    }
172
173    /// Creates a new function statement from a single identifier and a block.
174    pub fn from_name<S: Into<String>, B: Into<Block>>(name: S, block: B) -> Self {
175        Self {
176            name: FunctionName::from_name(name),
177            block: block.into(),
178            parameters: Vec::new(),
179            is_variadic: false,
180            variadic_type: None,
181            return_type: None,
182            generic_parameters: None,
183            attributes: Attributes::new(),
184            tokens: None,
185        }
186    }
187
188    /// Returns a reference to the attributes of this function statement.
189    pub fn attributes(&self) -> &Attributes {
190        &self.attributes
191    }
192
193    /// Associates attributes with this function statement. Replaces any existing attributes.
194    pub fn with_attributes(mut self, attributes: Attributes) -> Self {
195        self.attributes = attributes;
196        self
197    }
198
199    /// Returns a mutable reference to the attributes of this function statement.
200    pub fn mutate_attributes(&mut self) -> &mut Attributes {
201        &mut self.attributes
202    }
203
204    /// Adds an attribute to this function statement. Adds to any existing attributes.
205    pub fn with_attribute(mut self, attribute: impl Into<Attribute>) -> Self {
206        self.attributes.append_attribute(attribute.into());
207        self
208    }
209
210    /// Sets the tokens for this function statement.
211    pub fn with_tokens(mut self, tokens: FunctionBodyTokens) -> Self {
212        self.tokens = Some(tokens.into());
213        self
214    }
215
216    /// Sets the tokens for this function statement.
217    #[inline]
218    pub fn set_tokens(&mut self, tokens: FunctionBodyTokens) {
219        self.tokens = Some(tokens.into());
220    }
221
222    /// Returns the tokens for this function statement, if any.
223    #[inline]
224    pub fn get_tokens(&self) -> Option<&FunctionBodyTokens> {
225        self.tokens.as_deref()
226    }
227
228    /// Returns a mutable reference to the tokens, if any.
229    #[inline]
230    pub fn mutate_tokens(&mut self) -> Option<&mut FunctionBodyTokens> {
231        self.tokens.as_deref_mut()
232    }
233
234    /// Adds a parameter to this function.
235    pub fn with_parameter(mut self, parameter: impl Into<TypedIdentifier>) -> Self {
236        self.parameters.push(parameter.into());
237        self
238    }
239
240    /// Marks this function as variadic.
241    pub fn variadic(mut self) -> Self {
242        self.is_variadic = true;
243        self
244    }
245
246    /// Sets the variadic type for this function.
247    pub fn with_variadic_type(mut self, r#type: impl Into<FunctionVariadicType>) -> Self {
248        self.is_variadic = true;
249        self.variadic_type = Some(r#type.into());
250        self
251    }
252
253    /// Sets the variadic type for this function.
254    pub fn set_variadic_type(&mut self, r#type: impl Into<FunctionVariadicType>) {
255        self.is_variadic = true;
256        self.variadic_type = Some(r#type.into());
257    }
258
259    /// Returns the variadic type, if any.
260    #[inline]
261    pub fn get_variadic_type(&self) -> Option<&FunctionVariadicType> {
262        self.variadic_type.as_ref()
263    }
264
265    /// Returns whether this function has a variadic type.
266    #[inline]
267    pub fn has_variadic_type(&self) -> bool {
268        self.variadic_type.is_some()
269    }
270
271    /// Returns a mutable reference to the variadic type, if any.
272    #[inline]
273    pub fn mutate_variadic_type(&mut self) -> Option<&mut FunctionVariadicType> {
274        self.variadic_type.as_mut()
275    }
276
277    /// Sets the return type for this function.
278    pub fn with_return_type(mut self, return_type: impl Into<FunctionReturnType>) -> Self {
279        self.return_type = Some(return_type.into());
280        self
281    }
282
283    /// Sets the return type for this function.
284    pub fn set_return_type(&mut self, return_type: impl Into<FunctionReturnType>) {
285        self.return_type = Some(return_type.into());
286    }
287
288    /// Returns the return type, if any.
289    #[inline]
290    pub fn get_return_type(&self) -> Option<&FunctionReturnType> {
291        self.return_type.as_ref()
292    }
293
294    /// Returns whether this function has a return type.
295    #[inline]
296    pub fn has_return_type(&self) -> bool {
297        self.return_type.is_some()
298    }
299
300    /// Returns a mutable reference to the return type, if any.
301    #[inline]
302    pub fn mutate_return_type(&mut self) -> Option<&mut FunctionReturnType> {
303        self.return_type.as_mut()
304    }
305
306    /// Sets the generic parameters for this function.
307    pub fn with_generic_parameters(mut self, generic_parameters: GenericParameters) -> Self {
308        self.generic_parameters = Some(generic_parameters);
309        self
310    }
311
312    /// Sets the generic parameters for this function.
313    #[inline]
314    pub fn set_generic_parameters(&mut self, generic_parameters: GenericParameters) {
315        self.generic_parameters = Some(generic_parameters);
316    }
317
318    /// Returns the generic parameters, if any.
319    #[inline]
320    pub fn get_generic_parameters(&self) -> Option<&GenericParameters> {
321        self.generic_parameters.as_ref()
322    }
323
324    /// Returns a reference to the function's block.
325    #[inline]
326    pub fn get_block(&self) -> &Block {
327        &self.block
328    }
329
330    /// Returns a reference to the function's name.
331    #[inline]
332    pub fn get_name(&self) -> &FunctionName {
333        &self.name
334    }
335
336    /// Returns the number of parameters in this function.
337    #[inline]
338    pub fn parameters_count(&self) -> usize {
339        self.parameters.len()
340    }
341
342    /// Returns a reference to the function's parameters.
343    #[inline]
344    pub fn get_parameters(&self) -> &Vec<TypedIdentifier> {
345        &self.parameters
346    }
347
348    /// Returns an iterator over the function's parameters.
349    #[inline]
350    pub fn iter_parameters(&self) -> impl Iterator<Item = &TypedIdentifier> {
351        self.parameters.iter()
352    }
353
354    /// Returns a mutable iterator over the function's parameters.
355    #[inline]
356    pub fn iter_mut_parameters(&mut self) -> impl Iterator<Item = &mut TypedIdentifier> {
357        self.parameters.iter_mut()
358    }
359
360    /// Returns whether this function is variadic.
361    #[inline]
362    pub fn is_variadic(&self) -> bool {
363        self.is_variadic
364    }
365
366    /// Returns a mutable reference to the function's block.
367    #[inline]
368    pub fn mutate_block(&mut self) -> &mut Block {
369        &mut self.block
370    }
371
372    /// Returns a mutable reference to the function's name.
373    #[inline]
374    pub fn mutate_function_name(&mut self) -> &mut FunctionName {
375        &mut self.name
376    }
377
378    /// Returns a mutable reference to the function's parameters.
379    #[inline]
380    pub fn mutate_parameters(&mut self) -> &mut Vec<TypedIdentifier> {
381        &mut self.parameters
382    }
383
384    /// Removes the method from the function name and adds it as a
385    /// field, inserting 'self' as the first parameter.
386    ///
387    /// If the method is not present, this does nothing.
388    pub fn remove_method(&mut self) {
389        if let Some(method_name) = self.name.remove_method() {
390            self.name.push_field(method_name);
391            self.parameters.insert(0, TypedIdentifier::new("self"));
392        }
393    }
394
395    /// Returns whether this function has any parameters.
396    #[inline]
397    pub fn has_parameters(&self) -> bool {
398        !self.parameters.is_empty()
399    }
400
401    /// Removes all type information from the function, including return
402    /// type, variadic type, generic parameters, and parameter types.
403    pub fn clear_types(&mut self) {
404        self.return_type.take();
405        self.variadic_type.take();
406        self.generic_parameters.take();
407        for parameter in &mut self.parameters {
408            parameter.remove_type();
409        }
410        if let Some(tokens) = &mut self.tokens {
411            tokens.variable_arguments_colon.take();
412        }
413    }
414
415    /// Returns a mutable reference to the first token for this statement, creating it if missing.
416    pub fn mutate_first_token(&mut self) -> &mut Token {
417        self.set_default_tokens();
418        &mut self.tokens.as_deref_mut().unwrap().function
419    }
420
421    /// Returns a mutable reference to the last token for this statement,
422    /// creating it if missing.
423    pub fn mutate_last_token(&mut self) -> &mut Token {
424        self.set_default_tokens();
425        &mut self.tokens.as_deref_mut().unwrap().end
426    }
427
428    fn set_default_tokens(&mut self) {
429        if self.tokens.is_none() {
430            self.tokens = Some(
431                FunctionBodyTokens {
432                    function: Token::from_content("function"),
433                    opening_parenthese: Token::from_content("("),
434                    closing_parenthese: Token::from_content(")"),
435                    end: Token::from_content("end"),
436                    parameter_commas: Vec::new(),
437                    variable_arguments: None,
438                    variable_arguments_colon: None,
439                    return_type_colon: None,
440                }
441                .into(),
442            );
443        }
444    }
445
446    super::impl_token_fns!(
447        target = [name]
448        iter = [parameters, generic_parameters, tokens]
449    );
450}
darklua_core/nodes/statements/function.rs

darklua_core/nodes/statements/
function.rs
