darklua_core/nodes/statements/
local_assign.rs

1use crate::nodes::{Expression, Token, TypedIdentifier};
2
3/// Tokens associated with a local variable assignment statement.
4#[derive(Clone, Debug, PartialEq, Eq)]
5pub struct LocalAssignTokens {
6    pub local: Token,
7    /// The token for the equal sign, if any.
8    pub equal: Option<Token>,
9    /// The tokens for the commas between variables.
10    pub variable_commas: Vec<Token>,
11    /// The tokens for the commas between values.
12    pub value_commas: Vec<Token>,
13}
14
15impl LocalAssignTokens {
16    super::impl_token_fns!(
17        target = [local]
18        iter = [variable_commas, value_commas, equal]
19    );
20}
21
22/// Represents a local variable assignment statement.
23#[derive(Clone, Debug, PartialEq, Eq)]
24pub struct LocalAssignStatement {
25    variables: Vec<TypedIdentifier>,
26    values: Vec<Expression>,
27    tokens: Option<LocalAssignTokens>,
28}
29
30impl LocalAssignStatement {
31    /// Creates a new local assignment statement with the given variables and values.
32    pub fn new(variables: Vec<TypedIdentifier>, values: Vec<Expression>) -> Self {
33        Self {
34            variables,
35            values,
36            tokens: None,
37        }
38    }
39
40    /// Creates a new local assignment statement with a single variable and no values.
41    pub fn from_variable<S: Into<TypedIdentifier>>(variable: S) -> Self {
42        Self {
43            variables: vec![variable.into()],
44            values: Vec::new(),
45            tokens: None,
46        }
47    }
48
49    /// Sets the tokens for this local assignment statement.
50    pub fn with_tokens(mut self, tokens: LocalAssignTokens) -> Self {
51        self.tokens = Some(tokens);
52        self
53    }
54
55    /// Sets the tokens for this local assignment statement.
56    #[inline]
57    pub fn set_tokens(&mut self, tokens: LocalAssignTokens) {
58        self.tokens = Some(tokens);
59    }
60
61    /// Returns the tokens for this local assignment statement, if any.
62    #[inline]
63    pub fn get_tokens(&self) -> Option<&LocalAssignTokens> {
64        self.tokens.as_ref()
65    }
66
67    /// Returns a mutable reference to the tokens, if any.
68    #[inline]
69    pub fn mutate_tokens(&mut self) -> Option<&mut LocalAssignTokens> {
70        self.tokens.as_mut()
71    }
72
73    /// Adds a variable to this local assignment statement.
74    pub fn with_variable<S: Into<TypedIdentifier>>(mut self, variable: S) -> Self {
75        self.variables.push(variable.into());
76        self
77    }
78
79    /// Adds a value to this local assignment statement.
80    pub fn with_value<E: Into<Expression>>(mut self, value: E) -> Self {
81        self.values.push(value.into());
82        self
83    }
84
85    /// Converts this statement into a tuple of variables and values.
86    pub fn into_assignments(self) -> (Vec<TypedIdentifier>, Vec<Expression>) {
87        (self.variables, self.values)
88    }
89
90    /// Adds a new variable-value pair to this local assignment statement.
91    pub fn append_assignment<S: Into<TypedIdentifier>>(&mut self, variable: S, value: Expression) {
92        self.variables.push(variable.into());
93        self.values.push(value);
94    }
95
96    /// Applies a function to each variable-value pair.
97    pub fn for_each_assignment<F>(&mut self, mut callback: F)
98    where
99        F: FnMut(&mut TypedIdentifier, Option<&mut Expression>),
100    {
101        let mut values = self.values.iter_mut();
102        self.variables
103            .iter_mut()
104            .for_each(|variable| callback(variable, values.next()));
105    }
106
107    /// Returns the list of variables.
108    #[inline]
109    pub fn get_variables(&self) -> &Vec<TypedIdentifier> {
110        &self.variables
111    }
112
113    /// Returns an iterator over the variables.
114    #[inline]
115    pub fn iter_variables(&self) -> impl Iterator<Item = &TypedIdentifier> {
116        self.variables.iter()
117    }
118
119    /// Returns a mutable iterator over the variables.
120    #[inline]
121    pub fn iter_mut_variables(&mut self) -> impl Iterator<Item = &mut TypedIdentifier> {
122        self.variables.iter_mut()
123    }
124
125    /// Appends variables from another vector.
126    #[inline]
127    pub fn append_variables(&mut self, variables: &mut Vec<TypedIdentifier>) {
128        self.variables.append(variables);
129    }
130
131    /// Extends the values with elements from an iterator.
132    #[inline]
133    pub fn extend_values<T: IntoIterator<Item = Expression>>(&mut self, iter: T) {
134        self.values.extend(iter);
135    }
136
137    /// Returns a mutable iterator over the values.
138    #[inline]
139    pub fn iter_mut_values(&mut self) -> impl Iterator<Item = &mut Expression> {
140        self.values.iter_mut()
141    }
142
143    /// Returns an iterator over the values.
144    #[inline]
145    pub fn iter_values(&self) -> impl Iterator<Item = &Expression> {
146        self.values.iter()
147    }
148
149    /// Adds a variable to this local assignment statement.
150    #[inline]
151    pub fn push_variable(&mut self, variable: impl Into<TypedIdentifier>) {
152        self.variables.push(variable.into());
153    }
154
155    /// Adds a value to this local assignment statement.
156    #[inline]
157    pub fn push_value(&mut self, value: impl Into<Expression>) {
158        self.values.push(value.into());
159    }
160
161    /// Appends values from another vector.
162    #[inline]
163    pub fn append_values(&mut self, values: &mut Vec<Expression>) {
164        self.values.append(values);
165    }
166
167    /// Returns the last value, if any.
168    #[inline]
169    pub fn last_value(&self) -> Option<&Expression> {
170        self.values.last()
171    }
172
173    /// Removes and returns the last value, adjusting tokens as needed.
174    pub fn pop_value(&mut self) -> Option<Expression> {
175        let value = self.values.pop();
176        if let Some(tokens) = &mut self.tokens {
177            let length = self.values.len();
178            if length == 0 {
179                if !tokens.value_commas.is_empty() {
180                    tokens.value_commas.clear();
181                }
182                if tokens.equal.is_some() {
183                    tokens.equal = None;
184                }
185            } else {
186                tokens.value_commas.truncate(length.saturating_sub(1));
187            }
188        }
189        value
190    }
191
192    /// Removes and returns the value at the given index, adjusting tokens as needed.
193    pub fn remove_value(&mut self, index: usize) -> Option<Expression> {
194        if index < self.values.len() {
195            let value = self.values.remove(index);
196
197            if let Some(tokens) = &mut self.tokens {
198                if index < tokens.value_commas.len() {
199                    tokens.value_commas.remove(index);
200                }
201                if self.values.is_empty() && tokens.equal.is_some() {
202                    tokens.equal = None;
203                }
204            }
205
206            Some(value)
207        } else {
208            None
209        }
210    }
211
212    /// Removes and returns the variable at the given index, adjusting tokens as needed.
213    ///
214    /// Returns None if there is only one variable or if the index is out of bounds.
215    pub fn remove_variable(&mut self, index: usize) -> Option<TypedIdentifier> {
216        let len = self.variables.len();
217
218        if len > 1 && index < len {
219            let variable = self.variables.remove(index);
220
221            if let Some(tokens) = &mut self.tokens {
222                if index < tokens.variable_commas.len() {
223                    tokens.variable_commas.remove(index);
224                }
225            }
226
227            Some(variable)
228        } else {
229            None
230        }
231    }
232
233    /// Returns the number of values.
234    #[inline]
235    pub fn values_len(&self) -> usize {
236        self.values.len()
237    }
238
239    /// Returns the number of variables.
240    #[inline]
241    pub fn variables_len(&self) -> usize {
242        self.variables.len()
243    }
244
245    /// Returns whether this statement has any values.
246    #[inline]
247    pub fn has_values(&self) -> bool {
248        !self.values.is_empty()
249    }
250
251    /// Removes type annotations from all variables.
252    pub fn clear_types(&mut self) {
253        for variable in &mut self.variables {
254            variable.remove_type();
255        }
256    }
257
258    /// Returns a mutable reference to the first token for this statement, creating it if missing.
259    pub fn mutate_first_token(&mut self) -> &mut Token {
260        if self.tokens.is_none() {
261            self.tokens = Some(LocalAssignTokens {
262                local: Token::from_content("local"),
263                equal: (!self.values.is_empty()).then(|| Token::from_content("=")),
264                variable_commas: Vec::new(),
265                value_commas: Vec::new(),
266            });
267        }
268        &mut self.tokens.as_mut().unwrap().local
269    }
270
271    /// Returns a mutable reference to the last token for this statement,
272    /// creating it if missing.
273    pub fn mutate_last_token(&mut self) -> &mut Token {
274        if let Some(last_value) = self.values.last_mut() {
275            return last_value.mutate_last_token();
276        }
277        self.variables
278            .last_mut()
279            .expect("local assign must have at least one variable")
280            .mutate_or_insert_token()
281    }
282
283    super::impl_token_fns!(iter = [variables, tokens]);
284}
285
286#[cfg(test)]
287mod test {
288    use super::*;
289
290    mod pop_value {
291        use super::*;
292
293        #[test]
294        fn removes_the_equal_sign() {
295            let mut assign = LocalAssignStatement::from_variable("var")
296                .with_value(true)
297                .with_tokens(LocalAssignTokens {
298                    local: Token::from_content("local"),
299                    equal: Some(Token::from_content("=")),
300                    variable_commas: Vec::new(),
301                    value_commas: Vec::new(),
302                });
303
304            assign.pop_value();
305
306            pretty_assertions::assert_eq!(
307                assign,
308                LocalAssignStatement::from_variable("var").with_tokens(LocalAssignTokens {
309                    local: Token::from_content("local"),
310                    equal: None,
311                    variable_commas: Vec::new(),
312                    value_commas: Vec::new(),
313                })
314            );
315        }
316
317        #[test]
318        fn removes_the_last_comma_token() {
319            let mut assign = LocalAssignStatement::from_variable("var")
320                .with_value(true)
321                .with_value(false)
322                .with_tokens(LocalAssignTokens {
323                    local: Token::from_content("local"),
324                    equal: Some(Token::from_content("=")),
325                    variable_commas: Vec::new(),
326                    value_commas: vec![Token::from_content(",")],
327                });
328
329            assign.pop_value();
330
331            pretty_assertions::assert_eq!(
332                assign,
333                LocalAssignStatement::from_variable("var")
334                    .with_value(true)
335                    .with_tokens(LocalAssignTokens {
336                        local: Token::from_content("local"),
337                        equal: Some(Token::from_content("=")),
338                        variable_commas: Vec::new(),
339                        value_commas: Vec::new(),
340                    })
341            );
342        }
343
344        #[test]
345        fn removes_one_comma_token() {
346            let mut assign = LocalAssignStatement::from_variable("var")
347                .with_value(true)
348                .with_value(false)
349                .with_value(true)
350                .with_tokens(LocalAssignTokens {
351                    local: Token::from_content("local"),
352                    equal: Some(Token::from_content("=")),
353                    variable_commas: Vec::new(),
354                    value_commas: vec![Token::from_content(","), Token::from_content(",")],
355                });
356
357            assign.pop_value();
358
359            pretty_assertions::assert_eq!(
360                assign,
361                LocalAssignStatement::from_variable("var")
362                    .with_value(true)
363                    .with_value(false)
364                    .with_tokens(LocalAssignTokens {
365                        local: Token::from_content("local"),
366                        equal: Some(Token::from_content("=")),
367                        variable_commas: Vec::new(),
368                        value_commas: vec![Token::from_content(",")],
369                    })
370            );
371        }
372    }
373
374    mod remove_variable {
375        use super::*;
376
377        #[test]
378        fn single_variable_returns_none_without_mutating() {
379            let mut assign = LocalAssignStatement::from_variable("var").with_value(true);
380            let copy = assign.clone();
381
382            assert_eq!(assign.remove_variable(0), None);
383
384            pretty_assertions::assert_eq!(assign, copy);
385        }
386
387        #[test]
388        fn single_variable_remove_outside_of_bounds() {
389            let mut assign = LocalAssignStatement::from_variable("var");
390            let copy = assign.clone();
391
392            assert_eq!(assign.remove_variable(1), None);
393            pretty_assertions::assert_eq!(assign, copy);
394
395            assert_eq!(assign.remove_variable(3), None);
396            pretty_assertions::assert_eq!(assign, copy);
397        }
398
399        #[test]
400        fn two_variables_remove_first() {
401            let mut assign = LocalAssignStatement::from_variable("var")
402                .with_variable("var2")
403                .with_value(true)
404                .with_value(false);
405
406            assert_eq!(assign.remove_variable(0), Some(TypedIdentifier::new("var")));
407
408            pretty_assertions::assert_eq!(
409                assign,
410                LocalAssignStatement::from_variable("var2")
411                    .with_value(true)
412                    .with_value(false)
413            );
414        }
415
416        #[test]
417        fn two_variables_remove_second() {
418            let mut assign = LocalAssignStatement::from_variable("var")
419                .with_variable("var2")
420                .with_value(true)
421                .with_value(false);
422
423            assert_eq!(
424                assign.remove_variable(1),
425                Some(TypedIdentifier::new("var2"))
426            );
427
428            pretty_assertions::assert_eq!(
429                assign,
430                LocalAssignStatement::from_variable("var")
431                    .with_value(true)
432                    .with_value(false)
433            );
434        }
435    }
436}
darklua_core/nodes/statements/local_assign.rs

darklua_core/nodes/statements/
local_assign.rs
