darklua_core/nodes/statements/
local_function.rs

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

darklua_core/nodes/statements/
local_function.rs
