1use crate::nodes::{LastStatement, ReturnStatement, Statement, Token};
2
3#[derive(Clone, Debug, PartialEq, Eq)]
4pub struct BlockTokens {
5 pub semicolons: Vec<Option<Token>>,
6 pub last_semicolon: Option<Token>,
7 pub final_token: Option<Token>,
8}
9
10impl BlockTokens {
11 super::impl_token_fns!(
12 iter = [last_semicolon, final_token]
13 iter_flatten = [semicolons]
14 );
15}
16
17#[derive(Clone, Debug, PartialEq, Eq)]
18pub struct Block {
19 statements: Vec<Statement>,
20 last_statement: Option<LastStatement>,
21 tokens: Option<Box<BlockTokens>>,
22}
23
24impl Block {
25 pub fn new(statements: Vec<Statement>, last_statement: Option<LastStatement>) -> Self {
26 Self {
27 statements,
28 last_statement,
29 tokens: None,
30 }
31 }
32
33 pub fn with_tokens(mut self, tokens: BlockTokens) -> Self {
34 self.tokens = Some(tokens.into());
35 self
36 }
37
38 #[inline]
39 pub fn set_tokens(&mut self, tokens: BlockTokens) {
40 self.tokens = Some(tokens.into());
41 }
42
43 #[inline]
44 pub fn get_tokens(&self) -> Option<&BlockTokens> {
45 self.tokens.as_ref().map(|tokens| tokens.as_ref())
46 }
47
48 #[inline]
49 pub fn mutate_tokens(&mut self) -> Option<&mut BlockTokens> {
50 self.tokens.as_mut().map(|tokens| tokens.as_mut())
51 }
52
53 pub fn push_statement<T: Into<Statement>>(&mut self, statement: T) {
54 if let Some(tokens) = &mut self.tokens {
55 if self.statements.len() == tokens.semicolons.len() {
56 tokens.semicolons.push(None);
57 }
58 }
59 self.statements.push(statement.into());
60 }
61
62 pub fn with_statement<T: Into<Statement>>(mut self, statement: T) -> Self {
63 self.statements.push(statement.into());
64 self
65 }
66
67 pub fn insert_statement(&mut self, index: usize, statement: impl Into<Statement>) {
68 if index > self.statements.len() {
69 self.push_statement(statement.into());
70 } else {
71 self.statements.insert(index, statement.into());
72
73 if let Some(tokens) = &mut self.tokens {
74 if index <= tokens.semicolons.len() {
75 tokens.semicolons.insert(index, None);
76 }
77 }
78 }
79 }
80
81 #[inline]
82 pub fn set_last_statement(&mut self, last_statement: impl Into<LastStatement>) {
83 self.last_statement = Some(last_statement.into());
84 }
85
86 pub fn with_last_statement(mut self, last_statement: impl Into<LastStatement>) -> Self {
87 self.last_statement = Some(last_statement.into());
88 self
89 }
90
91 #[inline]
92 pub fn is_empty(&self) -> bool {
93 self.last_statement.is_none() && self.statements.is_empty()
94 }
95
96 #[inline]
97 pub fn statements_len(&self) -> usize {
98 self.statements.len()
99 }
100
101 #[inline]
102 pub fn iter_statements(&self) -> impl Iterator<Item = &Statement> {
103 self.statements.iter()
104 }
105
106 #[inline]
107 pub fn reverse_iter_statements(&self) -> impl Iterator<Item = &Statement> {
108 self.statements.iter().rev()
109 }
110
111 #[inline]
112 pub fn get_last_statement(&self) -> Option<&LastStatement> {
113 self.last_statement.as_ref()
114 }
115
116 pub fn filter_statements<F>(&mut self, mut f: F)
117 where
118 F: FnMut(&Statement) -> bool,
119 {
120 let mut i = 0;
121
122 while i != self.statements.len() {
123 if f(&self.statements[i]) {
124 i += 1;
125 } else {
126 self.statements.remove(i);
127
128 if let Some(tokens) = &mut self.tokens {
129 if i < tokens.semicolons.len() {
130 tokens.semicolons.remove(i);
131 }
132 }
133 }
134 }
135 }
136
137 pub fn filter_mut_statements<F>(&mut self, mut f: F)
138 where
139 F: FnMut(&mut Statement) -> bool,
140 {
141 let mut i = 0;
142
143 while i != self.statements.len() {
144 if f(&mut self.statements[i]) {
145 i += 1;
146 } else {
147 self.statements.remove(i);
148
149 if let Some(tokens) = &mut self.tokens {
150 if i < tokens.semicolons.len() {
151 tokens.semicolons.remove(i);
152 }
153 }
154 }
155 }
156 }
157
158 pub fn truncate(&mut self, length: usize) {
159 self.statements.truncate(length);
160 if let Some(tokens) = &mut self.tokens {
161 tokens.semicolons.truncate(length);
162 }
163 }
164
165 #[inline]
166 pub fn iter_mut_statements(&mut self) -> impl Iterator<Item = &mut Statement> {
167 self.statements.iter_mut()
168 }
169
170 #[inline]
171 pub fn mutate_statements(&mut self) -> &mut Vec<Statement> {
172 &mut self.statements
173 }
174
175 #[inline]
176 pub fn first_statement(&self) -> Option<&Statement> {
177 self.statements.first()
178 }
179
180 #[inline]
181 pub fn first_mut_statement(&mut self) -> Option<&mut Statement> {
182 self.statements.first_mut()
183 }
184
185 pub fn take_statements(&mut self) -> Vec<Statement> {
186 if let Some(tokens) = &mut self.tokens {
187 tokens.semicolons.clear();
188 }
189 self.statements.drain(..).collect()
190 }
191
192 pub fn take_last_statement(&mut self) -> Option<LastStatement> {
193 if let Some(tokens) = &mut self.tokens {
194 tokens.last_semicolon.take();
195 }
196 self.last_statement.take()
197 }
198
199 pub fn set_statements(&mut self, statements: Vec<Statement>) {
200 self.statements = statements;
201
202 if let Some(tokens) = &mut self.tokens {
203 tokens.semicolons.clear();
204 }
205 }
206
207 #[inline]
208 pub fn mutate_last_statement(&mut self) -> Option<&mut LastStatement> {
209 self.last_statement.as_mut()
210 }
211
212 #[inline]
213 pub fn replace_last_statement<S: Into<LastStatement>>(
214 &mut self,
215 statement: S,
216 ) -> Option<LastStatement> {
217 self.last_statement.replace(statement.into())
218 }
219
220 pub fn clear(&mut self) {
221 self.statements.clear();
222 self.last_statement.take();
223
224 if let Some(tokens) = &mut self.tokens {
225 tokens.semicolons.clear();
226 tokens.last_semicolon = None;
227 }
228 }
229
230 super::impl_token_fns!(iter = [tokens]);
231}
232
233impl Default for Block {
234 fn default() -> Self {
235 Self::new(Vec::new(), None)
236 }
237}
238
239impl<IntoStatement: Into<Statement>> From<IntoStatement> for Block {
240 fn from(statement: IntoStatement) -> Block {
241 Block::new(vec![statement.into()], None)
242 }
243}
244
245impl From<LastStatement> for Block {
246 fn from(statement: LastStatement) -> Block {
247 Block::new(Vec::new(), Some(statement))
248 }
249}
250
251impl From<ReturnStatement> for Block {
252 fn from(statement: ReturnStatement) -> Block {
253 Block::new(Vec::new(), Some(statement.into()))
254 }
255}
256
257#[cfg(test)]
258mod test {
259 use super::*;
260 use crate::{
261 nodes::{DoStatement, RepeatStatement},
262 Parser,
263 };
264
265 fn parse_block_with_tokens(lua: &str) -> Block {
266 let parser = Parser::default().preserve_tokens();
267 parser.parse(lua).expect("code should parse")
268 }
269
270 fn parse_statement_with_tokens(lua: &str) -> Statement {
271 let mut block = parse_block_with_tokens(lua);
272 assert!(block.get_last_statement().is_none());
273 let statements = block.take_statements();
274 assert_eq!(statements.len(), 1);
275 statements.into_iter().next().unwrap()
276 }
277
278 #[test]
279 fn default_block_is_empty() {
280 let block = Block::default();
281
282 assert!(block.is_empty());
283 }
284
285 #[test]
286 fn is_empty_is_true_when_block_has_no_statements_or_last_statement() {
287 let block = Block::new(Vec::new(), None);
288
289 assert!(block.is_empty());
290 }
291
292 #[test]
293 fn is_empty_is_false_when_block_has_a_last_statement() {
294 let block = Block::default().with_last_statement(LastStatement::new_break());
295
296 assert!(!block.is_empty());
297 }
298
299 #[test]
300 fn is_empty_is_false_when_block_a_statement() {
301 let block = Block::default().with_statement(DoStatement::default());
302
303 assert!(!block.is_empty());
304 }
305
306 #[test]
307 fn clear_removes_statements() {
308 let mut block = Block::default().with_statement(DoStatement::default());
309 block.clear();
310
311 assert!(block.is_empty());
312 }
313
314 #[test]
315 fn clear_removes_last_statement() {
316 let mut block = Block::default().with_last_statement(LastStatement::new_break());
317 block.clear();
318
319 assert!(block.is_empty());
320 assert_eq!(block.get_last_statement(), None);
321 }
322
323 #[test]
324 fn set_last_statement() {
325 let mut block = Block::default();
326 let continue_statement = LastStatement::new_continue();
327 block.set_last_statement(continue_statement.clone());
328
329 assert_eq!(block.get_last_statement(), Some(&continue_statement));
330 }
331
332 #[test]
333 fn insert_statement_at_index_0() {
334 let mut block = Block::default().with_statement(DoStatement::default());
335
336 let new_statement = RepeatStatement::new(Block::default(), false);
337 block.insert_statement(0, new_statement.clone());
338
339 assert_eq!(
340 block,
341 Block::default()
342 .with_statement(new_statement)
343 .with_statement(DoStatement::default())
344 );
345 }
346
347 #[test]
348 fn insert_statement_at_index_0_with_tokens() {
349 let mut block = parse_block_with_tokens("do end;");
350
351 block.insert_statement(0, RepeatStatement::new(Block::default(), false));
352
353 insta::assert_debug_snapshot!("insert_statement_at_index_0_with_tokens", block);
354 }
355
356 #[test]
357 fn insert_statement_at_upper_bound() {
358 let mut block = Block::default().with_statement(DoStatement::default());
359
360 let new_statement = RepeatStatement::new(Block::default(), false);
361 block.insert_statement(1, new_statement.clone());
362
363 assert_eq!(
364 block,
365 Block::default()
366 .with_statement(DoStatement::default())
367 .with_statement(new_statement)
368 );
369 }
370
371 #[test]
372 fn insert_statement_after_statement_upper_bound() {
373 let mut block = Block::default().with_statement(DoStatement::default());
374
375 let new_statement = RepeatStatement::new(Block::default(), false);
376 block.insert_statement(4, new_statement.clone());
377
378 assert_eq!(
379 block,
380 Block::default()
381 .with_statement(DoStatement::default())
382 .with_statement(new_statement)
383 );
384 }
385
386 #[test]
387 fn insert_statement_after_statement_upper_bound_with_tokens() {
388 let mut block = parse_block_with_tokens("do end;");
389
390 block.insert_statement(4, RepeatStatement::new(Block::default(), false));
391
392 insta::assert_debug_snapshot!(
393 "insert_statement_after_statement_upper_bound_with_tokens",
394 block
395 );
396 }
397
398 #[test]
399 fn push_statement_with_tokens() {
400 let mut block = parse_block_with_tokens("");
401
402 let new_statement = parse_statement_with_tokens("while true do end");
403 block.push_statement(new_statement);
404
405 pretty_assertions::assert_eq!(
406 block.get_tokens(),
407 Some(&BlockTokens {
408 semicolons: vec![None],
409 last_semicolon: None,
410 final_token: None,
411 })
412 );
413 }
414
415 #[test]
416 fn clean_removes_semicolon_tokens() {
417 let mut block = Block::default()
418 .with_statement(DoStatement::default())
419 .with_tokens(BlockTokens {
420 semicolons: vec![Some(Token::from_content(";"))],
421 last_semicolon: None,
422 final_token: None,
423 });
424 block.clear();
425
426 assert!(block.get_tokens().unwrap().semicolons.is_empty());
427 }
428
429 #[test]
430 fn clean_removes_last_semicolon_token() {
431 let mut block = Block::default()
432 .with_last_statement(LastStatement::new_break())
433 .with_tokens(BlockTokens {
434 semicolons: Vec::new(),
435 last_semicolon: Some(Token::from_content(";")),
436 final_token: None,
437 });
438 block.clear();
439
440 assert!(block.get_tokens().unwrap().last_semicolon.is_none());
441 }
442
443 #[test]
444 fn set_statements_clear_semicolon_tokens() {
445 let mut block = Block::default()
446 .with_statement(DoStatement::default())
447 .with_tokens(BlockTokens {
448 semicolons: vec![Some(Token::from_content(";"))],
449 last_semicolon: None,
450 final_token: None,
451 });
452 block.set_statements(Vec::new());
453
454 assert!(block.get_tokens().unwrap().semicolons.is_empty());
455 }
456
457 #[test]
458 fn take_last_statement_clear_semicolon_token() {
459 let mut block = Block::default()
460 .with_last_statement(LastStatement::new_break())
461 .with_tokens(BlockTokens {
462 semicolons: Vec::new(),
463 last_semicolon: Some(Token::from_content(";")),
464 final_token: None,
465 });
466
467 assert_eq!(
468 block.take_last_statement(),
469 Some(LastStatement::new_break())
470 );
471
472 assert!(block.get_tokens().unwrap().last_semicolon.is_none());
473 }
474
475 #[test]
476 fn filter_statements_does_not_panic_when_semicolons_do_not_match() {
477 let mut block = Block::default()
478 .with_statement(DoStatement::default())
479 .with_statement(DoStatement::default())
480 .with_tokens(BlockTokens {
481 semicolons: vec![Some(Token::from_content(";"))],
482 last_semicolon: None,
483 final_token: None,
484 });
485
486 block.filter_statements(|_statement| false);
487
488 pretty_assertions::assert_eq!(
489 block,
490 Block::default().with_tokens(BlockTokens {
491 semicolons: Vec::new(),
492 last_semicolon: None,
493 final_token: None,
494 })
495 );
496 }
497
498 #[test]
499 fn filter_mut_statements_does_not_panic_when_semicolons_do_not_match() {
500 let mut block = Block::default()
501 .with_statement(DoStatement::default())
502 .with_statement(DoStatement::default())
503 .with_tokens(BlockTokens {
504 semicolons: vec![Some(Token::from_content(";"))],
505 last_semicolon: None,
506 final_token: None,
507 });
508
509 block.filter_mut_statements(|_statement| false);
510
511 pretty_assertions::assert_eq!(
512 block,
513 Block::default().with_tokens(BlockTokens {
514 semicolons: Vec::new(),
515 last_semicolon: None,
516 final_token: None,
517 })
518 );
519 }
520}