//! This module contains functions for editing syntax trees. As the trees are //! immutable, all function here return a fresh copy of the tree, instead of //! doing an in-place modification. use std::{ fmt, iter, ops::{self, RangeInclusive}, }; use arrayvec::ArrayVec; use crate::{ algo::{self, neighbor, SyntaxRewriter}, ast::{ self, make::{self, tokens}, AstNode, TypeBoundsOwner, }, AstToken, Direction, InsertPosition, SmolStr, SyntaxElement, SyntaxKind, SyntaxKind::{ATTR, COMMENT, WHITESPACE}, SyntaxNode, SyntaxToken, T, }; impl ast::BinExpr { #[must_use] pub fn replace_op(&self, op: SyntaxKind) -> Option { let op_node: SyntaxElement = self.op_details()?.0.into(); let to_insert: Option = Some(make::token(op).into()); Some(self.replace_children(single_node(op_node), to_insert)) } } impl ast::Fn { #[must_use] pub fn with_body(&self, body: ast::BlockExpr) -> ast::Fn { let mut to_insert: ArrayVec<[SyntaxElement; 2]> = ArrayVec::new(); let old_body_or_semi: SyntaxElement = if let Some(old_body) = self.body() { old_body.syntax().clone().into() } else if let Some(semi) = self.semicolon_token() { to_insert.push(make::tokens::single_space().into()); semi.into() } else { to_insert.push(make::tokens::single_space().into()); to_insert.push(body.syntax().clone().into()); return self.insert_children(InsertPosition::Last, to_insert); }; to_insert.push(body.syntax().clone().into()); self.replace_children(single_node(old_body_or_semi), to_insert) } } fn make_multiline(node: N) -> N where N: AstNode + Clone, { let l_curly = match node.syntax().children_with_tokens().find(|it| it.kind() == T!['{']) { Some(it) => it, None => return node, }; let sibling = match l_curly.next_sibling_or_token() { Some(it) => it, None => return node, }; let existing_ws = match sibling.as_token() { None => None, Some(tok) if tok.kind() != WHITESPACE => None, Some(ws) => { if ws.text().contains('\n') { return node; } Some(ws.clone()) } }; let indent = leading_indent(node.syntax()).unwrap_or_default(); let ws = tokens::WsBuilder::new(&format!("\n{}", indent)); let to_insert = iter::once(ws.ws().into()); match existing_ws { None => node.insert_children(InsertPosition::After(l_curly), to_insert), Some(ws) => node.replace_children(single_node(ws), to_insert), } } impl ast::AssocItemList { #[must_use] pub fn append_items( &self, items: impl IntoIterator, ) -> ast::AssocItemList { let mut res = self.clone(); if !self.syntax().text().contains_char('\n') { res = make_multiline(res); } items.into_iter().for_each(|it| res = res.append_item(it)); res.fixup_trailing_whitespace().unwrap_or(res) } #[must_use] pub fn append_item(&self, item: ast::AssocItem) -> ast::AssocItemList { let (indent, position, whitespace) = match self.assoc_items().last() { Some(it) => ( leading_indent(it.syntax()).unwrap_or_default().to_string(), InsertPosition::After(it.syntax().clone().into()), "\n\n", ), None => match self.l_curly_token() { Some(it) => ( " ".to_string() + &leading_indent(self.syntax()).unwrap_or_default(), InsertPosition::After(it.into()), "\n", ), None => return self.clone(), }, }; let ws = tokens::WsBuilder::new(&format!("{}{}", whitespace, indent)); let to_insert: ArrayVec<[SyntaxElement; 2]> = [ws.ws().into(), item.syntax().clone().into()].into(); self.insert_children(position, to_insert) } /// Remove extra whitespace between last item and closing curly brace. fn fixup_trailing_whitespace(&self) -> Option { let first_token_after_items = self .assoc_items() .last()? .syntax() .next_sibling_or_token()?; let last_token_before_curly = self .r_curly_token()? .prev_sibling_or_token()?; if last_token_before_curly != first_token_after_items { // there is something more between last item and // right curly than just whitespace - bail out return None; } let whitespace = last_token_before_curly .clone() .into_token() .and_then(ast::Whitespace::cast)?; let text = whitespace.syntax().text(); let newline = text.rfind("\n")?; let keep = tokens::WsBuilder::new(&text[newline..]); Some(self.replace_children(first_token_after_items..=last_token_before_curly, std::iter::once(keep.ws().into()))) } } impl ast::RecordExprFieldList { #[must_use] pub fn append_field(&self, field: &ast::RecordExprField) -> ast::RecordExprFieldList { self.insert_field(InsertPosition::Last, field) } #[must_use] pub fn insert_field( &self, position: InsertPosition<&'_ ast::RecordExprField>, field: &ast::RecordExprField, ) -> ast::RecordExprFieldList { let is_multiline = self.syntax().text().contains_char('\n'); let ws; let space = if is_multiline { ws = tokens::WsBuilder::new(&format!( "\n{} ", leading_indent(self.syntax()).unwrap_or_default() )); ws.ws() } else { tokens::single_space() }; let mut to_insert: ArrayVec<[SyntaxElement; 4]> = ArrayVec::new(); to_insert.push(space.into()); to_insert.push(field.syntax().clone().into()); to_insert.push(make::token(T![,]).into()); macro_rules! after_l_curly { () => {{ let anchor = match self.l_curly_token() { Some(it) => it.into(), None => return self.clone(), }; InsertPosition::After(anchor) }}; } macro_rules! after_field { ($anchor:expr) => { if let Some(comma) = $anchor .syntax() .siblings_with_tokens(Direction::Next) .find(|it| it.kind() == T![,]) { InsertPosition::After(comma) } else { to_insert.insert(0, make::token(T![,]).into()); InsertPosition::After($anchor.syntax().clone().into()) } }; }; let position = match position { InsertPosition::First => after_l_curly!(), InsertPosition::Last => { if !is_multiline { // don't insert comma before curly to_insert.pop(); } match self.fields().last() { Some(it) => after_field!(it), None => after_l_curly!(), } } InsertPosition::Before(anchor) => { InsertPosition::Before(anchor.syntax().clone().into()) } InsertPosition::After(anchor) => after_field!(anchor), }; self.insert_children(position, to_insert) } } impl ast::TypeAlias { #[must_use] pub fn remove_bounds(&self) -> ast::TypeAlias { let colon = match self.colon_token() { Some(it) => it, None => return self.clone(), }; let end = match self.type_bound_list() { Some(it) => it.syntax().clone().into(), None => colon.clone().into(), }; self.replace_children(colon.into()..=end, iter::empty()) } } impl ast::TypeParam { #[must_use] pub fn remove_bounds(&self) -> ast::TypeParam { let colon = match self.colon_token() { Some(it) => it, None => return self.clone(), }; let end = match self.type_bound_list() { Some(it) => it.syntax().clone().into(), None => colon.clone().into(), }; self.replace_children(colon.into()..=end, iter::empty()) } } impl ast::Path { #[must_use] pub fn with_segment(&self, segment: ast::PathSegment) -> ast::Path { if let Some(old) = self.segment() { return self.replace_children( single_node(old.syntax().clone()), iter::once(segment.syntax().clone().into()), ); } self.clone() } } impl ast::PathSegment { #[must_use] pub fn with_type_args(&self, type_args: ast::GenericArgList) -> ast::PathSegment { self._with_type_args(type_args, false) } #[must_use] pub fn with_turbo_fish(&self, type_args: ast::GenericArgList) -> ast::PathSegment { self._with_type_args(type_args, true) } fn _with_type_args(&self, type_args: ast::GenericArgList, turbo: bool) -> ast::PathSegment { if let Some(old) = self.generic_arg_list() { return self.replace_children( single_node(old.syntax().clone()), iter::once(type_args.syntax().clone().into()), ); } let mut to_insert: ArrayVec<[SyntaxElement; 2]> = ArrayVec::new(); if turbo { to_insert.push(make::token(T![::]).into()); } to_insert.push(type_args.syntax().clone().into()); self.insert_children(InsertPosition::Last, to_insert) } } impl ast::Use { #[must_use] pub fn with_use_tree(&self, use_tree: ast::UseTree) -> ast::Use { if let Some(old) = self.use_tree() { return self.replace_descendant(old, use_tree); } self.clone() } pub fn remove(&self) -> SyntaxRewriter<'static> { let mut res = SyntaxRewriter::default(); res.delete(self.syntax()); let next_ws = self .syntax() .next_sibling_or_token() .and_then(|it| it.into_token()) .and_then(ast::Whitespace::cast); if let Some(next_ws) = next_ws { let ws_text = next_ws.syntax().text(); if ws_text.starts_with('\n') { let rest = &ws_text[1..]; if rest.is_empty() { res.delete(next_ws.syntax()) } else { res.replace(next_ws.syntax(), &make::tokens::whitespace(rest)); } } } res } } impl ast::UseTree { #[must_use] pub fn with_path(&self, path: ast::Path) -> ast::UseTree { if let Some(old) = self.path() { return self.replace_descendant(old, path); } self.clone() } #[must_use] pub fn with_use_tree_list(&self, use_tree_list: ast::UseTreeList) -> ast::UseTree { if let Some(old) = self.use_tree_list() { return self.replace_descendant(old, use_tree_list); } self.clone() } #[must_use] pub fn split_prefix(&self, prefix: &ast::Path) -> ast::UseTree { let suffix = if self.path().as_ref() == Some(prefix) && self.use_tree_list().is_none() { make::path_unqualified(make::path_segment_self()) } else { match split_path_prefix(&prefix) { Some(it) => it, None => return self.clone(), } }; let use_tree = make::use_tree( suffix, self.use_tree_list(), self.rename(), self.star_token().is_some(), ); let nested = make::use_tree_list(iter::once(use_tree)); return make::use_tree(prefix.clone(), Some(nested), None, false); fn split_path_prefix(prefix: &ast::Path) -> Option { let parent = prefix.parent_path()?; let segment = parent.segment()?; if algo::has_errors(segment.syntax()) { return None; } let mut res = make::path_unqualified(segment); for p in iter::successors(parent.parent_path(), |it| it.parent_path()) { res = make::path_qualified(res, p.segment()?); } Some(res) } } pub fn remove(&self) -> SyntaxRewriter<'static> { let mut res = SyntaxRewriter::default(); res.delete(self.syntax()); for &dir in [Direction::Next, Direction::Prev].iter() { if let Some(nb) = neighbor(self, dir) { self.syntax() .siblings_with_tokens(dir) .skip(1) .take_while(|it| it.as_node() != Some(nb.syntax())) .for_each(|el| res.delete(&el)); return res; } } res } } impl ast::MatchArmList { #[must_use] pub fn append_arms(&self, items: impl IntoIterator) -> ast::MatchArmList { let mut res = self.clone(); res = res.strip_if_only_whitespace(); if !res.syntax().text().contains_char('\n') { res = make_multiline(res); } items.into_iter().for_each(|it| res = res.append_arm(it)); res } fn strip_if_only_whitespace(&self) -> ast::MatchArmList { let mut iter = self.syntax().children_with_tokens().skip_while(|it| it.kind() != T!['{']); iter.next(); // Eat the curly let mut inner = iter.take_while(|it| it.kind() != T!['}']); if !inner.clone().all(|it| it.kind() == WHITESPACE) { return self.clone(); } let start = match inner.next() { Some(s) => s, None => return self.clone(), }; let end = match inner.last() { Some(s) => s, None => start.clone(), }; self.replace_children(start..=end, &mut iter::empty()) } #[must_use] pub fn remove_placeholder(&self) -> ast::MatchArmList { let placeholder = self.arms().find(|arm| matches!(arm.pat(), Some(ast::Pat::WildcardPat(_)))); if let Some(placeholder) = placeholder { self.remove_arm(&placeholder) } else { self.clone() } } #[must_use] fn remove_arm(&self, arm: &ast::MatchArm) -> ast::MatchArmList { let start = arm.syntax().clone(); let end = if let Some(comma) = start .siblings_with_tokens(Direction::Next) .skip(1) .skip_while(|it| it.kind().is_trivia()) .next() .filter(|it| it.kind() == T![,]) { comma } else { start.clone().into() }; self.replace_children(start.into()..=end, None) } #[must_use] pub fn append_arm(&self, item: ast::MatchArm) -> ast::MatchArmList { let r_curly = match self.syntax().children_with_tokens().find(|it| it.kind() == T!['}']) { Some(t) => t, None => return self.clone(), }; let position = InsertPosition::Before(r_curly.into()); let arm_ws = tokens::WsBuilder::new(" "); let match_indent = &leading_indent(self.syntax()).unwrap_or_default(); let match_ws = tokens::WsBuilder::new(&format!("\n{}", match_indent)); let to_insert: ArrayVec<[SyntaxElement; 3]> = [arm_ws.ws().into(), item.syntax().clone().into(), match_ws.ws().into()].into(); self.insert_children(position, to_insert) } } #[must_use] pub fn remove_attrs_and_docs(node: &N) -> N { N::cast(remove_attrs_and_docs_inner(node.syntax().clone())).unwrap() } fn remove_attrs_and_docs_inner(mut node: SyntaxNode) -> SyntaxNode { while let Some(start) = node.children_with_tokens().find(|it| it.kind() == ATTR || it.kind() == COMMENT) { let end = match &start.next_sibling_or_token() { Some(el) if el.kind() == WHITESPACE => el.clone(), Some(_) | None => start.clone(), }; node = algo::replace_children(&node, start..=end, &mut iter::empty()); } node } #[derive(Debug, Clone, Copy)] pub struct IndentLevel(pub u8); impl From for IndentLevel { fn from(level: u8) -> IndentLevel { IndentLevel(level) } } impl fmt::Display for IndentLevel { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let spaces = " "; let buf; let len = self.0 as usize * 4; let indent = if len <= spaces.len() { &spaces[..len] } else { buf = iter::repeat(' ').take(len).collect::(); &buf }; fmt::Display::fmt(indent, f) } } impl ops::Add for IndentLevel { type Output = IndentLevel; fn add(self, rhs: u8) -> IndentLevel { IndentLevel(self.0 + rhs) } } impl IndentLevel { pub fn from_node(node: &SyntaxNode) -> IndentLevel { let first_token = match node.first_token() { Some(it) => it, None => return IndentLevel(0), }; for ws in prev_tokens(first_token).filter_map(ast::Whitespace::cast) { let text = ws.syntax().text(); if let Some(pos) = text.rfind('\n') { let level = text[pos + 1..].chars().count() / 4; return IndentLevel(level as u8); } } IndentLevel(0) } /// XXX: this intentionally doesn't change the indent of the very first token. /// Ie, in something like /// ``` /// fn foo() { /// 92 /// } /// ``` /// if you indent the block, the `{` token would stay put. fn increase_indent(self, node: SyntaxNode) -> SyntaxNode { let mut rewriter = SyntaxRewriter::default(); node.descendants_with_tokens() .filter_map(|el| el.into_token()) .filter_map(ast::Whitespace::cast) .filter(|ws| { let text = ws.syntax().text(); text.contains('\n') }) .for_each(|ws| { let new_ws = make::tokens::whitespace(&format!("{}{}", ws.syntax(), self,)); rewriter.replace(ws.syntax(), &new_ws) }); rewriter.rewrite(&node) } fn decrease_indent(self, node: SyntaxNode) -> SyntaxNode { let mut rewriter = SyntaxRewriter::default(); node.descendants_with_tokens() .filter_map(|el| el.into_token()) .filter_map(ast::Whitespace::cast) .filter(|ws| { let text = ws.syntax().text(); text.contains('\n') }) .for_each(|ws| { let new_ws = make::tokens::whitespace( &ws.syntax().text().replace(&format!("\n{}", self), "\n"), ); rewriter.replace(ws.syntax(), &new_ws) }); rewriter.rewrite(&node) } } // FIXME: replace usages with IndentLevel above fn leading_indent(node: &SyntaxNode) -> Option { for token in prev_tokens(node.first_token()?) { if let Some(ws) = ast::Whitespace::cast(token.clone()) { let ws_text = ws.text(); if let Some(pos) = ws_text.rfind('\n') { return Some(ws_text[pos + 1..].into()); } } if token.text().contains('\n') { break; } } None } fn prev_tokens(token: SyntaxToken) -> impl Iterator { iter::successors(Some(token), |token| token.prev_token()) } pub trait AstNodeEdit: AstNode + Clone + Sized { #[must_use] fn insert_children( &self, position: InsertPosition, to_insert: impl IntoIterator, ) -> Self { let new_syntax = algo::insert_children(self.syntax(), position, to_insert); Self::cast(new_syntax).unwrap() } #[must_use] fn replace_children( &self, to_replace: RangeInclusive, to_insert: impl IntoIterator, ) -> Self { let new_syntax = algo::replace_children(self.syntax(), to_replace, to_insert); Self::cast(new_syntax).unwrap() } #[must_use] fn replace_descendant(&self, old: D, new: D) -> Self { self.replace_descendants(iter::once((old, new))) } #[must_use] fn replace_descendants( &self, replacement_map: impl IntoIterator, ) -> Self { let mut rewriter = SyntaxRewriter::default(); for (from, to) in replacement_map { rewriter.replace(from.syntax(), to.syntax()) } rewriter.rewrite_ast(self) } fn indent_level(&self) -> IndentLevel { IndentLevel::from_node(self.syntax()) } #[must_use] fn indent(&self, level: IndentLevel) -> Self { Self::cast(level.increase_indent(self.syntax().clone())).unwrap() } #[must_use] fn dedent(&self, level: IndentLevel) -> Self { Self::cast(level.decrease_indent(self.syntax().clone())).unwrap() } #[must_use] fn reset_indent(&self) -> Self { let level = IndentLevel::from_node(self.syntax()); self.dedent(level) } } impl AstNodeEdit for N {} fn single_node(element: impl Into) -> RangeInclusive { let element = element.into(); element.clone()..=element } #[test] fn test_increase_indent() { let arm_list = { let arm = make::match_arm(iter::once(make::wildcard_pat().into()), make::expr_unit()); make::match_arm_list(vec![arm.clone(), arm]) }; assert_eq!( arm_list.syntax().to_string(), "{ _ => (), _ => (), }" ); let indented = arm_list.indent(IndentLevel(2)); assert_eq!( indented.syntax().to_string(), "{ _ => (), _ => (), }" ); }