darklua_core/nodes/statements/
local_assign.rs

1use crate::nodes::{Expression, Token, TypedIdentifier};
2
3#[derive(Clone, Debug, PartialEq, Eq)]
4pub struct LocalAssignTokens {
5    pub local: Token,
6    pub equal: Option<Token>,
7    pub variable_commas: Vec<Token>,
8    pub value_commas: Vec<Token>,
9}
10
11impl LocalAssignTokens {
12    super::impl_token_fns!(
13        target = [local]
14        iter = [variable_commas, value_commas, equal]
15    );
16}
17
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub struct LocalAssignStatement {
20    variables: Vec<TypedIdentifier>,
21    values: Vec<Expression>,
22    tokens: Option<LocalAssignTokens>,
23}
24
25impl LocalAssignStatement {
26    pub fn new(variables: Vec<TypedIdentifier>, values: Vec<Expression>) -> Self {
27        Self {
28            variables,
29            values,
30            tokens: None,
31        }
32    }
33
34    pub fn from_variable<S: Into<TypedIdentifier>>(variable: S) -> Self {
35        Self {
36            variables: vec![variable.into()],
37            values: Vec::new(),
38            tokens: None,
39        }
40    }
41
42    pub fn with_tokens(mut self, tokens: LocalAssignTokens) -> Self {
43        self.tokens = Some(tokens);
44        self
45    }
46
47    #[inline]
48    pub fn set_tokens(&mut self, tokens: LocalAssignTokens) {
49        self.tokens = Some(tokens);
50    }
51
52    #[inline]
53    pub fn get_tokens(&self) -> Option<&LocalAssignTokens> {
54        self.tokens.as_ref()
55    }
56
57    #[inline]
58    pub fn mutate_tokens(&mut self) -> Option<&mut LocalAssignTokens> {
59        self.tokens.as_mut()
60    }
61
62    pub fn with_variable<S: Into<TypedIdentifier>>(mut self, variable: S) -> Self {
63        self.variables.push(variable.into());
64        self
65    }
66
67    pub fn with_value<E: Into<Expression>>(mut self, value: E) -> Self {
68        self.values.push(value.into());
69        self
70    }
71
72    pub fn into_assignments(self) -> (Vec<TypedIdentifier>, Vec<Expression>) {
73        (self.variables, self.values)
74    }
75
76    pub fn append_assignment<S: Into<TypedIdentifier>>(&mut self, variable: S, value: Expression) {
77        self.variables.push(variable.into());
78        self.values.push(value);
79    }
80
81    pub fn for_each_assignment<F>(&mut self, mut callback: F)
82    where
83        F: FnMut(&mut TypedIdentifier, Option<&mut Expression>),
84    {
85        let mut values = self.values.iter_mut();
86        self.variables
87            .iter_mut()
88            .for_each(|variable| callback(variable, values.next()));
89    }
90
91    #[inline]
92    pub fn get_variables(&self) -> &Vec<TypedIdentifier> {
93        &self.variables
94    }
95
96    #[inline]
97    pub fn iter_variables(&self) -> impl Iterator<Item = &TypedIdentifier> {
98        self.variables.iter()
99    }
100
101    #[inline]
102    pub fn iter_mut_variables(&mut self) -> impl Iterator<Item = &mut TypedIdentifier> {
103        self.variables.iter_mut()
104    }
105
106    #[inline]
107    pub fn append_variables(&mut self, variables: &mut Vec<TypedIdentifier>) {
108        self.variables.append(variables);
109    }
110
111    #[inline]
112    pub fn extend_values<T: IntoIterator<Item = Expression>>(&mut self, iter: T) {
113        self.values.extend(iter);
114    }
115
116    #[inline]
117    pub fn iter_mut_values(&mut self) -> impl Iterator<Item = &mut Expression> {
118        self.values.iter_mut()
119    }
120
121    #[inline]
122    pub fn iter_values(&self) -> impl Iterator<Item = &Expression> {
123        self.values.iter()
124    }
125
126    #[inline]
127    pub fn push_variable(&mut self, variable: impl Into<TypedIdentifier>) {
128        self.variables.push(variable.into());
129    }
130
131    #[inline]
132    pub fn push_value(&mut self, value: impl Into<Expression>) {
133        self.values.push(value.into());
134    }
135
136    #[inline]
137    pub fn append_values(&mut self, values: &mut Vec<Expression>) {
138        self.values.append(values);
139    }
140
141    #[inline]
142    pub fn last_value(&self) -> Option<&Expression> {
143        self.values.last()
144    }
145
146    pub fn pop_value(&mut self) -> Option<Expression> {
147        let value = self.values.pop();
148        if let Some(tokens) = &mut self.tokens {
149            let length = self.values.len();
150            if length == 0 {
151                if !tokens.value_commas.is_empty() {
152                    tokens.value_commas.clear();
153                }
154                if tokens.equal.is_some() {
155                    tokens.equal = None;
156                }
157            } else {
158                tokens.value_commas.truncate(length.saturating_sub(1));
159            }
160        }
161        value
162    }
163
164    pub fn remove_value(&mut self, index: usize) -> Option<Expression> {
165        if index < self.values.len() {
166            let value = self.values.remove(index);
167
168            if let Some(tokens) = &mut self.tokens {
169                if index < tokens.value_commas.len() {
170                    tokens.value_commas.remove(index);
171                }
172                if self.values.is_empty() && tokens.equal.is_some() {
173                    tokens.equal = None;
174                }
175            }
176
177            Some(value)
178        } else {
179            None
180        }
181    }
182
183    pub fn remove_variable(&mut self, index: usize) -> Option<TypedIdentifier> {
184        let len = self.variables.len();
185
186        if len > 1 && index < len {
187            let variable = self.variables.remove(index);
188
189            if let Some(tokens) = &mut self.tokens {
190                if index < tokens.variable_commas.len() {
191                    tokens.variable_commas.remove(index);
192                }
193            }
194
195            Some(variable)
196        } else {
197            None
198        }
199    }
200
201    #[inline]
202    pub fn values_len(&self) -> usize {
203        self.values.len()
204    }
205
206    #[inline]
207    pub fn variables_len(&self) -> usize {
208        self.variables.len()
209    }
210
211    #[inline]
212    pub fn has_values(&self) -> bool {
213        !self.values.is_empty()
214    }
215
216    pub fn clear_types(&mut self) {
217        for variable in &mut self.variables {
218            variable.remove_type();
219        }
220    }
221
222    super::impl_token_fns!(iter = [variables, tokens]);
223}
224
225#[cfg(test)]
226mod test {
227    use super::*;
228
229    mod pop_value {
230        use super::*;
231
232        #[test]
233        fn removes_the_equal_sign() {
234            let mut assign = LocalAssignStatement::from_variable("var")
235                .with_value(true)
236                .with_tokens(LocalAssignTokens {
237                    local: Token::from_content("local"),
238                    equal: Some(Token::from_content("=")),
239                    variable_commas: Vec::new(),
240                    value_commas: Vec::new(),
241                });
242
243            assign.pop_value();
244
245            pretty_assertions::assert_eq!(
246                assign,
247                LocalAssignStatement::from_variable("var").with_tokens(LocalAssignTokens {
248                    local: Token::from_content("local"),
249                    equal: None,
250                    variable_commas: Vec::new(),
251                    value_commas: Vec::new(),
252                })
253            );
254        }
255
256        #[test]
257        fn removes_the_last_comma_token() {
258            let mut assign = LocalAssignStatement::from_variable("var")
259                .with_value(true)
260                .with_value(false)
261                .with_tokens(LocalAssignTokens {
262                    local: Token::from_content("local"),
263                    equal: Some(Token::from_content("=")),
264                    variable_commas: Vec::new(),
265                    value_commas: vec![Token::from_content(",")],
266                });
267
268            assign.pop_value();
269
270            pretty_assertions::assert_eq!(
271                assign,
272                LocalAssignStatement::from_variable("var")
273                    .with_value(true)
274                    .with_tokens(LocalAssignTokens {
275                        local: Token::from_content("local"),
276                        equal: Some(Token::from_content("=")),
277                        variable_commas: Vec::new(),
278                        value_commas: Vec::new(),
279                    })
280            );
281        }
282
283        #[test]
284        fn removes_one_comma_token() {
285            let mut assign = LocalAssignStatement::from_variable("var")
286                .with_value(true)
287                .with_value(false)
288                .with_value(true)
289                .with_tokens(LocalAssignTokens {
290                    local: Token::from_content("local"),
291                    equal: Some(Token::from_content("=")),
292                    variable_commas: Vec::new(),
293                    value_commas: vec![Token::from_content(","), Token::from_content(",")],
294                });
295
296            assign.pop_value();
297
298            pretty_assertions::assert_eq!(
299                assign,
300                LocalAssignStatement::from_variable("var")
301                    .with_value(true)
302                    .with_value(false)
303                    .with_tokens(LocalAssignTokens {
304                        local: Token::from_content("local"),
305                        equal: Some(Token::from_content("=")),
306                        variable_commas: Vec::new(),
307                        value_commas: vec![Token::from_content(",")],
308                    })
309            );
310        }
311    }
312
313    mod remove_variable {
314        use super::*;
315
316        #[test]
317        fn single_variable_returns_none_without_mutating() {
318            let mut assign = LocalAssignStatement::from_variable("var").with_value(true);
319            let copy = assign.clone();
320
321            assert_eq!(assign.remove_variable(0), None);
322
323            pretty_assertions::assert_eq!(assign, copy);
324        }
325
326        #[test]
327        fn single_variable_remove_outside_of_bounds() {
328            let mut assign = LocalAssignStatement::from_variable("var");
329            let copy = assign.clone();
330
331            assert_eq!(assign.remove_variable(1), None);
332            pretty_assertions::assert_eq!(assign, copy);
333
334            assert_eq!(assign.remove_variable(3), None);
335            pretty_assertions::assert_eq!(assign, copy);
336        }
337
338        #[test]
339        fn two_variables_remove_first() {
340            let mut assign = LocalAssignStatement::from_variable("var")
341                .with_variable("var2")
342                .with_value(true)
343                .with_value(false);
344
345            assert_eq!(assign.remove_variable(0), Some(TypedIdentifier::new("var")));
346
347            pretty_assertions::assert_eq!(
348                assign,
349                LocalAssignStatement::from_variable("var2")
350                    .with_value(true)
351                    .with_value(false)
352            );
353        }
354
355        #[test]
356        fn two_variables_remove_second() {
357            let mut assign = LocalAssignStatement::from_variable("var")
358                .with_variable("var2")
359                .with_value(true)
360                .with_value(false);
361
362            assert_eq!(
363                assign.remove_variable(1),
364                Some(TypedIdentifier::new("var2"))
365            );
366
367            pretty_assertions::assert_eq!(
368                assign,
369                LocalAssignStatement::from_variable("var")
370                    .with_value(true)
371                    .with_value(false)
372            );
373        }
374    }
375}