use crate::cst::document::parse_document;
use crate::cst::green::{GreenChild, GreenNode};
use crate::cst::syntax::SyntaxKind;
use crate::error::Result;
use crate::prelude::*;
#[derive(Debug, Clone, Copy)]
pub struct FormatConfig {
pub indent_size: usize,
}
impl Default for FormatConfig {
fn default() -> Self {
Self { indent_size: 2 }
}
}
pub fn format(input: &str) -> Result<String> {
format_with_config(input, &FormatConfig::default())
}
pub fn format_with_config(input: &str, config: &FormatConfig) -> Result<String> {
if input.trim().is_empty() {
return Ok(String::new());
}
let doc = parse_document(input)?;
let mut formatter = Formatter::new(input, config);
formatter.format_node(doc.syntax(), 0)?;
Ok(formatter.finish())
}
struct Formatter<'a> {
source: &'a str,
config: &'a FormatConfig,
out: String,
indent_level: usize,
at_line_start: bool,
last_was_space: bool,
}
impl<'a> Formatter<'a> {
fn new(source: &'a str, config: &'a FormatConfig) -> Self {
Self {
source,
config,
out: String::with_capacity(source.len()),
indent_level: 0,
at_line_start: true,
last_was_space: false,
}
}
fn finish(self) -> String {
let mut out = self.out;
if !out.ends_with('\n') && !out.is_empty() {
out.push('\n');
}
out
}
fn indent(&mut self) {
if self.at_line_start {
for _ in 0..(self.indent_level * self.config.indent_size) {
self.out.push(' ');
}
self.at_line_start = false;
self.last_was_space = self.indent_level > 0;
}
}
fn newline(&mut self) {
if !self.at_line_start {
self.out.push('\n');
self.at_line_start = true;
self.last_was_space = false;
}
}
fn write_raw(&mut self, text: &str) {
if text.is_empty() {
return;
}
self.indent();
self.out.push_str(text);
if let Some(last_newline) = text.rfind('\n') {
let trailing = &text[last_newline + 1..];
self.at_line_start = trailing.is_empty();
self.last_was_space = trailing.ends_with(' ');
} else {
self.at_line_start = false;
self.last_was_space = text.ends_with(' ');
}
}
fn ensure_space(&mut self) {
if !self.at_line_start && !self.last_was_space {
self.out.push(' ');
self.last_was_space = true;
}
}
fn format_node(&mut self, node: &GreenNode, base: usize) -> Result<()> {
match node.kind() {
SyntaxKind::Document | SyntaxKind::Stream => {
self.format_children(node, base)?;
}
SyntaxKind::BlockMapping => {
self.format_block_mapping(node, base)?;
}
SyntaxKind::BlockSequence => {
self.format_block_sequence(node, base)?;
}
SyntaxKind::MappingEntry => {
self.format_mapping_entry(node, base)?;
}
SyntaxKind::SequenceItem => {
self.format_sequence_item(node, base)?;
}
SyntaxKind::FlowMapping | SyntaxKind::FlowSequence => {
self.write_verbatim(node, base);
}
_ => {
self.format_children(node, base)?;
}
}
Ok(())
}
fn format_children(&mut self, node: &GreenNode, base: usize) -> Result<()> {
let mut pos = base;
for child in node.children() {
match child {
GreenChild::Node(inner) => {
self.format_node(inner, pos)?;
}
GreenChild::Token { kind, len } => {
self.handle_token(*kind, &self.source[pos..pos + *len as usize]);
}
}
pos += child.text_len();
}
Ok(())
}
fn handle_token(&mut self, kind: SyntaxKind, text: &str) {
match kind {
SyntaxKind::Comment => {
self.ensure_space();
self.write_raw(text.trim_end());
self.newline();
}
SyntaxKind::Newline => {
self.newline();
}
SyntaxKind::Whitespace | SyntaxKind::Bom => {
}
SyntaxKind::DocStart => {
self.write_raw("---");
self.newline();
}
SyntaxKind::DocEnd => {
self.write_raw("...");
self.newline();
}
SyntaxKind::ColonIndicator => {
self.write_raw(":");
}
SyntaxKind::DashIndicator => {
if !self.at_line_start {
self.newline();
}
self.write_raw("-");
}
_ if kind.is_token() => {
let trimmed = if matches!(kind, SyntaxKind::PlainScalar) {
text.trim()
} else {
text
};
self.write_raw(trimmed);
}
_ => {}
}
}
fn format_block_mapping(&mut self, node: &GreenNode, base: usize) -> Result<()> {
let mut pos = base;
for child in node.children() {
match child {
GreenChild::Node(inner) if inner.kind() == SyntaxKind::MappingEntry => {
self.format_mapping_entry(inner, pos)?;
}
GreenChild::Token { kind, len } => {
self.handle_token(*kind, &self.source[pos..pos + *len as usize]);
}
GreenChild::Node(inner) => {
self.format_node(inner, pos)?;
}
}
pos += child.text_len();
}
Ok(())
}
fn format_block_sequence(&mut self, node: &GreenNode, base: usize) -> Result<()> {
let mut pos = base;
for child in node.children() {
match child {
GreenChild::Node(inner) if inner.kind() == SyntaxKind::SequenceItem => {
self.format_sequence_item(inner, pos)?;
}
GreenChild::Token { kind, len } => {
self.handle_token(*kind, &self.source[pos..pos + *len as usize]);
}
GreenChild::Node(inner) => {
self.format_node(inner, pos)?;
}
}
pos += child.text_len();
}
Ok(())
}
fn format_mapping_entry(&mut self, node: &GreenNode, base: usize) -> Result<()> {
let mut pos = base;
let mut saw_colon = false;
for child in node.children() {
match child {
GreenChild::Token { kind, len } => {
let text = &self.source[pos..pos + *len as usize];
if saw_colon
&& !matches!(
kind,
SyntaxKind::ColonIndicator
| SyntaxKind::Newline
| SyntaxKind::Whitespace
| SyntaxKind::Comment
)
{
self.ensure_space();
}
if matches!(kind, SyntaxKind::ColonIndicator) {
saw_colon = true;
}
self.handle_token(*kind, text);
}
GreenChild::Node(inner) => {
if saw_colon {
if matches!(
inner.kind(),
SyntaxKind::BlockMapping | SyntaxKind::BlockSequence
) {
self.newline();
self.indent_level += 1;
self.format_node(inner, pos)?;
self.indent_level -= 1;
} else {
self.ensure_space();
self.format_node(inner, pos)?;
}
} else {
self.format_node(inner, pos)?;
}
}
}
pos += child.text_len();
}
self.newline();
Ok(())
}
fn format_sequence_item(&mut self, node: &GreenNode, base: usize) -> Result<()> {
let mut pos = base;
let mut saw_dash = false;
for child in node.children() {
match child {
GreenChild::Token { kind, len } => {
let text = &self.source[pos..pos + *len as usize];
if saw_dash
&& !matches!(
kind,
SyntaxKind::DashIndicator
| SyntaxKind::Newline
| SyntaxKind::Whitespace
| SyntaxKind::Comment
)
{
self.ensure_space();
}
if matches!(kind, SyntaxKind::DashIndicator) {
saw_dash = true;
}
self.handle_token(*kind, text);
}
GreenChild::Node(inner) => {
if saw_dash {
if matches!(
inner.kind(),
SyntaxKind::BlockMapping | SyntaxKind::BlockSequence
) {
self.ensure_space();
self.indent_level += 1;
self.format_node(inner, pos)?;
self.indent_level -= 1;
} else {
self.ensure_space();
self.format_node(inner, pos)?;
}
} else {
self.format_node(inner, pos)?;
}
}
}
pos += child.text_len();
}
self.newline();
Ok(())
}
fn write_verbatim(&mut self, node: &GreenNode, base: usize) {
self.indent();
let text = node.text(&self.source[base..base + node.text_len()]);
self.out.push_str(&text);
self.at_line_start = text.ends_with('\n');
self.last_was_space = text.ends_with(' ');
}
}