4 files changed, 308 insertions, 92 deletions

diff --git a/crates/assists/src/handlers/extract_struct_from_enum_variant.rs b/crates/assists/src/handlers/extract_struct_from_enum_variant.rs
index 40028fc01..e3ef04932 100644
--- a/crates/assists/src/handlers/extract_struct_from_enum_variant.rs
+++ b/crates/assists/src/handlers/extract_struct_from_enum_variant.rs
@@ -2,12 +2,16 @@ use std::iter;
 
 use either::Either;
 use hir::{AsName, Module, ModuleDef, Name, Variant};
-use ide_db::helpers::{
-    insert_use::{insert_use, ImportScope},
-    mod_path_to_ast,
+use ide_db::{
+    defs::Definition,
+    helpers::{
+        insert_use::{insert_use, ImportScope},
+        mod_path_to_ast,
+    },
+    search::FileReference,
+    RootDatabase,
 };
-use ide_db::{defs::Definition, search::Reference, RootDatabase};
-use rustc_hash::{FxHashMap, FxHashSet};
+use rustc_hash::FxHashSet;
 use syntax::{
     algo::{find_node_at_offset, SyntaxRewriter},
     ast::{self, edit::IndentLevel, make, AstNode, NameOwner, VisibilityOwner},
@@ -58,29 +62,29 @@ pub(crate) fn extract_struct_from_enum_variant(
             let mut visited_modules_set = FxHashSet::default();
             let current_module = enum_hir.module(ctx.db());
             visited_modules_set.insert(current_module);
-            let mut rewriters = FxHashMap::default();
-            for reference in usages {
-                let rewriter = rewriters
-                    .entry(reference.file_range.file_id)
-                    .or_insert_with(SyntaxRewriter::default);
-                let source_file = ctx.sema.parse(reference.file_range.file_id);
-                update_reference(
-                    ctx,
-                    rewriter,
-                    reference,
-                    &source_file,
-                    &enum_module_def,
-                    &variant_hir_name,
-                    &mut visited_modules_set,
-                );
-            }
-            let mut rewriter =
-                rewriters.remove(&ctx.frange.file_id).unwrap_or_else(SyntaxRewriter::default);
-            for (file_id, rewriter) in rewriters {
+            let mut def_rewriter = None;
+            for (file_id, references) in usages {
+                let mut rewriter = SyntaxRewriter::default();
+                let source_file = ctx.sema.parse(file_id);
+                for reference in references {
+                    update_reference(
+                        ctx,
+                        &mut rewriter,
+                        reference,
+                        &source_file,
+                        &enum_module_def,
+                        &variant_hir_name,
+                        &mut visited_modules_set,
+                    );
+                }
+                if file_id == ctx.frange.file_id {
+                    def_rewriter = Some(rewriter);
+                    continue;
+                }
                 builder.edit_file(file_id);
                 builder.rewrite(rewriter);
             }
-            builder.edit_file(ctx.frange.file_id);
+            let mut rewriter = def_rewriter.unwrap_or_default();
             update_variant(&mut rewriter, &variant);
             extract_struct_def(
                 &mut rewriter,
@@ -90,6 +94,7 @@ pub(crate) fn extract_struct_from_enum_variant(
                 &variant.parent_enum().syntax().clone().into(),
                 enum_ast.visibility(),
             );
+            builder.edit_file(ctx.frange.file_id);
             builder.rewrite(rewriter);
         },
     )
@@ -205,13 +210,13 @@ fn update_variant(rewriter: &mut SyntaxRewriter, variant: &ast::Variant) -> Opti
 fn update_reference(
     ctx: &AssistContext,
     rewriter: &mut SyntaxRewriter,
-    reference: Reference,
+    reference: FileReference,
     source_file: &SourceFile,
     enum_module_def: &ModuleDef,
     variant_hir_name: &Name,
     visited_modules_set: &mut FxHashSet<Module>,
 ) -> Option<()> {
-    let offset = reference.file_range.range.start();
+    let offset = reference.range.start();
     let (segment, expr) = if let Some(path_expr) =
         find_node_at_offset::<ast::PathExpr>(source_file.syntax(), offset)
     {
diff --git a/crates/assists/src/handlers/inline_local_variable.rs b/crates/assists/src/handlers/inline_local_variable.rs
index d559be9cb..dc798daaa 100644
--- a/crates/assists/src/handlers/inline_local_variable.rs
+++ b/crates/assists/src/handlers/inline_local_variable.rs
@@ -1,4 +1,7 @@
-use ide_db::{defs::Definition, search::ReferenceKind};
+use ide_db::{
+    defs::Definition,
+    search::{FileReference, ReferenceKind},
+};
 use syntax::{
     ast::{self, AstNode, AstToken},
     TextRange,
@@ -44,8 +47,8 @@ pub(crate) fn inline_local_variable(acc: &mut Assists, ctx: &AssistContext) -> O
 
     let def = ctx.sema.to_def(&bind_pat)?;
     let def = Definition::Local(def);
-    let refs = def.usages(&ctx.sema).all();
-    if refs.is_empty() {
+    let usages = def.usages(&ctx.sema).all();
+    if usages.is_empty() {
         mark::hit!(test_not_applicable_if_variable_unused);
         return None;
     };
@@ -63,48 +66,45 @@ pub(crate) fn inline_local_variable(acc: &mut Assists, ctx: &AssistContext) -> O
         let_stmt.syntax().text_range()
     };
 
-    let mut wrap_in_parens = vec![true; refs.len()];
-
-    for (i, desc) in refs.iter().enumerate() {
-        let usage_node = ctx
-            .covering_node_for_range(desc.file_range.range)
-            .ancestors()
-            .find_map(ast::PathExpr::cast)?;
-        let usage_parent_option = usage_node.syntax().parent().and_then(ast::Expr::cast);
-        let usage_parent = match usage_parent_option {
-            Some(u) => u,
-            None => {
-                wrap_in_parens[i] = false;
-                continue;
-            }
-        };
-
-        wrap_in_parens[i] = match (&initializer_expr, usage_parent) {
-            (ast::Expr::CallExpr(_), _)
-            | (ast::Expr::IndexExpr(_), _)
-            | (ast::Expr::MethodCallExpr(_), _)
-            | (ast::Expr::FieldExpr(_), _)
-            | (ast::Expr::TryExpr(_), _)
-            | (ast::Expr::RefExpr(_), _)
-            | (ast::Expr::Literal(_), _)
-            | (ast::Expr::TupleExpr(_), _)
-            | (ast::Expr::ArrayExpr(_), _)
-            | (ast::Expr::ParenExpr(_), _)
-            | (ast::Expr::PathExpr(_), _)
-            | (ast::Expr::BlockExpr(_), _)
-            | (ast::Expr::EffectExpr(_), _)
-            | (_, ast::Expr::CallExpr(_))
-            | (_, ast::Expr::TupleExpr(_))
-            | (_, ast::Expr::ArrayExpr(_))
-            | (_, ast::Expr::ParenExpr(_))
-            | (_, ast::Expr::ForExpr(_))
-            | (_, ast::Expr::WhileExpr(_))
-            | (_, ast::Expr::BreakExpr(_))
-            | (_, ast::Expr::ReturnExpr(_))
-            | (_, ast::Expr::MatchExpr(_)) => false,
-            _ => true,
-        };
-    }
+    let wrap_in_parens = usages
+        .references
+        .values()
+        .flatten()
+        .map(|&FileReference { range, .. }| {
+            let usage_node =
+                ctx.covering_node_for_range(range).ancestors().find_map(ast::PathExpr::cast)?;
+            let usage_parent_option = usage_node.syntax().parent().and_then(ast::Expr::cast);
+            let usage_parent = match usage_parent_option {
+                Some(u) => u,
+                None => return Ok(false),
+            };
+
+            Ok(!matches!((&initializer_expr, usage_parent),
+                (ast::Expr::CallExpr(_), _)
+                | (ast::Expr::IndexExpr(_), _)
+                | (ast::Expr::MethodCallExpr(_), _)
+                | (ast::Expr::FieldExpr(_), _)
+                | (ast::Expr::TryExpr(_), _)
+                | (ast::Expr::RefExpr(_), _)
+                | (ast::Expr::Literal(_), _)
+                | (ast::Expr::TupleExpr(_), _)
+                | (ast::Expr::ArrayExpr(_), _)
+                | (ast::Expr::ParenExpr(_), _)
+                | (ast::Expr::PathExpr(_), _)
+                | (ast::Expr::BlockExpr(_), _)
+                | (ast::Expr::EffectExpr(_), _)
+                | (_, ast::Expr::CallExpr(_))
+                | (_, ast::Expr::TupleExpr(_))
+                | (_, ast::Expr::ArrayExpr(_))
+                | (_, ast::Expr::ParenExpr(_))
+                | (_, ast::Expr::ForExpr(_))
+                | (_, ast::Expr::WhileExpr(_))
+                | (_, ast::Expr::BreakExpr(_))
+                | (_, ast::Expr::ReturnExpr(_))
+                | (_, ast::Expr::MatchExpr(_))
+            ))
+        })
+        .collect::<Result<Vec<_>, _>>()?;
 
     let init_str = initializer_expr.syntax().text().to_string();
     let init_in_paren = format!("({})", &init_str);
@@ -116,15 +116,16 @@ pub(crate) fn inline_local_variable(acc: &mut Assists, ctx: &AssistContext) -> O
         target,
         move |builder| {
             builder.delete(delete_range);
-            for (desc, should_wrap) in refs.iter().zip(wrap_in_parens) {
+            for (reference, should_wrap) in usages.references.values().flatten().zip(wrap_in_parens)
+            {
                 let replacement =
                     if should_wrap { init_in_paren.clone() } else { init_str.clone() };
-                match desc.kind {
+                match reference.kind {
                     ReferenceKind::FieldShorthandForLocal => {
                         mark::hit!(inline_field_shorthand);
-                        builder.insert(desc.file_range.range.end(), format!(": {}", replacement))
+                        builder.insert(reference.range.end(), format!(": {}", replacement))
                     }
-                    _ => builder.replace(desc.file_range.range, replacement),
+                    _ => builder.replace(reference.range, replacement),
                 }
             }
         },
diff --git a/crates/assists/src/handlers/remove_unused_param.rs b/crates/assists/src/handlers/remove_unused_param.rs
index 56e8b5229..c961680e2 100644
--- a/crates/assists/src/handlers/remove_unused_param.rs
+++ b/crates/assists/src/handlers/remove_unused_param.rs
@@ -1,8 +1,8 @@
-use ide_db::{defs::Definition, search::Reference};
+use ide_db::{base_db::FileId, defs::Definition, search::FileReference};
 use syntax::{
     algo::find_node_at_range,
     ast::{self, ArgListOwner},
-    AstNode, SyntaxKind, SyntaxNode, TextRange, T,
+    AstNode, SourceFile, SyntaxKind, SyntaxNode, TextRange, T,
 };
 use test_utils::mark;
 use SyntaxKind::WHITESPACE;
@@ -58,32 +58,41 @@ pub(crate) fn remove_unused_param(acc: &mut Assists, ctx: &AssistContext) -> Opt
         param.syntax().text_range(),
         |builder| {
             builder.delete(range_to_remove(param.syntax()));
-            for usage in fn_def.usages(&ctx.sema).all() {
-                process_usage(ctx, builder, usage, param_position);
+            for (file_id, references) in fn_def.usages(&ctx.sema).all() {
+                process_usages(ctx, builder, file_id, references, param_position);
             }
         },
     )
 }
 
-fn process_usage(
+fn process_usages(
     ctx: &AssistContext,
     builder: &mut AssistBuilder,
-    usage: Reference,
+    file_id: FileId,
+    references: Vec<FileReference>,
     arg_to_remove: usize,
-) -> Option<()> {
-    let source_file = ctx.sema.parse(usage.file_range.file_id);
-    let call_expr: ast::CallExpr =
-        find_node_at_range(source_file.syntax(), usage.file_range.range)?;
+) {
+    let source_file = ctx.sema.parse(file_id);
+    builder.edit_file(file_id);
+    for usage in references {
+        if let Some(text_range) = process_usage(&source_file, usage, arg_to_remove) {
+            builder.delete(text_range);
+        }
+    }
+}
+
+fn process_usage(
+    source_file: &SourceFile,
+    FileReference { range, .. }: FileReference,
+    arg_to_remove: usize,
+) -> Option<TextRange> {
+    let call_expr: ast::CallExpr = find_node_at_range(source_file.syntax(), range)?;
     let call_expr_range = call_expr.expr()?.syntax().text_range();
-    if !call_expr_range.contains_range(usage.file_range.range) {
+    if !call_expr_range.contains_range(range) {
         return None;
     }
     let arg = call_expr.arg_list()?.args().nth(arg_to_remove)?;
-
-    builder.edit_file(usage.file_range.file_id);
-    builder.delete(range_to_remove(arg.syntax()));
-
-    Some(())
+    Some(range_to_remove(arg.syntax()))
 }
 
 fn range_to_remove(node: &SyntaxNode) -> TextRange {
diff --git a/crates/assists/src/handlers/reorder_impl.rs b/crates/assists/src/handlers/reorder_impl.rs
new file mode 100644
index 000000000..309f291c8
--- /dev/null
+++ b/crates/assists/src/handlers/reorder_impl.rs
@@ -0,0 +1,201 @@
+use itertools::Itertools;
+use rustc_hash::FxHashMap;
+
+use hir::{PathResolution, Semantics};
+use ide_db::RootDatabase;
+use syntax::{
+    algo,
+    ast::{self, NameOwner},
+    AstNode,
+};
+use test_utils::mark;
+
+use crate::{AssistContext, AssistId, AssistKind, Assists};
+
+// Assist: reorder_impl
+//
+// Reorder the methods of an `impl Trait`. The methods will be ordered
+// in the same order as in the trait definition.
+//
+// ```
+// trait Foo {
+//     fn a() {}
+//     fn b() {}
+//     fn c() {}
+// }
+//
+// struct Bar;
+// $0impl Foo for Bar {
+//     fn b() {}
+//     fn c() {}
+//     fn a() {}
+// }
+// ```
+// ->
+// ```
+// trait Foo {
+//     fn a() {}
+//     fn b() {}
+//     fn c() {}
+// }
+//
+// struct Bar;
+// impl Foo for Bar {
+//     fn a() {}
+//     fn b() {}
+//     fn c() {}
+// }
+// ```
+//
+pub(crate) fn reorder_impl(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    let impl_ast = ctx.find_node_at_offset::<ast::Impl>()?;
+    let items = impl_ast.assoc_item_list()?;
+    let methods = get_methods(&items);
+
+    let path = impl_ast
+        .trait_()
+        .and_then(|t| match t {
+            ast::Type::PathType(path) => Some(path),
+            _ => None,
+        })?
+        .path()?;
+
+    let ranks = compute_method_ranks(&path, ctx)?;
+    let sorted: Vec<_> = methods
+        .iter()
+        .cloned()
+        .sorted_by_key(|f| {
+            f.name().and_then(|n| ranks.get(&n.to_string()).copied()).unwrap_or(usize::max_value())
+        })
+        .collect();
+
+    // Don't edit already sorted methods:
+    if methods == sorted {
+        mark::hit!(not_applicable_if_sorted);
+        return None;
+    }
+
+    let target = items.syntax().text_range();
+    acc.add(AssistId("reorder_impl", AssistKind::RefactorRewrite), "Sort methods", target, |edit| {
+        let mut rewriter = algo::SyntaxRewriter::default();
+        for (old, new) in methods.iter().zip(&sorted) {
+            rewriter.replace(old.syntax(), new.syntax());
+        }
+        edit.rewrite(rewriter);
+    })
+}
+
+fn compute_method_ranks(path: &ast::Path, ctx: &AssistContext) -> Option<FxHashMap<String, usize>> {
+    let td = trait_definition(path, &ctx.sema)?;
+
+    Some(
+        td.items(ctx.db())
+            .iter()
+            .flat_map(|i| match i {
+                hir::AssocItem::Function(f) => Some(f),
+                _ => None,
+            })
+            .enumerate()
+            .map(|(idx, func)| ((func.name(ctx.db()).to_string(), idx)))
+            .collect(),
+    )
+}
+
+fn trait_definition(path: &ast::Path, sema: &Semantics<RootDatabase>) -> Option<hir::Trait> {
+    match sema.resolve_path(path)? {
+        PathResolution::Def(hir::ModuleDef::Trait(trait_)) => Some(trait_),
+        _ => None,
+    }
+}
+
+fn get_methods(items: &ast::AssocItemList) -> Vec<ast::Fn> {
+    items
+        .assoc_items()
+        .flat_map(|i| match i {
+            ast::AssocItem::Fn(f) => Some(f),
+            _ => None,
+        })
+        .filter(|f| f.name().is_some())
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use test_utils::mark;
+
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    #[test]
+    fn not_applicable_if_sorted() {
+        mark::check!(not_applicable_if_sorted);
+        check_assist_not_applicable(
+            reorder_impl,
+            r#"
+trait Bar {
+    fn a() {}
+    fn z() {}
+    fn b() {}
+}
+struct Foo;
+$0impl Bar for Foo {
+    fn a() {}
+    fn z() {}
+    fn b() {}
+}
+        "#,
+        )
+    }
+
+    #[test]
+    fn not_applicable_if_empty() {
+        check_assist_not_applicable(
+            reorder_impl,
+            r#"
+trait Bar {};
+struct Foo;
+$0impl Bar for Foo {}
+        "#,
+        )
+    }
+
+    #[test]
+    fn reorder_impl_trait_methods() {
+        check_assist(
+            reorder_impl,
+            r#"
+trait Bar {
+    fn a() {}
+    fn c() {}
+    fn b() {}
+    fn d() {}
+}
+
+struct Foo;
+$0impl Bar for Foo {
+    fn d() {}
+    fn b() {}
+    fn c() {}
+    fn a() {}
+}
+        "#,
+            r#"
+trait Bar {
+    fn a() {}
+    fn c() {}
+    fn b() {}
+    fn d() {}
+}
+
+struct Foo;
+impl Bar for Foo {
+    fn a() {}
+    fn c() {}
+    fn b() {}
+    fn d() {}
+}
+        "#,
+        )
+    }
+}