Make selections in assists with trailing/leading whitespace more forgiving

crates/ide_assists/src/assist_context.rs

@@ -14,7 +14,7 @@ use ide_db::{
 };
 use syntax::{
     algo::{self, find_node_at_offset, find_node_at_range},
-    AstNode, AstToken, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxNodePtr,
+    AstNode, AstToken, Direction, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxNodePtr,
     SyntaxToken, TextRange, TextSize, TokenAtOffset,
 };
 use text_edit::{TextEdit, TextEditBuilder};
@@ -57,6 +57,7 @@ pub(crate) struct AssistContext<'a> {
     pub(crate) config: &'a AssistConfig,
     pub(crate) sema: Semantics<'a, RootDatabase>,
     pub(crate) frange: FileRange,
+    trimmed_range: TextRange,
     source_file: SourceFile,
 }
 
@@ -67,7 +68,20 @@ impl<'a> AssistContext<'a> {
         frange: FileRange,
     ) -> AssistContext<'a> {
         let source_file = sema.parse(frange.file_id);
-        AssistContext { config, sema, frange, source_file }
+
+        let start = frange.range.start();
+        let end = frange.range.end();
+        let left = source_file.syntax().token_at_offset(start);
+        let right = source_file.syntax().token_at_offset(end);
+        let left =
+            left.right_biased().and_then(|t| algo::skip_whitespace_token(t, Direction::Next));
+        let right =
+            right.left_biased().and_then(|t| algo::skip_whitespace_token(t, Direction::Prev));
+        let left = left.map(|t| t.text_range().start()).unwrap_or(start).clamp(start, end);
+        let right = right.map(|t| t.text_range().end()).unwrap_or(end).clamp(start, end);
+        let trimmed_range = TextRange::new(left, right);
+
+        AssistContext { config, sema, frange, source_file, trimmed_range }
     }
 
     pub(crate) fn db(&self) -> &RootDatabase {
@@ -79,6 +93,12 @@ impl<'a> AssistContext<'a> {
         self.frange.range.start()
     }
 
+    /// Returns the selected range trimmed for whitespace tokens, that is the range will be snapped
+    /// to the nearest enclosed token.
+    pub(crate) fn selection_trimmed(&self) -> TextRange {
+        self.trimmed_range
+    }
+
     pub(crate) fn token_at_offset(&self) -> TokenAtOffset<SyntaxToken> {
         self.source_file.syntax().token_at_offset(self.offset())
     }
@@ -92,13 +112,15 @@ impl<'a> AssistContext<'a> {
         find_node_at_offset(self.source_file.syntax(), self.offset())
     }
     pub(crate) fn find_node_at_range<N: AstNode>(&self) -> Option<N> {
-        find_node_at_range(self.source_file.syntax(), self.frange.range)
+        find_node_at_range(self.source_file.syntax(), self.trimmed_range)
     }
     pub(crate) fn find_node_at_offset_with_descend<N: AstNode>(&self) -> Option<N> {
         self.sema.find_node_at_offset_with_descend(self.source_file.syntax(), self.offset())
     }
+    /// Returns the element covered by the selection range, this excludes trailing whitespace in the selection.
     pub(crate) fn covering_element(&self) -> SyntaxElement {
-        self.source_file.syntax().covering_element(self.frange.range)
+        self.source_file.syntax().covering_element(self.selection_trimmed())
+        // self.source_file.syntax().covering_element(self.frange.range)
     }
 }
 

crates/ide_assists/src/handlers/add_return_type.rs

@@ -107,11 +107,11 @@ fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrRepla
             let ret_range = TextRange::new(rparen_pos, ret_range_end);
             (FnType::Function, tail_expr, ret_range, action)
         };
-    let frange = ctx.frange.range;
+    let range = ctx.selection_trimmed();
-    if return_type_range.contains_range(frange) {
+    if return_type_range.contains_range(range) {
         cov_mark::hit!(cursor_in_ret_position);
         cov_mark::hit!(cursor_in_ret_position_closure);
-    } else if tail_expr.syntax().text_range().contains_range(frange) {
+    } else if tail_expr.syntax().text_range().contains_range(range) {
         cov_mark::hit!(cursor_on_tail);
         cov_mark::hit!(cursor_on_tail_closure);
     } else {

crates/ide_assists/src/handlers/apply_demorgan.rs

@@ -33,7 +33,7 @@ pub(crate) fn apply_demorgan(acc: &mut Assists, ctx: &AssistContext) -> Option<(
         _ => return None,
     };
 
-    let cursor_in_range = op_range.contains_range(ctx.frange.range);
+    let cursor_in_range = op_range.contains_range(ctx.selection_trimmed());
     if !cursor_in_range {
         return None;
     }

crates/ide_assists/src/handlers/convert_iter_for_each_to_for.rs

@@ -199,7 +199,7 @@ fn validate_method_call_expr(
     expr: ast::MethodCallExpr,
 ) -> Option<(ast::Expr, ast::Expr)> {
     let name_ref = expr.name_ref()?;
-    if name_ref.syntax().text_range().intersect(ctx.frange.range).is_none() {
+    if !name_ref.syntax().text_range().contains_range(ctx.selection_trimmed()) {
         cov_mark::hit!(test_for_each_not_applicable_invalid_cursor_pos);
         return None;
     }

crates/ide_assists/src/handlers/extract_function.rs

@@ -56,7 +56,7 @@ type FxIndexSet<T> = indexmap::IndexSet<T, BuildHasherDefault<FxHasher>>;
 // }
 // ```
 pub(crate) fn extract_function(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let range = ctx.frange.range;
+    let range = ctx.selection_trimmed();
     if range.is_empty() {
         return None;
     }

crates/ide_assists/src/handlers/extract_variable.rs

@@ -31,6 +31,7 @@ pub(crate) fn extract_variable(acc: &mut Assists, ctx: &AssistContext) -> Option
     if ctx.frange.range.is_empty() {
         return None;
     }
+
     let node = match ctx.covering_element() {
         NodeOrToken::Node(it) => it,
         NodeOrToken::Token(it) if it.kind() == COMMENT => {
@@ -238,7 +239,7 @@ fn foo() {
             extract_variable,
             r#"
 fn foo() {
-    $01 + 1$0;
+  $0  1 + 1$0;
 }"#,
             r#"
 fn foo() {
@@ -247,12 +248,12 @@ fn foo() {
         );
         check_assist(
             extract_variable,
-            "
+            r"
 fn foo() {
     $0{ let x = 0; x }$0
     something_else();
 }",
-            "
+            r"
 fn foo() {
     let $0var_name = { let x = 0; x };
     something_else();
@@ -264,11 +265,11 @@ fn foo() {
     fn test_extract_var_part_of_expr_stmt() {
         check_assist(
             extract_variable,
-            "
+            r"
 fn foo() {
     $01$0 + 1;
 }",
-            "
+            r"
 fn foo() {
     let $0var_name = 1;
     var_name + 1;

crates/ide_assists/src/handlers/flip_binexpr.rs

@@ -23,7 +23,7 @@ pub(crate) fn flip_binexpr(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
     let rhs = expr.rhs()?.syntax().clone();
     let op_range = expr.op_token()?.text_range();
     // The assist should be applied only if the cursor is on the operator
-    let cursor_in_range = op_range.contains_range(ctx.frange.range);
+    let cursor_in_range = op_range.contains_range(ctx.selection_trimmed());
     if !cursor_in_range {
         return None;
     }

crates/ide_assists/src/handlers/inline_call.rs

@@ -199,7 +199,7 @@ pub(crate) fn inline_call(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
     let param_list = fn_source.value.param_list()?;
 
     let FileRange { file_id, range } = fn_source.syntax().original_file_range(ctx.sema.db);
-    if file_id == ctx.frange.file_id && range.contains(ctx.frange.range.start()) {
+    if file_id == ctx.frange.file_id && range.contains(ctx.offset()) {
         cov_mark::hit!(inline_call_recursive);
         return None;
     }

crates/ide_assists/src/handlers/inline_local_variable.rs

@@ -1,7 +1,7 @@
 use either::Either;
 use hir::{PathResolution, Semantics};
 use ide_db::{
-    base_db::{FileId, FileRange},
+    base_db::FileId,
     defs::Definition,
     search::{FileReference, UsageSearchResult},
     RootDatabase,
@@ -33,7 +33,8 @@ use crate::{
 // }
 // ```
 pub(crate) fn inline_local_variable(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let FileRange { file_id, range } = ctx.frange;
+    let file_id = ctx.frange.file_id;
+    let range = ctx.selection_trimmed();
     let InlineData { let_stmt, delete_let, references, target } =
         if let Some(let_stmt) = ctx.find_node_at_offset::<ast::LetStmt>() {
             inline_let(&ctx.sema, let_stmt, range, file_id)

crates/ide_assists/src/handlers/invert_if.rs

@@ -29,7 +29,7 @@ pub(crate) fn invert_if(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let if_keyword = ctx.find_token_syntax_at_offset(T![if])?;
     let expr = ast::IfExpr::cast(if_keyword.parent()?)?;
     let if_range = if_keyword.text_range();
-    let cursor_in_range = if_range.contains_range(ctx.frange.range);
+    let cursor_in_range = if_range.contains_range(ctx.selection_trimmed());
     if !cursor_in_range {
         return None;
     }

crates/ide_assists/src/handlers/move_from_mod_rs.rs

@@ -27,7 +27,7 @@ pub(crate) fn move_from_mod_rs(acc: &mut Assists, ctx: &AssistContext) -> Option
     let source_file = ctx.find_node_at_offset::<ast::SourceFile>()?;
     let module = ctx.sema.to_module_def(ctx.frange.file_id)?;
     // Enable this assist if the user select all "meaningful" content in the source file
-    let trimmed_selected_range = trimmed_text_range(&source_file, ctx.frange.range);
+    let trimmed_selected_range = trimmed_text_range(&source_file, ctx.selection_trimmed());
     let trimmed_file_range = trimmed_text_range(&source_file, source_file.syntax().text_range());
     if !module.is_mod_rs(ctx.db()) {
         cov_mark::hit!(not_mod_rs);

crates/ide_assists/src/handlers/move_to_mod_rs.rs

@@ -27,7 +27,7 @@ pub(crate) fn move_to_mod_rs(acc: &mut Assists, ctx: &AssistContext) -> Option<(
     let source_file = ctx.find_node_at_offset::<ast::SourceFile>()?;
     let module = ctx.sema.to_module_def(ctx.frange.file_id)?;
     // Enable this assist if the user select all "meaningful" content in the source file
-    let trimmed_selected_range = trimmed_text_range(&source_file, ctx.frange.range);
+    let trimmed_selected_range = trimmed_text_range(&source_file, ctx.selection_trimmed());
     let trimmed_file_range = trimmed_text_range(&source_file, source_file.syntax().text_range());
     if module.is_mod_rs(ctx.db()) {
         cov_mark::hit!(already_mod_rs);

crates/ide_assists/src/handlers/replace_if_let_with_match.rs

@@ -48,7 +48,7 @@ pub(crate) fn replace_if_let_with_match(acc: &mut Assists, ctx: &AssistContext)
         if_expr.syntax().text_range().start(),
         if_expr.then_branch()?.syntax().text_range().start(),
     );
-    let cursor_in_range = available_range.contains_range(ctx.frange.range);
+    let cursor_in_range = available_range.contains_range(ctx.selection_trimmed());
     if !cursor_in_range {
         return None;
     }

crates/syntax/src/algo.rs

@@ -53,6 +53,16 @@ pub fn skip_trivia_token(mut token: SyntaxToken, direction: Direction) -> Option
     }
     Some(token)
 }
+/// Skip to next non `whitespace` token
+pub fn skip_whitespace_token(mut token: SyntaxToken, direction: Direction) -> Option<SyntaxToken> {
+    while token.kind() == SyntaxKind::WHITESPACE {
+        token = match direction {
+            Direction::Next => token.next_token()?,
+            Direction::Prev => token.prev_token()?,
+        }
+    }
+    Some(token)
+}
 
 /// Finds the first sibling in the given direction which is not `trivia`
 pub fn non_trivia_sibling(element: SyntaxElement, direction: Direction) -> Option<SyntaxElement> {