Merge pull request #18575 from Giga-Bowser/flip-assists

Veykril · web-flow · commit 4a418837f2e3 · 2024-12-07T12:27:40.000Z
minor: Migrate `flip_*` assists to `SyntaxEditor`
diff --git a/crates/ide-assists/src/handlers/flip_binexpr.rs b/crates/ide-assists/src/handlers/flip_binexpr.rs
@@ -1,4 +1,7 @@
-use syntax::ast::{self, AstNode, BinExpr};
+use syntax::{
+    ast::{self, syntax_factory::SyntaxFactory, AstNode, BinExpr},
+    SyntaxKind, T,
+};
 
 use crate::{AssistContext, AssistId, AssistKind, Assists};
 
@@ -19,22 +22,17 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
 // ```
 pub(crate) fn flip_binexpr(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
     let expr = ctx.find_node_at_offset::<BinExpr>()?;
-    let rhs = expr.rhs()?.syntax().clone();
-    let lhs = expr.lhs()?.syntax().clone();
-
-    let lhs = if let Some(bin_expr) = BinExpr::cast(lhs.clone()) {
-        if bin_expr.op_kind() == expr.op_kind() {
-            bin_expr.rhs()?.syntax().clone()
-        } else {
-            lhs
-        }
-    } else {
-        lhs
+    let lhs = expr.lhs()?;
+    let rhs = expr.rhs()?;
+
+    let lhs = match &lhs {
+        ast::Expr::BinExpr(bin_expr) if bin_expr.op_kind() == expr.op_kind() => bin_expr.rhs()?,
+        _ => lhs,
     };
 
-    let op_range = expr.op_token()?.text_range();
+    let op_token = expr.op_token()?;
     // The assist should be applied only if the cursor is on the operator
-    let cursor_in_range = op_range.contains_range(ctx.selection_trimmed());
+    let cursor_in_range = op_token.text_range().contains_range(ctx.selection_trimmed());
     if !cursor_in_range {
         return None;
     }
@@ -47,13 +45,18 @@ pub(crate) fn flip_binexpr(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option
     acc.add(
         AssistId("flip_binexpr", AssistKind::RefactorRewrite),
         "Flip binary expression",
-        op_range,
-        |edit| {
-            if let FlipAction::FlipAndReplaceOp(new_op) = action {
-                edit.replace(op_range, new_op);
-            }
-            edit.replace(lhs.text_range(), rhs.text());
-            edit.replace(rhs.text_range(), lhs.text());
+        op_token.text_range(),
+        |builder| {
+            let mut editor = builder.make_editor(&expr.syntax().parent().unwrap());
+            let make = SyntaxFactory::new();
+            if let FlipAction::FlipAndReplaceOp(binary_op) = action {
+                editor.replace(op_token, make.token(binary_op))
+            };
+            // FIXME: remove `clone_for_update` when `SyntaxEditor` handles it for us
+            editor.replace(lhs.syntax(), rhs.syntax().clone_for_update());
+            editor.replace(rhs.syntax(), lhs.syntax().clone_for_update());
+            editor.add_mappings(make.finish_with_mappings());
+            builder.add_file_edits(ctx.file_id(), editor);
         },
     )
 }
@@ -62,7 +65,7 @@ enum FlipAction {
     // Flip the expression
     Flip,
     // Flip the expression and replace the operator with this string
-    FlipAndReplaceOp(&'static str),
+    FlipAndReplaceOp(SyntaxKind),
     // Do not flip the expression
     DontFlip,
 }
@@ -73,10 +76,10 @@ impl From<ast::BinaryOp> for FlipAction {
             ast::BinaryOp::Assignment { .. } => FlipAction::DontFlip,
             ast::BinaryOp::CmpOp(ast::CmpOp::Ord { ordering, strict }) => {
                 let rev_op = match (ordering, strict) {
-                    (ast::Ordering::Less, true) => ">",
-                    (ast::Ordering::Less, false) => ">=",
-                    (ast::Ordering::Greater, true) => "<",
-                    (ast::Ordering::Greater, false) => "<=",
+                    (ast::Ordering::Less, true) => T![>],
+                    (ast::Ordering::Less, false) => T![>=],
+                    (ast::Ordering::Greater, true) => T![<],
+                    (ast::Ordering::Greater, false) => T![<=],
                 };
                 FlipAction::FlipAndReplaceOp(rev_op)
             }
diff --git a/crates/ide-assists/src/handlers/flip_comma.rs b/crates/ide-assists/src/handlers/flip_comma.rs
@@ -1,7 +1,8 @@
-use ide_db::base_db::SourceDatabase;
-use syntax::TextSize;
 use syntax::{
-    algo::non_trivia_sibling, ast, AstNode, Direction, SyntaxKind, SyntaxToken, TextRange, T,
+    algo::non_trivia_sibling,
+    ast::{self, syntax_factory::SyntaxFactory},
+    syntax_editor::{Element, SyntaxMapping},
+    AstNode, Direction, NodeOrToken, SyntaxElement, SyntaxKind, SyntaxToken, T,
 };
 
 use crate::{AssistContext, AssistId, AssistKind, Assists};
@@ -25,8 +26,6 @@ pub(crate) fn flip_comma(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<(
     let comma = ctx.find_token_syntax_at_offset(T![,])?;
     let prev = non_trivia_sibling(comma.clone().into(), Direction::Prev)?;
     let next = non_trivia_sibling(comma.clone().into(), Direction::Next)?;
-    let (mut prev_text, mut next_text) = (prev.to_string(), next.to_string());
-    let (mut prev_range, mut next_range) = (prev.text_range(), next.text_range());
 
     // Don't apply a "flip" in case of a last comma
     // that typically comes before punctuation
@@ -40,53 +39,85 @@ pub(crate) fn flip_comma(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<(
         return None;
     }
 
-    if let Some(parent) = comma.parent().and_then(ast::TokenTree::cast) {
-        // An attribute. It often contains a path followed by a token tree (e.g. `align(2)`), so we have
-        // to be smarter.
-        let prev_start =
-            match comma.siblings_with_tokens(Direction::Prev).skip(1).find(|it| it.kind() == T![,])
-            {
-                Some(it) => position_after_token(it.as_token().unwrap()),
-                None => position_after_token(&parent.left_delimiter_token()?),
-            };
-        let prev_end = prev.text_range().end();
-        let next_start = next.text_range().start();
-        let next_end =
-            match comma.siblings_with_tokens(Direction::Next).skip(1).find(|it| it.kind() == T![,])
-            {
-                Some(it) => position_before_token(it.as_token().unwrap()),
-                None => position_before_token(&parent.right_delimiter_token()?),
-            };
-        prev_range = TextRange::new(prev_start, prev_end);
-        next_range = TextRange::new(next_start, next_end);
-        let file_text = ctx.db().file_text(ctx.file_id().file_id());
-        prev_text = file_text[prev_range].to_owned();
-        next_text = file_text[next_range].to_owned();
-    }
+    // FIXME: remove `clone_for_update` when `SyntaxEditor` handles it for us
+    let prev = match prev {
+        SyntaxElement::Node(node) => node.clone_for_update().syntax_element(),
+        _ => prev,
+    };
+    let next = match next {
+        SyntaxElement::Node(node) => node.clone_for_update().syntax_element(),
+        _ => next,
+    };
 
     acc.add(
         AssistId("flip_comma", AssistKind::RefactorRewrite),
         "Flip comma",
         comma.text_range(),
-        |edit| {
-            edit.replace(prev_range, next_text);
-            edit.replace(next_range, prev_text);
+        |builder| {
+            let parent = comma.parent().unwrap();
+            let mut editor = builder.make_editor(&parent);
+
+            if let Some(parent) = ast::TokenTree::cast(parent) {
+                // An attribute. It often contains a path followed by a
+                // token tree (e.g. `align(2)`), so we have to be smarter.
+                let (new_tree, mapping) = flip_tree(parent.clone(), comma);
+                editor.replace(parent.syntax(), new_tree.syntax());
+                editor.add_mappings(mapping);
+            } else {
+                editor.replace(prev.clone(), next.clone());
+                editor.replace(next.clone(), prev.clone());
+            }
+
+            builder.add_file_edits(ctx.file_id(), editor);
         },
     )
 }
 
-fn position_before_token(token: &SyntaxToken) -> TextSize {
-    match non_trivia_sibling(token.clone().into(), Direction::Prev) {
-        Some(prev_token) => prev_token.text_range().end(),
-        None => token.text_range().start(),
-    }
-}
-
-fn position_after_token(token: &SyntaxToken) -> TextSize {
-    match non_trivia_sibling(token.clone().into(), Direction::Next) {
-        Some(prev_token) => prev_token.text_range().start(),
-        None => token.text_range().end(),
-    }
+fn flip_tree(tree: ast::TokenTree, comma: SyntaxToken) -> (ast::TokenTree, SyntaxMapping) {
+    let mut tree_iter = tree.token_trees_and_tokens();
+    let before: Vec<_> =
+        tree_iter.by_ref().take_while(|it| it.as_token() != Some(&comma)).collect();
+    let after: Vec<_> = tree_iter.collect();
+
+    let not_ws = |element: &NodeOrToken<_, SyntaxToken>| match element {
+        NodeOrToken::Token(token) => token.kind() != SyntaxKind::WHITESPACE,
+        NodeOrToken::Node(_) => true,
+    };
+
+    let is_comma = |element: &NodeOrToken<_, SyntaxToken>| match element {
+        NodeOrToken::Token(token) => token.kind() == T![,],
+        NodeOrToken::Node(_) => false,
+    };
+
+    let prev_start_untrimmed = match before.iter().rposition(is_comma) {
+        Some(pos) => pos + 1,
+        None => 1,
+    };
+    let prev_end = 1 + before.iter().rposition(not_ws).unwrap();
+    let prev_start = prev_start_untrimmed
+        + before[prev_start_untrimmed..prev_end].iter().position(not_ws).unwrap();
+
+    let next_start = after.iter().position(not_ws).unwrap();
+    let next_end_untrimmed = match after.iter().position(is_comma) {
+        Some(pos) => pos,
+        None => after.len() - 1,
+    };
+    let next_end = 1 + after[..next_end_untrimmed].iter().rposition(not_ws).unwrap();
+
+    let result = [
+        &before[1..prev_start],
+        &after[next_start..next_end],
+        &before[prev_end..],
+        &[NodeOrToken::Token(comma)],
+        &after[..next_start],
+        &before[prev_start..prev_end],
+        &after[next_end..after.len() - 1],
+    ]
+    .concat();
+
+    let make = SyntaxFactory::new();
+    let new_token_tree = make.token_tree(tree.left_delimiter_token().unwrap().kind(), result);
+    (new_token_tree, make.finish_with_mappings())
 }
 
 #[cfg(test)]
@@ -147,4 +178,9 @@ mod tests {
             r#"#[foo(bar, qux, baz(1 + 1), other)] struct Foo;"#,
         );
     }
+
+    #[test]
+    fn flip_comma_attribute_incomplete() {
+        check_assist_not_applicable(flip_comma, r#"#[repr(align(2),$0)] struct Foo;"#);
+    }
 }
diff --git a/crates/ide-assists/src/handlers/flip_trait_bound.rs b/crates/ide-assists/src/handlers/flip_trait_bound.rs
@@ -23,21 +23,23 @@ pub(crate) fn flip_trait_bound(acc: &mut Assists, ctx: &AssistContext<'_>) -> Op
     let plus = ctx.find_token_syntax_at_offset(T![+])?;
 
     // Make sure we're in a `TypeBoundList`
-    ast::TypeBoundList::cast(plus.parent()?)?;
+    let parent = ast::TypeBoundList::cast(plus.parent()?)?;
 
     let (before, after) = (
-        non_trivia_sibling(plus.clone().into(), Direction::Prev)?,
-        non_trivia_sibling(plus.clone().into(), Direction::Next)?,
+        non_trivia_sibling(plus.clone().into(), Direction::Prev)?.into_node()?,
+        non_trivia_sibling(plus.clone().into(), Direction::Next)?.into_node()?,
     );
 
     let target = plus.text_range();
     acc.add(
         AssistId("flip_trait_bound", AssistKind::RefactorRewrite),
         "Flip trait bounds",
         target,
-        |edit| {
-            edit.replace(before.text_range(), after.to_string());
-            edit.replace(after.text_range(), before.to_string());
+        |builder| {
+            let mut editor = builder.make_editor(parent.syntax());
+            editor.replace(before.clone(), after.clone_for_update());
+            editor.replace(after.clone(), before.clone_for_update());
+            builder.add_file_edits(ctx.file_id(), editor);
         },
     )
 }
diff --git a/crates/syntax/src/ast/syntax_factory/constructors.rs b/crates/syntax/src/ast/syntax_factory/constructors.rs
@@ -4,7 +4,7 @@ use itertools::Itertools;
 use crate::{
     ast::{self, make, HasName, HasTypeBounds},
     syntax_editor::SyntaxMappingBuilder,
-    AstNode,
+    AstNode, NodeOrToken, SyntaxKind, SyntaxNode, SyntaxToken,
 };
 
 use super::SyntaxFactory;
@@ -80,6 +80,23 @@ impl SyntaxFactory {
         ast
     }
 
+    pub fn expr_bin(&self, lhs: ast::Expr, op: ast::BinaryOp, rhs: ast::Expr) -> ast::BinExpr {
+        let ast::Expr::BinExpr(ast) =
+            make::expr_bin_op(lhs.clone(), op, rhs.clone()).clone_for_update()
+        else {
+            unreachable!()
+        };
+
+        if let Some(mut mapping) = self.mappings() {
+            let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
+            builder.map_node(lhs.syntax().clone(), ast.lhs().unwrap().syntax().clone());
+            builder.map_node(rhs.syntax().clone(), ast.rhs().unwrap().syntax().clone());
+            builder.finish(&mut mapping);
+        }
+
+        ast
+    }
+
     pub fn expr_path(&self, path: ast::Path) -> ast::Expr {
         let ast::Expr::PathExpr(ast) = make::expr_path(path.clone()).clone_for_update() else {
             unreachable!()
@@ -151,4 +168,34 @@ impl SyntaxFactory {
 
         ast
     }
+
+    pub fn token_tree(
+        &self,
+        delimiter: SyntaxKind,
+        tt: Vec<NodeOrToken<ast::TokenTree, SyntaxToken>>,
+    ) -> ast::TokenTree {
+        let tt: Vec<_> = tt.into_iter().collect();
+        let input: Vec<_> = tt.iter().cloned().filter_map(only_nodes).collect();
+
+        let ast = make::token_tree(delimiter, tt).clone_for_update();
+
+        if let Some(mut mapping) = self.mappings() {
+            let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
+            builder.map_children(
+                input.into_iter(),
+                ast.token_trees_and_tokens().filter_map(only_nodes),
+            );
+            builder.finish(&mut mapping);
+        }
+
+        return ast;
+
+        fn only_nodes(element: NodeOrToken<ast::TokenTree, SyntaxToken>) -> Option<SyntaxNode> {
+            element.as_node().map(|it| it.syntax().clone())
+        }
+    }
+
+    pub fn token(&self, kind: SyntaxKind) -> SyntaxToken {
+        make::token(kind)
+    }
 }
diff --git a/crates/syntax/src/syntax_editor.rs b/crates/syntax/src/syntax_editor.rs
diff --git a/crates/syntax/src/syntax_editor/edit_algo.rs b/crates/syntax/src/syntax_editor/edit_algo.rs
