use rowan::NodeOrToken;
use crate::formatter::ir::Ir;
use crate::syntax::{SyntaxKind, SyntaxNode, SyntaxToken};
pub fn lower(root: &SyntaxNode) -> Ir {
lower_node(root)
}
fn lower_node(node: &SyntaxNode) -> Ir {
match node.kind() {
SyntaxKind::BINARY_EXPR | SyntaxKind::ASSIGNMENT_EXPR => lower_binary(node),
SyntaxKind::ARROW_EXPR => lower_arrow(node),
SyntaxKind::COMPARISON_EXPR => lower_comparison(node),
SyntaxKind::RANGE_EXPR => lower_range(node),
SyntaxKind::TYPE_ANNOTATION => lower_type_annotation(node),
SyntaxKind::MATRIX_EXPR => lower_matrix(node),
SyntaxKind::ARG_LIST => lower_arg_list(node),
SyntaxKind::TUPLE_EXPR | SyntaxKind::VECT_EXPR | SyntaxKind::BRACES => {
lower_collection(node)
}
SyntaxKind::KEYWORD_ARG => lower_keyword_arg(node),
SyntaxKind::PARAMETERS => lower_parameters(node),
_ => lower_transparent(node),
}
}
fn lower_transparent(node: &SyntaxNode) -> Ir {
Ir::concat(node.children_with_tokens().map(|el| match el {
NodeOrToken::Node(child) => lower_node(&child),
NodeOrToken::Token(tok) => Ir::text(tok.text().to_string()),
}))
}
fn lower_binary(node: &SyntaxNode) -> Ir {
let mut operands: Vec<SyntaxNode> = Vec::new();
let mut op: Option<SyntaxToken> = None;
for el in node.children_with_tokens() {
match el {
NodeOrToken::Node(child) => operands.push(child),
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::WHITESPACE => {}
SyntaxKind::COMMENT | SyntaxKind::BLOCK_COMMENT | SyntaxKind::NEWLINE => {
return lower_transparent(node);
}
_ if op.is_none() => op = Some(tok),
_ => return lower_transparent(node),
},
}
}
let (Some(op), [lhs, rhs]) = (op, operands.as_slice()) else {
return lower_transparent(node);
};
let lhs = lower_node(lhs);
let rhs = lower_node(rhs);
let op_text = Ir::text(op.text().to_string());
if is_tight_binop(op.kind()) {
Ir::concat([lhs, op_text, rhs])
} else {
Ir::concat([lhs, Ir::text(" "), op_text, Ir::text(" "), rhs])
}
}
fn lower_arrow(node: &SyntaxNode) -> Ir {
let mut operands: Vec<SyntaxNode> = Vec::new();
let mut op: Option<SyntaxToken> = None;
for el in node.children_with_tokens() {
match el {
NodeOrToken::Node(child) => operands.push(child),
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::WHITESPACE => {}
SyntaxKind::ARROW if op.is_none() => op = Some(tok),
_ => return lower_transparent(node),
},
}
}
let (Some(_), [lhs, rhs]) = (op, operands.as_slice()) else {
return lower_transparent(node);
};
Ir::concat([lower_node(lhs), Ir::text(" -> "), lower_node(rhs)])
}
fn lower_comparison(node: &SyntaxNode) -> Ir {
let mut parts: Vec<Ir> = Vec::new();
let mut expect_operand = true;
let mut operand_count = 0usize;
for el in node.children_with_tokens() {
match el {
NodeOrToken::Node(child) => {
if !expect_operand {
return lower_transparent(node);
}
if operand_count > 0 {
parts.push(Ir::text(" "));
}
parts.push(lower_node(&child));
operand_count += 1;
expect_operand = false;
}
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::WHITESPACE => {}
SyntaxKind::COMMENT | SyntaxKind::BLOCK_COMMENT | SyntaxKind::NEWLINE => {
return lower_transparent(node);
}
_ => {
if expect_operand {
return lower_transparent(node);
}
parts.push(Ir::text(" "));
parts.push(Ir::text(tok.text().to_string()));
expect_operand = true;
}
},
}
}
if expect_operand || operand_count < 2 {
return lower_transparent(node);
}
Ir::concat(parts)
}
fn lower_range(node: &SyntaxNode) -> Ir {
let mut parts: Vec<Ir> = Vec::new();
let mut expect_operand = true;
let mut operand_count = 0usize;
for el in node.children_with_tokens() {
match el {
NodeOrToken::Node(child) => {
if !expect_operand {
return lower_transparent(node);
}
parts.push(lower_node(&child));
operand_count += 1;
expect_operand = false;
}
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::WHITESPACE => {}
SyntaxKind::COLON if !expect_operand => {
parts.push(Ir::text(":"));
expect_operand = true;
}
_ => return lower_transparent(node),
},
}
}
if expect_operand || operand_count < 2 {
return lower_transparent(node);
}
Ir::concat(parts)
}
fn lower_type_annotation(node: &SyntaxNode) -> Ir {
let mut parts: Vec<Ir> = Vec::new();
let mut seen_colons = false;
for el in node.children_with_tokens() {
match el {
NodeOrToken::Node(child) => parts.push(lower_node(&child)),
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::WHITESPACE => {}
SyntaxKind::COLON_COLON if !seen_colons => {
seen_colons = true;
parts.push(Ir::text("::"));
}
_ => return lower_transparent(node),
},
}
}
if !seen_colons {
return lower_transparent(node);
}
Ir::concat(parts)
}
fn lower_arg_list(node: &SyntaxNode) -> Ir {
if has_newline_token(node) {
return lower_multiline_bracket(node);
}
let mut parts: Vec<Ir> = Vec::new();
let mut first_item = true;
let mut pending_comma = false;
for el in node.children_with_tokens() {
match el {
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::LPAREN
| SyntaxKind::RPAREN
| SyntaxKind::LBRACKET
| SyntaxKind::RBRACKET => parts.push(Ir::text(tok.text().to_string())),
SyntaxKind::WHITESPACE => {}
SyntaxKind::COMMA => {
if pending_comma {
return lower_transparent(node);
}
pending_comma = true;
}
_ => return lower_transparent(node),
},
NodeOrToken::Node(child) => match child.kind() {
SyntaxKind::ARG | SyntaxKind::KEYWORD_ARG => {
if !first_item {
if !pending_comma {
return lower_transparent(node);
}
parts.push(Ir::text(", "));
}
parts.push(lower_node(&child));
first_item = false;
pending_comma = false;
}
SyntaxKind::PARAMETERS => {
if pending_comma {
return lower_transparent(node);
}
parts.push(lower_node(&child));
first_item = false;
}
_ => return lower_transparent(node),
},
}
}
Ir::concat(parts)
}
fn lower_collection(node: &SyntaxNode) -> Ir {
if has_newline_token(node) {
return lower_multiline_bracket(node);
}
let keep_singleton_comma = node.kind() == SyntaxKind::TUPLE_EXPR;
let mut parts: Vec<Ir> = Vec::new();
let mut item_count = 0usize;
let mut pending_comma = false;
for el in node.children_with_tokens() {
match el {
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::LPAREN | SyntaxKind::LBRACKET | SyntaxKind::LBRACE => {
parts.push(Ir::text(tok.text().to_string()))
}
SyntaxKind::RPAREN | SyntaxKind::RBRACKET | SyntaxKind::RBRACE => {
if pending_comma && keep_singleton_comma && item_count == 1 {
parts.push(Ir::text(","));
}
parts.push(Ir::text(tok.text().to_string()));
}
SyntaxKind::WHITESPACE => {}
SyntaxKind::COMMA => {
if pending_comma {
return lower_transparent(node);
}
pending_comma = true;
}
_ => return lower_transparent(node),
},
NodeOrToken::Node(child) => match child.kind() {
SyntaxKind::ARG => {
if item_count > 0 {
if !pending_comma {
return lower_transparent(node);
}
parts.push(Ir::text(", "));
}
parts.push(lower_node(&child));
item_count += 1;
pending_comma = false;
}
_ => return lower_transparent(node),
},
}
}
Ir::concat(parts)
}
fn lower_keyword_arg(node: &SyntaxNode) -> Ir {
let mut parts: Vec<Ir> = Vec::new();
let mut seen_eq = false;
for el in node.children_with_tokens() {
match el {
NodeOrToken::Node(child) => parts.push(lower_node(&child)),
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::WHITESPACE => {}
SyntaxKind::EQ if !seen_eq => {
seen_eq = true;
parts.push(Ir::text(" = "));
}
_ => return lower_transparent(node),
},
}
}
if !seen_eq {
return lower_transparent(node);
}
Ir::concat(parts)
}
fn lower_parameters(node: &SyntaxNode) -> Ir {
let mut parts: Vec<Ir> = Vec::new();
let mut first_item = true;
let mut pending_comma = false;
let mut seen_semi = false;
for el in node.children_with_tokens() {
match el {
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::SEMICOLON if !seen_semi => {
seen_semi = true;
parts.push(Ir::text(";"));
}
SyntaxKind::WHITESPACE => {}
SyntaxKind::COMMA => {
if pending_comma {
return lower_transparent(node);
}
pending_comma = true;
}
_ => return lower_transparent(node),
},
NodeOrToken::Node(child) => {
if !matches!(child.kind(), SyntaxKind::ARG | SyntaxKind::KEYWORD_ARG) {
return lower_transparent(node);
}
if first_item {
parts.push(Ir::text(" "));
} else {
if !pending_comma {
return lower_transparent(node);
}
parts.push(Ir::text(", "));
}
parts.push(lower_node(&child));
first_item = false;
pending_comma = false;
}
}
}
if !seen_semi {
return lower_transparent(node);
}
Ir::concat(parts)
}
fn has_newline_token(node: &SyntaxNode) -> bool {
node.descendants_with_tokens()
.any(|el| el.kind() == SyntaxKind::NEWLINE)
}
enum Sep {
Space,
Newline { blanks: usize },
}
const MAX_BLANK_LINES: usize = 2;
fn adds_trailing_comma(node: &SyntaxNode) -> bool {
match node.kind() {
SyntaxKind::ARG_LIST => node
.parent()
.is_some_and(|p| p.kind() == SyntaxKind::INDEX_EXPR),
_ => true,
}
}
fn lower_multiline_bracket(node: &SyntaxNode) -> Ir {
let mut open: Option<String> = None;
let mut close: Option<String> = None;
let mut items: Vec<Ir> = Vec::new();
let mut seps: Vec<Sep> = Vec::new();
let mut newlines = 0usize;
let mut comma = false;
let mut leading_comma = false;
let mut leading_blanks = 0usize;
for el in node.children_with_tokens() {
match el {
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::LPAREN | SyntaxKind::LBRACKET | SyntaxKind::LBRACE => {
open = Some(tok.text().to_string())
}
SyntaxKind::RPAREN | SyntaxKind::RBRACKET | SyntaxKind::RBRACE => {
close = Some(tok.text().to_string())
}
SyntaxKind::WHITESPACE => {}
SyntaxKind::NEWLINE => newlines += 1,
SyntaxKind::COMMA => {
if comma {
return lower_transparent(node);
}
if items.is_empty() {
leading_comma = true;
}
comma = true;
}
_ => return lower_transparent(node),
},
NodeOrToken::Node(child) => match child.kind() {
SyntaxKind::ARG | SyntaxKind::KEYWORD_ARG => {
if items.is_empty() {
leading_blanks = newlines.saturating_sub(1).min(MAX_BLANK_LINES);
if leading_comma {
return lower_transparent(node);
}
} else {
if !comma {
return lower_transparent(node);
}
seps.push(if newlines >= 1 {
Sep::Newline {
blanks: (newlines - 1).min(MAX_BLANK_LINES),
}
} else {
Sep::Space
});
}
items.push(lower_node(&child));
newlines = 0;
comma = false;
}
_ => return lower_transparent(node),
},
}
}
let trailing_blanks = newlines.saturating_sub(1).min(MAX_BLANK_LINES);
let trailing_comma = comma;
let (Some(open), Some(close)) = (open, close) else {
return lower_transparent(node);
};
if items.is_empty() {
return lower_transparent(node);
}
let want_trailing = if adds_trailing_comma(node) {
true
} else {
trailing_comma
};
let n = items.len();
let mut inner: Vec<Ir> = Vec::with_capacity(n * 2 + 1 + leading_blanks + trailing_blanks);
for _ in 0..leading_blanks {
inner.push(Ir::BlankLine);
}
inner.push(Ir::HardLine); for (i, item) in items.into_iter().enumerate() {
inner.push(item);
if i + 1 < n {
inner.push(Ir::text(","));
match seps[i] {
Sep::Newline { blanks } => {
for _ in 0..blanks {
inner.push(Ir::BlankLine);
}
inner.push(Ir::HardLine);
}
Sep::Space => inner.push(Ir::text(" ")),
}
} else if want_trailing {
inner.push(Ir::text(","));
}
}
for _ in 0..trailing_blanks {
inner.push(Ir::BlankLine);
}
Ir::concat([
Ir::text(open),
Ir::indent(Ir::concat(inner)),
Ir::HardLine, Ir::text(close),
])
}
fn lower_matrix(node: &SyntaxNode) -> Ir {
if !has_newline_token(node) {
return lower_transparent(node);
}
let mut open: Option<String> = None;
let mut close: Option<String> = None;
let mut lines: Vec<Vec<(bool, Ir)>> = vec![Vec::new()];
for el in node.children_with_tokens() {
match el {
NodeOrToken::Token(tok) => match tok.kind() {
SyntaxKind::LBRACKET => open = Some(tok.text().to_string()),
SyntaxKind::RBRACKET => close = Some(tok.text().to_string()),
SyntaxKind::NEWLINE => lines.push(Vec::new()),
SyntaxKind::WHITESPACE => lines
.last_mut()
.unwrap()
.push((true, Ir::text(tok.text().to_string()))),
SyntaxKind::SEMICOLON => lines.last_mut().unwrap().push((false, Ir::text(";"))),
_ => return lower_transparent(node),
},
NodeOrToken::Node(child) => match child.kind() {
SyntaxKind::MATRIX_ROW | SyntaxKind::ARG => {
lines.last_mut().unwrap().push((false, lower_node(&child)))
}
_ => return lower_transparent(node),
},
}
}
let (Some(open), Some(close)) = (open, close) else {
return lower_transparent(node);
};
for line in &mut lines {
while line.first().is_some_and(|(ws, _)| *ws) {
line.remove(0);
}
while line.last().is_some_and(|(ws, _)| *ws) {
line.pop();
}
}
let first = lines.iter().position(|l| !l.is_empty());
let last = lines.iter().rposition(|l| !l.is_empty());
let (Some(first), Some(last)) = (first, last) else {
return lower_transparent(node);
};
let leading_blanks = first.saturating_sub(1).min(MAX_BLANK_LINES);
let trailing_blanks = (lines.len() - 1 - last)
.saturating_sub(1)
.min(MAX_BLANK_LINES);
let content = &lines[first..=last];
let mut inner: Vec<Ir> =
Vec::with_capacity(content.len() * 2 + leading_blanks + trailing_blanks);
for _ in 0..leading_blanks {
inner.push(Ir::BlankLine);
}
let mut pending_blanks = 0usize;
for line in content {
if line.is_empty() {
pending_blanks += 1;
continue;
}
for _ in 0..pending_blanks.min(MAX_BLANK_LINES) {
inner.push(Ir::BlankLine);
}
pending_blanks = 0;
inner.push(Ir::HardLine); inner.extend(line.iter().map(|(_, ir)| ir.clone()));
}
for _ in 0..trailing_blanks {
inner.push(Ir::BlankLine);
}
Ir::concat([
Ir::text(open),
Ir::indent(Ir::concat(inner)),
Ir::HardLine, Ir::text(close),
])
}
fn is_tight_binop(kind: SyntaxKind) -> bool {
matches!(kind, SyntaxKind::CARET | SyntaxKind::COLON)
}