1use crate::allocator::AstArena;
2use crate::ast::{
3 Attribute, Block, ClassMember, Expression, ExpressionKind, Function, GenericType,
4 GenericTypePack, Local, Statement, StatementTag, StringQuoteStyle, TableAccess, TableItem,
5 Type, TypeKind, TypeList, TypeOrPack, TypePack, TypePackKind,
6};
7use crate::ast_names::AstNameTable;
8use crate::cst::{
9 CstAttrList, CstAttribute, CstExprConstantInteger, CstExprConstantNumber,
10 CstExprConstantString, CstExprExplicitTypeInstantiation, CstExprGroup, CstExprIfElse,
11 CstExprIndexExpr, CstExprTypeAssertion, CstNode, CstNodeMap, CstStatCompoundAssign, CstStatDo,
12 CstStatFunction, CstStatLocalFunction, CstStatRepeat, CstStatReturn, CstStatTypeAlias,
13 CstStatTypeFunction, CstStringQuoteStyle, CstTypeGroup, CstTypeInstantiation,
14 CstTypePackExplicit, CstTypePackGeneric, CstTypeSingletonString, CstTypeTableItemKind,
15 TableSeparator,
16};
17use crate::location::{Location, Position};
18use crate::parser::{ParseError, ParseMessage, ParseOptions, parse_bytes};
19use luau_common::{BString, ByteSlice, LuauEscapeExt, flags};
20use std::io::Write;
21
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct PrettyPrintResult {
24 pub code: BString,
25 pub error_location: Location,
26 pub parse_error: ParseMessage,
27}
28
29pub fn pretty_print(
30 source: impl AsRef<[u8]>,
31 options: ParseOptions,
32 with_types: bool,
33 ignore_parse_errors: bool,
34) -> PrettyPrintResult {
35 let source = source.as_ref();
36 let arena = AstArena::new();
37 let mut names = AstNameTable::new(&arena);
38 let mut options = options;
39 options = options.with_cst_data(true);
40
41 match parse_bytes(source, &arena, &mut names, options) {
42 Ok(result) => {
43 if let Some(error) = result.metadata.errors.first()
44 && !ignore_parse_errors
45 {
46 return parse_error_result(error);
47 }
48
49 let code = if with_types {
50 pretty_print_with_types_and_cst(result.root, &result.metadata.cst_nodes)
51 } else {
52 pretty_print_and_cst(result.root, &result.metadata.cst_nodes)
53 };
54
55 PrettyPrintResult {
56 code,
57 error_location: Location::zero(),
58 parse_error: ParseMessage::from(""),
59 }
60 }
61 Err(errors) => parse_error_result(errors.first()),
62 }
63}
64
65fn parse_error_result(error: &ParseError) -> PrettyPrintResult {
66 PrettyPrintResult {
67 code: BString::new(Vec::new()),
68 error_location: error.location,
69 parse_error: error.message.clone(),
70 }
71}
72
73pub fn pretty_print_with_types(root: Block) -> BString {
74 Printer::new(true, None).finish_root_block(root)
75}
76
77pub fn pretty_print_block(root: Block) -> BString {
78 Printer::new(false, None).finish_root_block(root)
79}
80
81fn pretty_print_and_cst(root: Block, cst_nodes: &CstNodeMap<'_>) -> BString {
82 Printer::new(false, Some(cst_nodes)).finish_root_block(root)
83}
84
85pub fn pretty_print_with_types_and_cst(root: Block, cst_nodes: &CstNodeMap<'_>) -> BString {
86 Printer::new(true, Some(cst_nodes)).finish_root_block(root)
87}
88
89#[derive(Debug, Clone, Copy)]
90pub enum PrintableAstNodeRef<'ast> {
91 Statement(Statement<'ast>),
92 Expression(Expression<'ast>),
93 Type(Type<'ast>),
94}
95
96pub fn to_string(node: PrintableAstNodeRef<'_>) -> BString {
97 match node {
98 PrintableAstNodeRef::Statement(statement) => {
99 Printer::new(true, None).finish_root_statement(statement)
100 }
101 PrintableAstNodeRef::Expression(expression) => {
102 Printer::new(true, None).finish_expression(expression)
103 }
104 PrintableAstNodeRef::Type(annotation) => Printer::new(true, None).finish_type(annotation),
105 }
106}
107
108pub fn dump(node: PrintableAstNodeRef<'_>) {
109 let mut stdout = std::io::stdout().lock();
110 let code = to_string(node);
111 let _ = stdout.write_all(code.as_slice());
112 let _ = stdout.write_all(b"\n");
113}
114
115fn expression_cst<'map, 'ast>(
116 cst_nodes: Option<&'map CstNodeMap<'ast>>,
117 expression: Expression,
118) -> Option<&'map CstNode<'ast>> {
119 let nodes = cst_nodes?;
120 nodes.get_expression(expression)
121}
122
123fn statement_cst<'map, 'ast>(
124 cst_nodes: Option<&'map CstNodeMap<'ast>>,
125 statement: Statement,
126) -> Option<&'map CstNode<'ast>> {
127 let nodes = cst_nodes?;
128 nodes.get_statement(statement)
129}
130
131fn function_cst<'map, 'ast>(
132 cst_nodes: Option<&'map CstNodeMap<'ast>>,
133 function: &Function,
134) -> Option<&'map CstNode<'ast>> {
135 let nodes = cst_nodes?;
136 nodes.get_function(function)
137}
138
139fn attribute_cst<'map, 'ast>(
140 cst_nodes: Option<&'map CstNodeMap<'ast>>,
141 attribute: &Attribute,
142) -> Option<&'map CstNode<'ast>> {
143 let nodes = cst_nodes?;
144 nodes.get_attribute(attribute)
145}
146
147fn type_cst<'map, 'ast>(
148 cst_nodes: Option<&'map CstNodeMap<'ast>>,
149 ty: Type,
150) -> Option<&'map CstNode<'ast>> {
151 let nodes = cst_nodes?;
152 nodes.get_type(ty)
153}
154
155fn type_pack_cst<'map, 'ast>(
156 cst_nodes: Option<&'map CstNodeMap<'ast>>,
157 pack: TypePack,
158) -> Option<&'map CstNode<'ast>> {
159 let nodes = cst_nodes?;
160 nodes.get_type_pack(pack)
161}
162
163fn generic_type_cst<'map, 'ast>(
164 cst_nodes: Option<&'map CstNodeMap<'ast>>,
165 generic: &GenericType,
166) -> Option<&'map CstNode<'ast>> {
167 let nodes = cst_nodes?;
168 nodes.get_generic_type(generic)
169}
170
171fn generic_type_pack_cst<'map, 'ast>(
172 cst_nodes: Option<&'map CstNodeMap<'ast>>,
173 generic: &GenericTypePack,
174) -> Option<&'map CstNode<'ast>> {
175 let nodes = cst_nodes?;
176 nodes.get_generic_type_pack(generic)
177}
178
179struct Printer<'cst, 'ast> {
180 output: Vec<u8>,
181 position: Position,
182 last_byte: Option<u8>,
183 write_types: bool,
184 cst_nodes: Option<&'cst CstNodeMap<'ast>>,
185}
186
187#[derive(Clone, Copy)]
188enum GenericPackEllipsisMode {
189 PreserveMissingFromCst,
190 AlwaysWrite,
191}
192
193mod expressions;
194mod statements;
195mod types;
196
197impl<'cst, 'ast> Printer<'cst, 'ast> {
198 fn new(write_types: bool, cst_nodes: Option<&'cst CstNodeMap<'ast>>) -> Self {
199 Self {
200 output: Vec::new(),
201 position: Position::zero(),
202 last_byte: None,
203 write_types,
204 cst_nodes,
205 }
206 }
207
208 fn finish_expression(mut self, expression: Expression) -> BString {
209 self.position = expression.location.begin;
210 self.write_expression(expression);
211 BString::new(self.output)
212 }
213
214 fn finish_type(mut self, annotation: Type) -> BString {
215 self.position = annotation.location.begin;
216 self.write_type(annotation);
217 BString::new(self.output)
218 }
219
220 fn finish_root_statement(mut self, statement: Statement) -> BString {
221 self.position = Position::zero();
222
223 match statement.tag {
224 StatementTag::Block => self.write_root_block(statement.as_block_unchecked()),
225 _ => self.write_statement(statement),
226 }
227
228 BString::new(self.output)
229 }
230
231 fn finish_root_block(mut self, block: Block) -> BString {
232 self.position = Position::zero();
233 self.write_root_block(block);
234 BString::new(self.output)
235 }
236
237 fn advance(&mut self, position: Position) {
238 while self.position.line < position.line {
239 self.write_byte(b'\n');
240 self.position.line += 1;
241 self.position.column = 0;
242 }
243
244 while self.position.column < position.column {
245 self.write_byte(b' ');
246 self.position.column += 1;
247 }
248 }
249
250 fn write_byte(&mut self, byte: u8) {
251 self.output.push(byte);
252 self.last_byte = Some(byte);
253 }
254
255 fn write_bytes(&mut self, bytes: &[u8]) {
256 if bytes.is_empty() {
257 return;
258 }
259
260 self.output.extend_from_slice(bytes);
261 self.last_byte = bytes.last().copied();
262 self.position.column += bytes.len() as u32;
263 }
264
265 fn write_multiline(&mut self, bytes: &[u8]) {
266 for byte in bytes {
267 self.write_byte(*byte);
268 if *byte == b'\n' {
269 self.position.line += 1;
270 self.position.column = 0;
271 } else {
272 self.position.column += 1;
273 }
274 }
275 }
276
277 fn newline(&mut self) {
278 self.write_byte(b'\n');
279 self.position.line += 1;
280 self.position.column = 0;
281 }
282
283 fn keyword(&mut self, keyword: &str) {
284 self.identifierish(keyword.as_bytes());
285 }
286
287 fn identifier(&mut self, bytes: &[u8]) {
288 self.identifierish(bytes);
289 }
290
291 fn identifierish(&mut self, bytes: &[u8]) {
292 if self.last_byte.is_some_and(is_identifier_char) {
293 self.symbol(" ");
294 }
295 self.write_bytes(bytes);
296 }
297
298 fn symbol(&mut self, symbol: &str) {
299 self.write_bytes(symbol.as_bytes());
300 }
301
302 fn maybe_advance_and_write(
303 &mut self,
304 position: Option<Position>,
305 symbol: &str,
306 always_write: bool,
307 ) {
308 if let Some(position) = position
309 && position.has_value()
310 {
311 self.advance(position);
312 self.symbol(symbol);
313 } else if always_write {
314 self.symbol(symbol);
315 }
316 }
317
318 fn maybe_space(&mut self, position: Position, reserve: u32) {
319 if self.position.column + reserve < position.column {
320 self.symbol(" ");
321 }
322 }
323
324 fn advance_before(&mut self, position: Position, token_length: u32) {
325 self.advance(Position::new(
326 position.line,
327 position.column.saturating_sub(token_length),
328 ));
329 }
330
331 fn write_expression_list_with_commas(
332 &mut self,
333 expressions: &[Expression],
334 comma_positions: Option<&[Position]>,
335 ) {
336 for (index, expression) in expressions.iter().enumerate() {
337 if index > 0 {
338 if let Some(position) =
339 comma_positions.and_then(|positions| positions.get(index - 1))
340 {
341 self.advance(*position);
342 }
343 self.symbol(",");
344 }
345 self.write_expression(*expression);
346 }
347 }
348
349 fn write_attribute(&mut self, attribute: &Attribute) {
350 self.advance(attribute.location.begin);
351 match attribute_cst(self.cst_nodes, attribute) {
352 Some(CstNode::Attribute(CstAttribute::Simple { has_at })) => {
353 if *has_at {
354 self.symbol("@");
355 }
356 self.identifier(attribute.name.bytes());
357 }
358 Some(CstNode::Attribute(CstAttribute::Parametrized {
359 open_paren_position,
360 close_paren_position,
361 argument_commas,
362 })) => {
363 self.identifier(attribute.name.bytes());
364 if let Some(position) = open_paren_position {
365 self.maybe_advance_and_write(Some(*position), "(", false);
366 }
367 self.write_expression_list_with_commas(attribute.args, Some(argument_commas));
368 if let Some(position) = close_paren_position {
369 self.maybe_advance_and_write(Some(*position), ")", false);
370 }
371 }
372 _ => {
373 self.symbol("@");
374 self.identifier(attribute.name.bytes());
375 }
376 }
377 }
378
379 fn write_attributes(&mut self, attributes: &[&Attribute], attr_lists: Option<&[CstAttrList]>) {
380 let Some(attr_lists) = attr_lists else {
381 for attribute in attributes {
382 self.write_attribute(attribute);
383 }
384 return;
385 };
386
387 let mut attribute_index = 0;
388 let mut attr_list_index = 0;
389
390 while attribute_index < attributes.len() || attr_list_index < attr_lists.len() {
391 if attr_list_index == attr_lists.len()
392 || (attribute_index < attributes.len()
393 && attributes[attribute_index].location.begin
394 < attr_lists[attr_list_index].at_bracket_position)
395 {
396 self.write_attribute(attributes[attribute_index]);
397 attribute_index += 1;
398 continue;
399 }
400
401 let attr_list = &attr_lists[attr_list_index];
402 self.advance(attr_list.at_bracket_position);
403 self.symbol("@[");
404 for comma in &attr_list.comma_positions {
405 if attribute_index < attributes.len() {
406 self.write_attribute(attributes[attribute_index]);
407 attribute_index += 1;
408 }
409 self.advance(*comma);
410 self.symbol(",");
411 }
412 if attribute_index < attributes.len() {
413 self.write_attribute(attributes[attribute_index]);
414 attribute_index += 1;
415 }
416 self.maybe_advance_and_write(Some(attr_list.close_bracket_position), "]", false);
417 attr_list_index += 1;
418 }
419 }
420
421 fn write_generic_parameters(
422 &mut self,
423 generics: &[&GenericType],
424 generic_packs: &[&GenericTypePack],
425 cst: Option<(Position, &[Position], Position)>,
426 generic_pack_ellipsis_mode: GenericPackEllipsisMode,
427 ) {
428 if generics.is_empty() && generic_packs.is_empty() {
429 return;
430 }
431
432 if let Some((open, _, _)) = cst {
433 self.advance(open);
434 } else {
435 let first_location = generics
436 .first()
437 .map(|generic| generic.location.begin)
438 .or_else(|| generic_packs.first().map(|generic| generic.location.begin));
439 if let Some(first_location) = first_location
440 && first_location.column > 0
441 {
442 self.advance(Position::new(
443 first_location.line,
444 first_location.column.saturating_sub(1),
445 ));
446 }
447 }
448 self.symbol("<");
449 let mut first = true;
450 let mut comma_index = 0;
451 let cst_nodes = self.cst_nodes;
452 for generic in generics {
453 if !first {
454 if let Some((_, commas, _)) = cst
455 && let Some(position) = commas.get(comma_index)
456 {
457 self.advance(*position);
458 }
459 comma_index += 1;
460 self.symbol(",");
461 }
462 first = false;
463 self.advance(generic.location.begin);
464 self.identifier(generic.name.bytes());
465 if let Some(default) = generic.default_value {
466 if let Some(CstNode::GenericType(cst)) = generic_type_cst(cst_nodes, generic)
467 && let Some(position) = cst.default_equals
468 {
469 self.advance(position);
470 } else {
471 self.maybe_space(default.location.begin, 2);
472 }
473 self.symbol("=");
474 self.write_type(default);
475 }
476 }
477 for generic_pack in generic_packs {
478 if !first {
479 if let Some((_, commas, _)) = cst
480 && let Some(position) = commas.get(comma_index)
481 {
482 self.advance(*position);
483 }
484 comma_index += 1;
485 self.symbol(",");
486 }
487 first = false;
488 self.advance(generic_pack.location.begin);
489 self.identifier(generic_pack.name.bytes());
490 if let Some(CstNode::GenericTypePack(cst)) =
491 generic_type_pack_cst(cst_nodes, generic_pack)
492 {
493 match generic_pack_ellipsis_mode {
494 GenericPackEllipsisMode::PreserveMissingFromCst => {
495 self.maybe_advance_and_write(Some(cst.ellipsis), "...", false);
496 }
497 GenericPackEllipsisMode::AlwaysWrite => {
498 if cst.ellipsis.has_value() {
499 self.advance(cst.ellipsis);
500 }
501 self.symbol("...");
502 }
503 }
504 } else {
505 self.symbol("...");
506 }
507 if let Some(default) = generic_pack.default_value {
508 if let Some(CstNode::GenericTypePack(cst)) =
509 generic_type_pack_cst(cst_nodes, generic_pack)
510 && let Some(position) = cst.default_equals
511 {
512 self.advance(position);
513 } else {
514 self.maybe_space(default.location.begin, 2);
515 }
516 self.symbol("=");
517 self.write_type_pack(default, false);
518 }
519 }
520 if let Some((_, _, close)) = cst {
521 self.maybe_advance_and_write(Some(close), ">", false);
522 } else {
523 self.symbol(">");
524 }
525 }
526
527 fn write_type_instantiation(
528 &mut self,
529 type_args: &[TypeOrPack],
530 cst: Option<&CstTypeInstantiation>,
531 ) {
532 if let Some(cst) = cst {
533 self.maybe_advance_and_write(Some(cst.left_arrow_1), "<", false);
534 self.maybe_advance_and_write(Some(cst.left_arrow_2), "<", false);
535 } else {
536 self.symbol("<");
537 self.symbol("<");
538 }
539 for (index, type_arg) in type_args.iter().enumerate() {
540 if index > 0 {
541 if let Some(position) = cst.and_then(|cst| cst.comma_positions.get(index - 1)) {
542 self.advance(*position);
543 }
544 self.symbol(",");
545 }
546 self.write_type_or_pack(*type_arg);
547 }
548 if let Some(cst) = cst {
549 self.maybe_advance_and_write(Some(cst.right_arrow_1), ">", false);
550 self.maybe_advance_and_write(Some(cst.right_arrow_2), ">", false);
551 } else {
552 self.symbol(">");
553 self.symbol(">");
554 }
555 }
556
557 fn write_source_string_content(&mut self, bytes: &[u8]) {
558 self.write_multiline(bytes);
559 }
560
561 fn write_source_string(
562 &mut self,
563 bytes: &[u8],
564 quote_style: CstStringQuoteStyle,
565 block_depth: u32,
566 ) {
567 match quote_style {
568 CstStringQuoteStyle::QuotedRaw => {
569 self.symbol("[");
570 for _ in 0..block_depth {
571 self.symbol("=");
572 }
573 self.symbol("[");
574 self.write_multiline(bytes);
575 self.symbol("]");
576 for _ in 0..block_depth {
577 self.symbol("=");
578 }
579 self.symbol("]");
580 }
581 CstStringQuoteStyle::QuotedDouble => {
582 self.symbol("\"");
583 self.write_multiline(bytes);
584 self.symbol("\"");
585 }
586 CstStringQuoteStyle::QuotedSingle => {
587 self.symbol("'");
588 self.write_multiline(bytes);
589 self.symbol("'");
590 }
591 CstStringQuoteStyle::QuotedInterp => {
592 self.symbol("`");
593 self.write_multiline(bytes);
594 self.symbol("`");
595 }
596 }
597 }
598
599 fn write_string(&mut self, bytes: &[u8]) {
600 let quote = if bytes.contains(&b'\'') { "\"" } else { "'" };
601 self.symbol(quote);
602 self.write_bytes(bytes.escape_luau().as_bytes());
603 self.symbol(quote);
604 }
605
606 fn write_table_record_key(&mut self, key: Expression) {
607 self.advance(key.location.begin);
608 match key.kind() {
609 ExpressionKind::String {
610 value,
611 quote_style: StringQuoteStyle::Unquoted,
612 } => self.identifier(value.as_bytes()),
613 _ => self.write_expression(key),
614 }
615 }
616}
617
618fn is_identifier_char(byte: u8) -> bool {
619 byte.is_ascii_alphanumeric() || byte == b'_'
620}