use crate::allocator::AstArena;
use crate::ast::{
Attribute, Block, ClassMember, Expression, ExpressionKind, Function, GenericType,
GenericTypePack, Local, Statement, StatementTag, StringQuoteStyle, TableAccess, TableItem,
Type, TypeKind, TypeList, TypeOrPack, TypePack, TypePackKind,
};
use crate::ast_names::AstNameTable;
use crate::cst::{
CstAttrList, CstAttribute, CstExprConstantInteger, CstExprConstantNumber,
CstExprConstantString, CstExprExplicitTypeInstantiation, CstExprGroup, CstExprIfElse,
CstExprIndexExpr, CstExprTypeAssertion, CstNode, CstNodeMap, CstStatCompoundAssign, CstStatDo,
CstStatFunction, CstStatLocalFunction, CstStatRepeat, CstStatReturn, CstStatTypeAlias,
CstStatTypeFunction, CstStringQuoteStyle, CstTypeGroup, CstTypeInstantiation,
CstTypePackExplicit, CstTypePackGeneric, CstTypeSingletonString, CstTypeTableItemKind,
TableSeparator,
};
use crate::location::{Location, Position};
use crate::parser::{ParseError, ParseMessage, ParseOptions, parse_bytes};
use luau_common::{BString, ByteSlice, LuauEscapeExt, flags};
use std::io::Write;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PrettyPrintResult {
pub code: BString,
pub error_location: Location,
pub parse_error: ParseMessage,
}
pub fn pretty_print(
source: impl AsRef<[u8]>,
options: ParseOptions,
with_types: bool,
ignore_parse_errors: bool,
) -> PrettyPrintResult {
let source = source.as_ref();
let arena = AstArena::new();
let mut names = AstNameTable::new(&arena);
let mut options = options;
options = options.with_cst_data(true);
match parse_bytes(source, &arena, &mut names, options) {
Ok(result) => {
if let Some(error) = result.metadata.errors.first()
&& !ignore_parse_errors
{
return parse_error_result(error);
}
let code = if with_types {
pretty_print_with_types_and_cst(result.root, &result.metadata.cst_nodes)
} else {
pretty_print_and_cst(result.root, &result.metadata.cst_nodes)
};
PrettyPrintResult {
code,
error_location: Location::zero(),
parse_error: ParseMessage::from(""),
}
}
Err(errors) => parse_error_result(errors.first()),
}
}
fn parse_error_result(error: &ParseError) -> PrettyPrintResult {
PrettyPrintResult {
code: BString::new(Vec::new()),
error_location: error.location,
parse_error: error.message.clone(),
}
}
pub fn pretty_print_with_types(root: Block) -> BString {
Printer::new(true, None).finish_root_block(root)
}
pub fn pretty_print_block(root: Block) -> BString {
Printer::new(false, None).finish_root_block(root)
}
fn pretty_print_and_cst(root: Block, cst_nodes: &CstNodeMap<'_>) -> BString {
Printer::new(false, Some(cst_nodes)).finish_root_block(root)
}
pub fn pretty_print_with_types_and_cst(root: Block, cst_nodes: &CstNodeMap<'_>) -> BString {
Printer::new(true, Some(cst_nodes)).finish_root_block(root)
}
#[derive(Debug, Clone, Copy)]
pub enum PrintableAstNodeRef<'ast> {
Statement(Statement<'ast>),
Expression(Expression<'ast>),
Type(Type<'ast>),
}
pub fn to_string(node: PrintableAstNodeRef<'_>) -> BString {
match node {
PrintableAstNodeRef::Statement(statement) => {
Printer::new(true, None).finish_root_statement(statement)
}
PrintableAstNodeRef::Expression(expression) => {
Printer::new(true, None).finish_expression(expression)
}
PrintableAstNodeRef::Type(annotation) => Printer::new(true, None).finish_type(annotation),
}
}
pub fn dump(node: PrintableAstNodeRef<'_>) {
let mut stdout = std::io::stdout().lock();
let code = to_string(node);
let _ = stdout.write_all(code.as_slice());
let _ = stdout.write_all(b"\n");
}
fn expression_cst<'map, 'ast>(
cst_nodes: Option<&'map CstNodeMap<'ast>>,
expression: Expression,
) -> Option<&'map CstNode<'ast>> {
let nodes = cst_nodes?;
nodes.get_expression(expression)
}
fn statement_cst<'map, 'ast>(
cst_nodes: Option<&'map CstNodeMap<'ast>>,
statement: Statement,
) -> Option<&'map CstNode<'ast>> {
let nodes = cst_nodes?;
nodes.get_statement(statement)
}
fn function_cst<'map, 'ast>(
cst_nodes: Option<&'map CstNodeMap<'ast>>,
function: &Function,
) -> Option<&'map CstNode<'ast>> {
let nodes = cst_nodes?;
nodes.get_function(function)
}
fn attribute_cst<'map, 'ast>(
cst_nodes: Option<&'map CstNodeMap<'ast>>,
attribute: &Attribute,
) -> Option<&'map CstNode<'ast>> {
let nodes = cst_nodes?;
nodes.get_attribute(attribute)
}
fn type_cst<'map, 'ast>(
cst_nodes: Option<&'map CstNodeMap<'ast>>,
ty: Type,
) -> Option<&'map CstNode<'ast>> {
let nodes = cst_nodes?;
nodes.get_type(ty)
}
fn type_pack_cst<'map, 'ast>(
cst_nodes: Option<&'map CstNodeMap<'ast>>,
pack: TypePack,
) -> Option<&'map CstNode<'ast>> {
let nodes = cst_nodes?;
nodes.get_type_pack(pack)
}
fn generic_type_cst<'map, 'ast>(
cst_nodes: Option<&'map CstNodeMap<'ast>>,
generic: &GenericType,
) -> Option<&'map CstNode<'ast>> {
let nodes = cst_nodes?;
nodes.get_generic_type(generic)
}
fn generic_type_pack_cst<'map, 'ast>(
cst_nodes: Option<&'map CstNodeMap<'ast>>,
generic: &GenericTypePack,
) -> Option<&'map CstNode<'ast>> {
let nodes = cst_nodes?;
nodes.get_generic_type_pack(generic)
}
struct Printer<'cst, 'ast> {
output: Vec<u8>,
position: Position,
last_byte: Option<u8>,
write_types: bool,
cst_nodes: Option<&'cst CstNodeMap<'ast>>,
}
#[derive(Clone, Copy)]
enum GenericPackEllipsisMode {
PreserveMissingFromCst,
AlwaysWrite,
}
mod expressions;
mod statements;
mod types;
impl<'cst, 'ast> Printer<'cst, 'ast> {
fn new(write_types: bool, cst_nodes: Option<&'cst CstNodeMap<'ast>>) -> Self {
Self {
output: Vec::new(),
position: Position::zero(),
last_byte: None,
write_types,
cst_nodes,
}
}
fn finish_expression(mut self, expression: Expression) -> BString {
self.position = expression.location.begin;
self.write_expression(expression);
BString::new(self.output)
}
fn finish_type(mut self, annotation: Type) -> BString {
self.position = annotation.location.begin;
self.write_type(annotation);
BString::new(self.output)
}
fn finish_root_statement(mut self, statement: Statement) -> BString {
self.position = Position::zero();
match statement.tag {
StatementTag::Block => self.write_root_block(statement.as_block_unchecked()),
_ => self.write_statement(statement),
}
BString::new(self.output)
}
fn finish_root_block(mut self, block: Block) -> BString {
self.position = Position::zero();
self.write_root_block(block);
BString::new(self.output)
}
fn advance(&mut self, position: Position) {
while self.position.line < position.line {
self.write_byte(b'\n');
self.position.line += 1;
self.position.column = 0;
}
while self.position.column < position.column {
self.write_byte(b' ');
self.position.column += 1;
}
}
fn write_byte(&mut self, byte: u8) {
self.output.push(byte);
self.last_byte = Some(byte);
}
fn write_bytes(&mut self, bytes: &[u8]) {
if bytes.is_empty() {
return;
}
self.output.extend_from_slice(bytes);
self.last_byte = bytes.last().copied();
self.position.column += bytes.len() as u32;
}
fn write_multiline(&mut self, bytes: &[u8]) {
for byte in bytes {
self.write_byte(*byte);
if *byte == b'\n' {
self.position.line += 1;
self.position.column = 0;
} else {
self.position.column += 1;
}
}
}
fn newline(&mut self) {
self.write_byte(b'\n');
self.position.line += 1;
self.position.column = 0;
}
fn keyword(&mut self, keyword: &str) {
self.identifierish(keyword.as_bytes());
}
fn identifier(&mut self, bytes: &[u8]) {
self.identifierish(bytes);
}
fn identifierish(&mut self, bytes: &[u8]) {
if self.last_byte.is_some_and(is_identifier_char) {
self.symbol(" ");
}
self.write_bytes(bytes);
}
fn symbol(&mut self, symbol: &str) {
self.write_bytes(symbol.as_bytes());
}
fn maybe_advance_and_write(
&mut self,
position: Option<Position>,
symbol: &str,
always_write: bool,
) {
if let Some(position) = position
&& position.has_value()
{
self.advance(position);
self.symbol(symbol);
} else if always_write {
self.symbol(symbol);
}
}
fn maybe_space(&mut self, position: Position, reserve: u32) {
if self.position.column + reserve < position.column {
self.symbol(" ");
}
}
fn advance_before(&mut self, position: Position, token_length: u32) {
self.advance(Position::new(
position.line,
position.column.saturating_sub(token_length),
));
}
fn write_expression_list_with_commas(
&mut self,
expressions: &[Expression],
comma_positions: Option<&[Position]>,
) {
for (index, expression) in expressions.iter().enumerate() {
if index > 0 {
if let Some(position) =
comma_positions.and_then(|positions| positions.get(index - 1))
{
self.advance(*position);
}
self.symbol(",");
}
self.write_expression(*expression);
}
}
fn write_attribute(&mut self, attribute: &Attribute) {
self.advance(attribute.location.begin);
match attribute_cst(self.cst_nodes, attribute) {
Some(CstNode::Attribute(CstAttribute::Simple { has_at })) => {
if *has_at {
self.symbol("@");
}
self.identifier(attribute.name.bytes());
}
Some(CstNode::Attribute(CstAttribute::Parametrized {
open_paren_position,
close_paren_position,
argument_commas,
})) => {
self.identifier(attribute.name.bytes());
if let Some(position) = open_paren_position {
self.maybe_advance_and_write(Some(*position), "(", false);
}
self.write_expression_list_with_commas(attribute.args, Some(argument_commas));
if let Some(position) = close_paren_position {
self.maybe_advance_and_write(Some(*position), ")", false);
}
}
_ => {
self.symbol("@");
self.identifier(attribute.name.bytes());
}
}
}
fn write_attributes(&mut self, attributes: &[&Attribute], attr_lists: Option<&[CstAttrList]>) {
let Some(attr_lists) = attr_lists else {
for attribute in attributes {
self.write_attribute(attribute);
}
return;
};
let mut attribute_index = 0;
let mut attr_list_index = 0;
while attribute_index < attributes.len() || attr_list_index < attr_lists.len() {
if attr_list_index == attr_lists.len()
|| (attribute_index < attributes.len()
&& attributes[attribute_index].location.begin
< attr_lists[attr_list_index].at_bracket_position)
{
self.write_attribute(attributes[attribute_index]);
attribute_index += 1;
continue;
}
let attr_list = &attr_lists[attr_list_index];
self.advance(attr_list.at_bracket_position);
self.symbol("@[");
for comma in &attr_list.comma_positions {
if attribute_index < attributes.len() {
self.write_attribute(attributes[attribute_index]);
attribute_index += 1;
}
self.advance(*comma);
self.symbol(",");
}
if attribute_index < attributes.len() {
self.write_attribute(attributes[attribute_index]);
attribute_index += 1;
}
self.maybe_advance_and_write(Some(attr_list.close_bracket_position), "]", false);
attr_list_index += 1;
}
}
fn write_generic_parameters(
&mut self,
generics: &[&GenericType],
generic_packs: &[&GenericTypePack],
cst: Option<(Position, &[Position], Position)>,
generic_pack_ellipsis_mode: GenericPackEllipsisMode,
) {
if generics.is_empty() && generic_packs.is_empty() {
return;
}
if let Some((open, _, _)) = cst {
self.advance(open);
} else {
let first_location = generics
.first()
.map(|generic| generic.location.begin)
.or_else(|| generic_packs.first().map(|generic| generic.location.begin));
if let Some(first_location) = first_location
&& first_location.column > 0
{
self.advance(Position::new(
first_location.line,
first_location.column.saturating_sub(1),
));
}
}
self.symbol("<");
let mut first = true;
let mut comma_index = 0;
let cst_nodes = self.cst_nodes;
for generic in generics {
if !first {
if let Some((_, commas, _)) = cst
&& let Some(position) = commas.get(comma_index)
{
self.advance(*position);
}
comma_index += 1;
self.symbol(",");
}
first = false;
self.advance(generic.location.begin);
self.identifier(generic.name.bytes());
if let Some(default) = generic.default_value {
if let Some(CstNode::GenericType(cst)) = generic_type_cst(cst_nodes, generic)
&& let Some(position) = cst.default_equals
{
self.advance(position);
} else {
self.maybe_space(default.location.begin, 2);
}
self.symbol("=");
self.write_type(default);
}
}
for generic_pack in generic_packs {
if !first {
if let Some((_, commas, _)) = cst
&& let Some(position) = commas.get(comma_index)
{
self.advance(*position);
}
comma_index += 1;
self.symbol(",");
}
first = false;
self.advance(generic_pack.location.begin);
self.identifier(generic_pack.name.bytes());
if let Some(CstNode::GenericTypePack(cst)) =
generic_type_pack_cst(cst_nodes, generic_pack)
{
match generic_pack_ellipsis_mode {
GenericPackEllipsisMode::PreserveMissingFromCst => {
self.maybe_advance_and_write(Some(cst.ellipsis), "...", false);
}
GenericPackEllipsisMode::AlwaysWrite => {
if cst.ellipsis.has_value() {
self.advance(cst.ellipsis);
}
self.symbol("...");
}
}
} else {
self.symbol("...");
}
if let Some(default) = generic_pack.default_value {
if let Some(CstNode::GenericTypePack(cst)) =
generic_type_pack_cst(cst_nodes, generic_pack)
&& let Some(position) = cst.default_equals
{
self.advance(position);
} else {
self.maybe_space(default.location.begin, 2);
}
self.symbol("=");
self.write_type_pack(default, false);
}
}
if let Some((_, _, close)) = cst {
self.maybe_advance_and_write(Some(close), ">", false);
} else {
self.symbol(">");
}
}
fn write_type_instantiation(
&mut self,
type_args: &[TypeOrPack],
cst: Option<&CstTypeInstantiation>,
) {
if let Some(cst) = cst {
self.maybe_advance_and_write(Some(cst.left_arrow_1), "<", false);
self.maybe_advance_and_write(Some(cst.left_arrow_2), "<", false);
} else {
self.symbol("<");
self.symbol("<");
}
for (index, type_arg) in type_args.iter().enumerate() {
if index > 0 {
if let Some(position) = cst.and_then(|cst| cst.comma_positions.get(index - 1)) {
self.advance(*position);
}
self.symbol(",");
}
self.write_type_or_pack(*type_arg);
}
if let Some(cst) = cst {
self.maybe_advance_and_write(Some(cst.right_arrow_1), ">", false);
self.maybe_advance_and_write(Some(cst.right_arrow_2), ">", false);
} else {
self.symbol(">");
self.symbol(">");
}
}
fn write_source_string_content(&mut self, bytes: &[u8]) {
self.write_multiline(bytes);
}
fn write_source_string(
&mut self,
bytes: &[u8],
quote_style: CstStringQuoteStyle,
block_depth: u32,
) {
match quote_style {
CstStringQuoteStyle::QuotedRaw => {
self.symbol("[");
for _ in 0..block_depth {
self.symbol("=");
}
self.symbol("[");
self.write_multiline(bytes);
self.symbol("]");
for _ in 0..block_depth {
self.symbol("=");
}
self.symbol("]");
}
CstStringQuoteStyle::QuotedDouble => {
self.symbol("\"");
self.write_multiline(bytes);
self.symbol("\"");
}
CstStringQuoteStyle::QuotedSingle => {
self.symbol("'");
self.write_multiline(bytes);
self.symbol("'");
}
CstStringQuoteStyle::QuotedInterp => {
self.symbol("`");
self.write_multiline(bytes);
self.symbol("`");
}
}
}
fn write_string(&mut self, bytes: &[u8]) {
let quote = if bytes.contains(&b'\'') { "\"" } else { "'" };
self.symbol(quote);
self.write_bytes(bytes.escape_luau().as_bytes());
self.symbol(quote);
}
fn write_table_record_key(&mut self, key: Expression) {
self.advance(key.location.begin);
match key.kind() {
ExpressionKind::String {
value,
quote_style: StringQuoteStyle::Unquoted,
} => self.identifier(value.as_bytes()),
_ => self.write_expression(key),
}
}
}
fn is_identifier_char(byte: u8) -> bool {
byte.is_ascii_alphanumeric() || byte == b'_'
}