1use std::fmt::{self, Display, Write};
2
3use crate::ast::*;
4use crate::token::{CommentKind, NumberLiteral, StringLiteralKind, TokenValue};
5
6macro_rules! impl_display {
7 ($( fn $name:ident(&mut $s:ident, $arg:ident: &$ast_name:ident<'_>) -> fmt::Result $body:block )+) => {
8 impl<W: Write> AstPrinter<'_, W> {
9 $(
10 fn $name(&mut $s, $arg: &$ast_name<'_>) -> fmt::Result $body
11 )+
12 }
13
14 $(
15 impl Display for $ast_name<'_> {
16 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
17 AstPrinter::new(f).$name(self)
18 }
19 }
20 )+
21 }
22}
23
24fn get_utf8_length(first_byte: u8) -> Option<usize> {
25 match first_byte {
26 0b0000_0000..=0b0111_1111 => Some(1),
27 0b1100_0000..=0b1101_1111 => Some(2),
28 0b1110_0000..=0b1110_1111 => Some(3),
29 0b1111_0000..=0b1111_0111 => Some(4),
30 _ => None,
31 }
32}
33
34struct AstPrinter<'w, W: Write> {
35 indent_level: usize,
36 writer: &'w mut W,
37}
38
39impl<'w, W: Write> AstPrinter<'w, W> {
40 fn new(writer: &'w mut W) -> Self {
41 Self {
42 indent_level: 0,
43 writer,
44 }
45 }
46
47 fn print_new_line(&mut self) -> fmt::Result {
48 writeln!(self.writer)?;
49
50 for _ in 0..self.indent_level {
51 write!(self.writer, " ")?
52 }
53
54 Ok(())
55 }
56
57 fn print_token(&mut self, token: &TokenReference<'_>) -> fmt::Result {
58 match &token.token.value {
59 TokenValue::Symbol(symbol) => write!(self.writer, "{}", symbol.as_str()),
60 TokenValue::Number(num_lit) => self.print_number_token(*num_lit),
61 TokenValue::String { value, kind } => self.print_string_token(value.as_ref(), *kind),
62 TokenValue::Comment { value, kind } => self.print_comment(value, *kind),
63 TokenValue::Whitespace(_) => unreachable!(),
64 TokenValue::Ident(s) => write!(self.writer, "{}", s),
65 TokenValue::Eof => unreachable!(),
66 }
67 }
68
69 fn print_number_token(&mut self, num_lit: NumberLiteral) -> fmt::Result {
70 match num_lit {
71 NumberLiteral::Float(v) => write!(self.writer, "{}", v),
72 NumberLiteral::Integer(v) => write!(self.writer, "{}", v),
73 }
74 }
75
76 fn print_string_token(&mut self, value: &[u8], kind: StringLiteralKind) -> fmt::Result {
77 match kind {
78 StringLiteralKind::SingleQuoted => write!(self.writer, "'")?,
79 StringLiteralKind::DoubleQuoted => write!(self.writer, "\"")?,
80 StringLiteralKind::Bracketed { level } => {
81 if level == 0 {
82 write!(self.writer, "[[")?
83 } else {
84 write!(self.writer, "[{:=>level$}[", "=", level = level)?
85 }
86 }
87 }
88
89 let quoted = !matches!(kind, StringLiteralKind::Bracketed { .. });
90
91 let iter = value.iter().enumerate();
92 let mut first = true;
93
94 for (i, &b) in iter {
95 let c = get_utf8_length(b)
96 .and_then(|len| value.get(i..(i + len)))
97 .and_then(|slice| std::str::from_utf8(slice).ok())
98 .map(|s| s.chars().next().unwrap())
99 .unwrap_or_else(|| b.into());
100
101 match c {
102 '\x07' if quoted => write!(self.writer, "\\a"),
103 '\x08' if quoted => write!(self.writer, "\\b"),
104 '\x0c' if quoted => write!(self.writer, "\\f"),
105 '\n' if quoted => write!(self.writer, "\\n"),
106 '\n' if first => write!(self.writer, "\n\n"),
107 '\r' if quoted => write!(self.writer, "\\r"),
108 '\r' => panic!("\\r cannot be represented in bracketed string literals"),
109 '\t' if quoted => write!(self.writer, "\\t"),
110 '\x0b' if quoted => write!(self.writer, "\\v"),
111 '\\' if quoted => write!(self.writer, "\\\\"),
112 '\'' if matches!(kind, StringLiteralKind::SingleQuoted) => {
113 write!(self.writer, "\\'")
114 }
115 '"' if matches!(kind, StringLiteralKind::DoubleQuoted) => {
116 write!(self.writer, "\\\"")
117 }
118 '\0'..='\x1f' | '\u{007f}'..='\u{00ff}' if quoted => {
119 write!(self.writer, "\\x{:02x}", c as u32)
120 }
121 _ => write!(self.writer, "{}", c),
122 }?;
123
124 first = false;
125 }
126
127 match kind {
128 StringLiteralKind::SingleQuoted => write!(self.writer, "'"),
129 StringLiteralKind::DoubleQuoted => write!(self.writer, "\""),
130 StringLiteralKind::Bracketed { level } => {
131 if level == 0 {
132 write!(self.writer, "]]")
133 } else {
134 write!(self.writer, "]{:=>level$}]", "=", level = level)
135 }
136 }
137 }
138 }
139
140 fn print_comment(&mut self, value: &str, kind: CommentKind) -> fmt::Result {
141 match kind {
142 CommentKind::Unbracketed => {
143 write!(self.writer, "--{}", value)?;
144 self.print_new_line()
145 }
146 CommentKind::Bracketed { level: 0 } => {
147 write!(self.writer, "--[[{}]]", value)
148 }
149 CommentKind::Bracketed { level } => {
150 write!(self.writer, "--[{1:=>level$}[{}]{1:=>level$}]", value, "=", level = level)
151 }
152 }
153 }
154
155 fn print_punctuated<T, R, F>(&mut self, punc: &Punctuated<'_, T>, f: F, put_space: bool) -> fmt::Result
156 where
157 F: Fn(&mut Self, &R) -> fmt::Result,
158 T: std::borrow::Borrow<R>,
159 {
160 for (v, sep) in &punc.pairs {
161 f(self, v.borrow())?;
162
163 if let Some(sep) = sep {
164 self.print_token(sep)?;
165
166 if put_space {
167 write!(self.writer, " ")?;
168 }
169 }
170 }
171
172 Ok(())
173 }
174
175 fn print_parenthesized_list<T, R, F>(&mut self, plist: &ParenthesizedList<'_, T>, f: F) -> fmt::Result
176 where
177 F: Fn(&mut Self, &R) -> fmt::Result,
178 T: std::borrow::Borrow<R>,
179 {
180 self.print_token(&plist.brackets.0)?;
181 self.print_punctuated(&plist.list, f, true)?;
182 self.print_token(&plist.brackets.1)
183 }
184
185 fn print_block_indent(&mut self, block: &Block<'_>) -> fmt::Result {
186 self.indent_level += 1;
187
188 if !block.statements.is_empty() {
189 self.print_new_line()?;
190 }
191
192 self.print_block(block)?;
193 self.indent_level -= 1;
194
195 Ok(())
196 }
197}
198
199impl_display![
200 fn print_nil_lit(&mut self, nil_lit: &NilLit<'_>) -> fmt::Result {
201 self.print_token(&nil_lit.0)
202 }
203
204 fn print_boolean_lit(&mut self, boolean_lit: &BooleanLit<'_>) -> fmt::Result {
205 self.print_token(&boolean_lit.0)
206 }
207
208 fn print_number_lit(&mut self, number_lit: &NumberLit<'_>) -> fmt::Result {
209 self.print_token(&number_lit.0)
210 }
211
212 fn print_string_lit(&mut self, string_lit: &StringLit<'_>) -> fmt::Result {
213 self.print_token(&string_lit.0)
214 }
215
216 fn print_name(&mut self, name: &Name<'_>) -> fmt::Result {
217 self.print_token(&name.0)
218 }
219
220 fn print_vararg(&mut self, vararg: &Vararg<'_>) -> fmt::Result {
221 self.print_token(&vararg.0)
222 }
223
224 fn print_bin_op_expr(&mut self, bin_op: &BinOpExpr<'_>) -> fmt::Result {
225 self.print_expr(&bin_op.left)?;
226 write!(self.writer, " {} ", bin_op.op.kind().as_symbol().as_str())?;
227 self.print_expr(&bin_op.right)
228 }
229
230 fn print_un_op_expr(&mut self, un_op: &UnOpExpr<'_>) -> fmt::Result {
231 write!(self.writer, "{}", un_op.op.kind().as_symbol().as_str())?;
232 self.print_expr(&un_op.right)
233 }
234
235 fn print_parenthesized_expr(&mut self, pexpr: &ParenthesizedExpr<'_>) -> fmt::Result {
236 self.print_token(&pexpr.brackets.0)?;
237 self.print_expr(&pexpr.expr)?;
238 self.print_token(&pexpr.brackets.1)
239 }
240
241 fn print_var_field(&mut self, var_field: &VarField<'_>) -> fmt::Result {
242 match var_field {
243 VarField::Expr { brackets, key } => {
244 self.print_token(&brackets.0)?;
245 self.print_expr(key)?;
246 self.print_token(&brackets.1)
247 }
248 VarField::Name { period, key } => {
249 self.print_token(period)?;
250 self.print_name(key)
251 }
252 }
253 }
254
255 fn print_var(&mut self, var: &Var<'_>) -> fmt::Result {
256 match var {
257 Var::Name(name) => self.print_name(name),
258 Var::Field(prefix_expr, var_field) => {
259 self.print_prefix_expr(prefix_expr)?;
260 self.print_var_field(var_field)
261 }
262 }
263 }
264
265 fn print_function_callee(&mut self, callee: &FunctionCallee<'_>) -> fmt::Result {
266 match callee {
267 FunctionCallee::Expr(prefix_expr) => self.print_prefix_expr(prefix_expr),
268 FunctionCallee::Method { object, colon, name } => {
269 self.print_prefix_expr(object)?;
270 self.print_token(colon)?;
271 self.print_name(name)
272 }
273 }
274 }
275
276 fn print_function_args(&mut self, args: &FunctionArgs<'_>) -> fmt::Result {
277 match args {
278 FunctionArgs::TableConstructor(tcons) => {
279 write!(self.writer, " ")?;
280 self.print_table_cons(tcons)
281 }
282 FunctionArgs::StringLit(string_lit) => {
283 write!(self.writer, " ")?;
284 self.print_string_lit(string_lit)
285 }
286 FunctionArgs::ParenthesizedList(plist) => {
287 self.print_parenthesized_list(plist, Self::print_expr)
288 }
289 }
290 }
291
292 fn print_function_call(&mut self, function_call: &FunctionCall<'_>) -> fmt::Result {
293 self.print_function_callee(&function_call.callee)?;
294 self.print_function_args(&function_call.args)
295 }
296
297 fn print_prefix_expr(&mut self, prefix_expr: &PrefixExpr<'_>) -> fmt::Result {
298 match prefix_expr {
299 PrefixExpr::Parenthesized(pexpr) => self.print_parenthesized_expr(pexpr),
300 PrefixExpr::Var(var) => self.print_var(var),
301 PrefixExpr::Call(call) => self.print_function_call(call),
302 }
303 }
304
305 fn print_expr(&mut self, expr: &Expr<'_>) -> fmt::Result {
306 match expr {
307 Expr::Nil(nil_lit) => self.print_nil_lit(nil_lit),
308 Expr::Boolean(boolean_lit) => self.print_boolean_lit(boolean_lit),
309 Expr::Number(number_lit) => self.print_number_lit(number_lit),
310 Expr::String(string_lit) => self.print_string_lit(string_lit),
311 Expr::Vararg(vararg) => self.print_vararg(vararg),
312 Expr::UnOp(un_op_expr) => self.print_un_op_expr(un_op_expr),
313 Expr::BinOp(bin_op_expr) => self.print_bin_op_expr(bin_op_expr),
314 Expr::TableConstructor(tcons) => self.print_table_cons(tcons),
315 Expr::Prefix(prefix_expr) => self.print_prefix_expr(prefix_expr),
316 Expr::Function(function_expr) => self.print_function_expr(function_expr),
317 }
318 }
319
320 fn print_table_cons(&mut self, tcons: &TableConstructor<'_>) -> fmt::Result {
321 self.print_token(&tcons.brackets.0)?;
322
323 for field in &tcons.fields {
324 self.print_table_field(field)?;
325 }
326
327 self.print_token(&tcons.brackets.1)
328 }
329
330 fn print_table_key(&mut self, key: &TableKey<'_>) -> fmt::Result {
331 match key {
332 TableKey::Expr { brackets, key } => {
333 self.print_token(&brackets.0)?;
334 self.print_expr(key)
335 }
336 TableKey::Name { key } => self.print_name(key),
337 }
338 }
339
340 fn print_table_field(&mut self, field: &TableField<'_>) -> fmt::Result {
341 if let Some((key, token)) = &field.key {
342 self.print_table_key(key)?;
343 write!(self.writer, " ")?;
344 self.print_token(token)?;
345 write!(self.writer, " ")?;
346 }
347
348 self.print_expr(&field.value)?;
349
350 if let Some(sep) = &field.separator {
351 self.print_token(sep)?;
352 write!(self.writer, " ")?;
353 }
354
355 Ok(())
356 }
357
358
359 fn print_block(&mut self, block: &Block<'_>) -> fmt::Result {
360 for (i, stmt) in block.statements.iter().enumerate() {
361 match (i, stmt) {
362 (0, _) | (_, Statement::Empty(_)) => {},
363 _ => self.print_new_line()?,
364 }
365
366 self.print_stmt(stmt)?;
367 }
368
369 Ok(())
370 }
371
372 fn print_stmt(&mut self, stmt: &Statement<'_>) -> fmt::Result {
373 match stmt {
374 Statement::Empty(empty_stmt) => self.print_empty_stmt(empty_stmt),
375 Statement::Block(block_stmt) => self.print_block_stmt(block_stmt),
376 Statement::Return(return_stmt) => self.print_return_stmt(return_stmt),
377 Statement::Break(break_stmt) => self.print_break_stmt(break_stmt),
378 Statement::Assignment(assignment_stmt) => self.print_assignment_stmt(assignment_stmt),
379 Statement::FunctionCall(call) => self.print_function_call(call),
380 Statement::While(while_stmt) => self.print_while_stmt(while_stmt),
381 Statement::For(for_stmt) => self.print_for_stmt(for_stmt),
382 Statement::Repeat(repeat_stmt) => self.print_repeat_stmt(repeat_stmt),
383 Statement::LocalDeclaration(local_decl) => self.print_local_decl(local_decl),
384 Statement::FunctionDeclaration(func_decl) => self.print_function_decl(func_decl),
385 Statement::Label(label_stmt) => self.print_label_stmt(label_stmt),
386 Statement::Goto(goto_stmt) => self.print_goto_stmt(goto_stmt),
387 Statement::If(if_stmt) => self.print_if_stmt(if_stmt),
388 }
389 }
390
391 fn print_empty_stmt(&mut self, stmt: &EmptyStat<'_>) -> fmt::Result {
392 self.print_token(&stmt.0)
393 }
394
395 fn print_block_stmt(&mut self, stmt: &BlockStat<'_>) -> fmt::Result {
396 self.print_token(&stmt.do_)?;
397 self.print_block_indent(&stmt.block)?;
398 self.print_token(&stmt.end)
399 }
400
401 fn print_return_stmt(&mut self, stmt: &ReturnStat<'_>) -> fmt::Result {
402 self.print_token(&stmt.return_)?;
403 write!(self.writer, " ")?;
404 self.print_punctuated(&stmt.exprs, Self::print_expr, true)?;
405
406 if let Some(semi) = &stmt.semi {
407 self.print_token(semi)?;
408 }
409
410 Ok(())
411 }
412
413 fn print_break_stmt(&mut self, stmt: &BreakStat<'_>) -> fmt::Result {
414 self.print_token(&stmt.0)
415 }
416
417 fn print_assignment_stmt(&mut self, stmt: &AssignmentStat<'_>) -> fmt::Result {
418 self.print_punctuated(&stmt.vars, Self::print_var, true)?;
419 write!(self.writer, " ")?;
420 self.print_token(&stmt.assign)?;
421 write!(self.writer, " ")?;
422 self.print_punctuated(&stmt.exprs, Self::print_expr, true)
423 }
424
425 fn print_while_stmt(&mut self, stmt: &WhileStat<'_>) -> fmt::Result {
426 self.print_token(&stmt.while_)?;
427 write!(self.writer, " ")?;
428 self.print_expr(&stmt.condition)?;
429 write!(self.writer, " ")?;
430 self.print_token(&stmt.do_)?;
431 self.print_block_indent(&stmt.block)?;
432 self.print_new_line()?;
433 self.print_token(&stmt.end)
434 }
435
436 fn print_for_stmt(&mut self, stmt: &ForStat<'_>) -> fmt::Result {
437 match stmt {
438 ForStat::Generic(generic_for) => self.print_generic_for(generic_for),
439 ForStat::Numerical(numerical_for) => self.print_numerical_for(numerical_for),
440 }
441 }
442
443 fn print_generic_for(&mut self, generic_for: &GenericFor<'_>) -> fmt::Result {
444 self.print_token(&generic_for.for_)?;
445 write!(self.writer, " ")?;
446 self.print_punctuated(&generic_for.names, Self::print_name, true)?;
447 write!(self.writer, " ")?;
448 self.print_token(&generic_for.in_)?;
449 write!(self.writer, " ")?;
450 self.print_punctuated(&generic_for.exprs, Self::print_expr, true)?;
451 write!(self.writer, " ")?;
452 self.print_token(&generic_for.do_)?;
453 self.print_block_indent(&generic_for.block)?;
454 self.print_new_line()?;
455 self.print_token(&generic_for.end)
456 }
457
458 fn print_numerical_for(&mut self, numerical_for: &NumericalFor<'_>) -> fmt::Result {
459 self.print_token(&numerical_for.for_)?;
460 write!(self.writer, " ")?;
461 self.print_name(&numerical_for.name)?;
462 write!(self.writer, " ")?;
463 self.print_token(&numerical_for.assign)?;
464 write!(self.writer, " ")?;
465 self.print_expr(&numerical_for.from)?;
466 self.print_token(&numerical_for.comma)?;
467 write!(self.writer, " ")?;
468 self.print_expr(&numerical_for.to)?;
469
470 if let Some((comma, step)) = &numerical_for.step {
471 self.print_token(comma)?;
472 write!(self.writer, " ")?;
473 self.print_expr(step)?;
474 }
475
476 write!(self.writer, " ")?;
477 self.print_token(&numerical_for.do_)?;
478 self.print_block_indent(&numerical_for.block)?;
479 self.print_new_line()?;
480 self.print_token(&numerical_for.end)
481 }
482
483 fn print_repeat_stmt(&mut self, stmt: &RepeatStat<'_>) -> fmt::Result {
484 self.print_token(&stmt.repeat)?;
485 self.print_block_indent(&stmt.block)?;
486 self.print_new_line()?;
487 self.print_token(&stmt.until)?;
488 write!(self.writer, " ")?;
489 self.print_expr(&stmt.condition)
490 }
491
492 fn print_local_decl(&mut self, stmt: &LocalDeclarationStat<'_>) -> fmt::Result {
493 self.print_token(&stmt.local)?;
494 write!(self.writer, " ")?;
495 self.print_punctuated(&stmt.names, Self::print_name, true)?;
496
497 if let Some(def) = &stmt.definition {
498 self.print_local_def(def)?;
499 }
500
501 Ok(())
502 }
503
504 fn print_local_def(&mut self, def: &LocalDefinition<'_>) -> fmt::Result {
505 write!(self.writer, " ")?;
506 self.print_token(&def.assign)?;
507 write!(self.writer, " ")?;
508 self.print_punctuated(&def.exprs, Self::print_expr, true)
509 }
510
511 fn print_function_decl(&mut self, stmt: &FunctionDeclarationStat<'_>) -> fmt::Result {
512 match stmt {
513 FunctionDeclarationStat::Local { local, function, name, body } => {
514 self.print_token(local)?;
515 write!(self.writer, " ")?;
516 self.print_token(function)?;
517 write!(self.writer, " ")?;
518 self.print_name(name)?;
519 self.print_function_body(body)
520 }
521 FunctionDeclarationStat::Nonlocal { function, name, body } => {
522 self.print_token(function)?;
523 write!(self.writer, " ")?;
524 self.print_function_name(name)?;
525 self.print_function_body(body)
526 }
527 }
528 }
529
530 fn print_function_body(&mut self, body: &FunctionBody<'_>) -> fmt::Result {
531 self.print_token(&body.params.brackets.0)?;
532 self.print_punctuated(&body.params.list, Self::print_name, true)?;
533
534 if let Some(vararg) = &body.vararg {
535 self.print_vararg(vararg)?;
536 }
537
538 self.print_token(&body.params.brackets.1)?;
539 self.print_block_indent(&body.block)?;
540 self.print_new_line()?;
541 self.print_token(&body.end)
542 }
543
544 fn print_function_name(&mut self, name: &FunctionName<'_>) -> fmt::Result {
545 match name {
546 FunctionName::PlainName(name) => self.print_name(name),
547 FunctionName::Indexed(punc) => self.print_punctuated(punc, Self::print_name, false),
548 FunctionName::Method { receiver, colon, method } => {
549 self.print_punctuated(receiver, Self::print_name, false)?;
550 self.print_token(colon)?;
551 self.print_name(method)
552 }
553 }
554 }
555
556 fn print_function_expr(&mut self, expr: &FunctionExpr<'_>) -> fmt::Result {
557 self.print_token(&expr.function)?;
558 self.print_function_body(&expr.body)
559 }
560
561 fn print_label_stmt(&mut self, stmt: &LabelStat<'_>) -> fmt::Result {
562 self.print_token(&stmt.preceding)?;
563 self.print_name(&stmt.name)?;
564 self.print_token(&stmt.following)
565 }
566
567 fn print_goto_stmt(&mut self, stmt: &GotoStat<'_>) -> fmt::Result {
568 self.print_token(&stmt.goto)?;
569 write!(self.writer, " ")?;
570 self.print_name(&stmt.label)
571 }
572
573 fn print_if_stmt(&mut self, stmt: &IfStat<'_>) -> fmt::Result {
574 self.print_token(&stmt.if_)?;
575 write!(self.writer, " ")?;
576 self.print_expr(&stmt.condition)?;
577 write!(self.writer, " ")?;
578 self.print_token(&stmt.then)?;
579 self.print_block_indent(&stmt.block)?;
580 self.print_new_line()?;
581
582 for elseif in &stmt.elseifs {
583 self.print_else_if(elseif)?;
584 }
585
586 if let Some(else_) = &stmt.else_ {
587 self.print_else(else_)?;
588 }
589
590 self.print_token(&stmt.end)
591 }
592
593 fn print_else_if(&mut self, elseif: &ElseIf<'_>) -> fmt::Result {
594 self.print_token(&elseif.elseif)?;
595 write!(self.writer, " ")?;
596 self.print_expr(&elseif.condition)?;
597 write!(self.writer, " ")?;
598 self.print_token(&elseif.then)?;
599 self.print_block_indent(&elseif.block)
600 }
601
602 fn print_else(&mut self, else_: &Else<'_>) -> fmt::Result {
603 self.print_token(&else_.else_)?;
604 self.print_block_indent(&else_.block)
605 }
606];