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