use async_lsp::lsp_types::{Position, Range, TextEdit};
use comrak::{
Arena, ComrakOptions, ExtensionOptions,
nodes::{AstNode, NodeValue, TableAlignment},
parse_document,
};
use ropey::RopeSlice;
use unicode_width::UnicodeWidthStr;
pub fn format(rope: RopeSlice, range: Range) -> Vec<TextEdit> {
let tables = parse_tables(rope, range.start);
tables
.iter()
.map(
|Table {
header,
alignments,
rows,
range,
col_widths,
}| {
let separator = gen_separator(alignments, col_widths);
let rows: Vec<String> = [header.clone(), separator]
.iter()
.chain(rows.iter())
.map(|row| format_row(row, col_widths, alignments))
.collect();
let new_text = rows.join("\n");
TextEdit {
range: *range,
new_text,
}
},
)
.collect()
}
#[derive(Clone, Debug, Default)]
struct Table {
header: Vec<String>,
rows: Vec<Vec<String>>,
alignments: Vec<TableAlignment>,
col_widths: Vec<usize>,
range: Range,
}
fn parse_tables(rope: RopeSlice, start_line: Position) -> Vec<Table> {
let arena = Arena::new();
let options = ComrakOptions {
extension: ExtensionOptions {
table: true,
strikethrough: false,
tagfilter: false,
autolink: false,
tasklist: false,
superscript: false,
header_ids: None,
footnotes: false,
description_lists: false,
front_matter_delimiter: None,
multiline_block_quotes: false,
alerts: false,
math_dollars: false,
math_code: false,
wikilinks_title_after_pipe: false,
wikilinks_title_before_pipe: false,
underline: false,
subscript: false,
spoiler: false,
greentext: false,
image_url_rewriter: None,
link_url_rewriter: None,
},
..Default::default()
};
let root = parse_document(&arena, &rope.to_string(), &options);
let tables = find_table_nodes(root);
tables
.iter()
.filter_map(|(alignments, table_node)| {
let table_range = get_table_range(table_node, start_line)?;
let mut header: Vec<String> = Vec::new();
let mut rows: Vec<Vec<String>> = Vec::new();
for node in table_node.children() {
match &node.data.borrow().value {
NodeValue::TableRow(true) if header.is_empty() => {
header = extract_row_cells(node, rope);
}
NodeValue::TableRow(false) => {
rows.push(extract_row_cells(node, rope));
}
_ => {}
}
}
if alignments.is_empty() {
return None;
}
let col_widths = calculate_column_widths(&header, &rows, alignments);
let table = Table {
header,
rows,
alignments: alignments.clone(),
col_widths,
range: table_range,
};
Some(table)
})
.collect()
}
fn find_table_nodes<'a>(root: &'a AstNode<'a>) -> Vec<(Vec<TableAlignment>, &'a AstNode<'a>)> {
let mut tables = Vec::new();
let mut stack = vec![root];
while let Some(node) = stack.pop() {
if let NodeValue::Table(table) = &node.data.borrow().value {
tables.push((table.alignments.clone(), node))
}
stack.extend(node.children());
}
tables
}
fn get_table_range(table_node: &AstNode, start: Position) -> Option<Range> {
let pos = table_node.data.borrow().sourcepos;
Some(Range {
start: Position {
line: pos.start.line as u32 - 1 + start.line,
character: pos.start.column as u32 - 1 + start.character,
},
end: Position {
line: pos.end.line as u32 - 1 + start.line,
character: pos.end.column as u32 + start.character,
},
})
}
fn extract_row_cells<'a>(row_node: &'a AstNode<'a>, rope: RopeSlice) -> Vec<String> {
let mut cells = Vec::new();
for cell in row_node.children() {
if let NodeValue::TableCell = cell.data.borrow().value {
let pos = cell.data.borrow().sourcepos;
let start_byte = pos.start.column - 1;
let end_byte = pos.end.column;
if let Some(slice) = rope
.line(pos.start.line - 1)
.get_byte_slice(start_byte..end_byte)
{
cells.push(slice.to_string().trim().to_string());
} else {
cells.push(String::new());
}
}
}
cells
}
fn calculate_column_widths(
header: &[String],
rows: &[Vec<String>],
alignments: &[TableAlignment],
) -> Vec<usize> {
let mut widths: Vec<usize> = alignments
.iter()
.map(get_alignment_cell_minimum_width)
.collect();
rows.iter()
.cloned()
.chain(vec![header.to_vec()])
.for_each(|row| {
row.iter().enumerate().for_each(|(i, cell)| {
widths[i] = widths[i].max(cell.width())
});
});
widths
}
fn format_row(cells: &[String], col_widths: &[usize], alignments: &[TableAlignment]) -> String {
let cells: Vec<String> = cells
.iter()
.zip(col_widths)
.zip(alignments)
.map(|((cell, width), alignment)| {
let cell_width = cell.width();
let pad = width - cell_width;
match alignment {
TableAlignment::Right => format!(" {}{cell} ", " ".repeat(pad)),
TableAlignment::Center => {
if cell_width >= *width {
format!(" {cell} ")
} else {
let left_pad = pad / 2;
let right_pad = pad - left_pad;
format!(
" {}{}{} ",
" ".repeat(left_pad),
cell,
" ".repeat(right_pad)
)
}
}
_ => format!(" {cell}{} ", " ".repeat(pad)),
}
})
.collect();
format!("|{}|", cells.join("|"))
}
fn gen_separator(alignments: &[TableAlignment], col_widths: &[usize]) -> Vec<String> {
alignments
.iter()
.zip(col_widths)
.map(|(&alignment, &width)| {
let min = get_alignment_cell_minimum_width(&alignment);
let width = width.max(min);
match alignment {
TableAlignment::Left => format!(":{}", "-".repeat(width - 1)),
TableAlignment::Right => format!("{}:", "-".repeat(width - 1)),
TableAlignment::Center => format!(":{}:", "-".repeat(width - 2)),
TableAlignment::None => "-".repeat(width),
}
})
.collect()
}
fn get_alignment_cell_minimum_width(alignment: &TableAlignment) -> usize {
match alignment {
TableAlignment::Center => 5,
TableAlignment::Left | TableAlignment::Right => 4,
TableAlignment::None => 3,
}
}