4 files changed, 298 insertions, 38 deletions

diff --git a/crates/ide_assists/src/handlers/extract_function.rs b/crates/ide_assists/src/handlers/extract_function.rs
index 059414274..78a57fbdc 100644
--- a/crates/ide_assists/src/handlers/extract_function.rs
+++ b/crates/ide_assists/src/handlers/extract_function.rs
@@ -599,7 +599,12 @@ fn extraction_target(node: &SyntaxNode, selection_range: TextRange) -> Option<Fu
     // we have selected a few statements in a block
     // so covering_element returns the whole block
     if node.kind() == BLOCK_EXPR {
-        let body = FunctionBody::from_range(node.clone(), selection_range);
+        // Extract the full statements.
+        let statements_range = node
+            .children()
+            .filter(|c| selection_range.intersect(c.text_range()).is_some())
+            .fold(selection_range, |acc, c| acc.cover(c.text_range()));
+        let body = FunctionBody::from_range(node.clone(), statements_range);
         if body.is_some() {
             return body;
         }
@@ -610,7 +615,8 @@ fn extraction_target(node: &SyntaxNode, selection_range: TextRange) -> Option<Fu
     // so we try to expand covering_element to parent and repeat the previous
     if let Some(parent) = node.parent() {
         if parent.kind() == BLOCK_EXPR {
-            let body = FunctionBody::from_range(parent, selection_range);
+            // Extract the full statement.
+            let body = FunctionBody::from_range(parent, node.text_range());
             if body.is_some() {
                 return body;
             }
@@ -1785,6 +1791,60 @@ fn $0fun_name() -> i32 {
     }
 
     #[test]
+    fn extract_partial_block_single_line() {
+        check_assist(
+            extract_function,
+            r#"
+fn foo() {
+    let n = 1;
+    let mut v = $0n * n;$0
+    v += 1;
+}"#,
+            r#"
+fn foo() {
+    let n = 1;
+    let mut v = fun_name(n);
+    v += 1;
+}
+
+fn $0fun_name(n: i32) -> i32 {
+    let mut v = n * n;
+    v
+}"#,
+        );
+    }
+
+    #[test]
+    fn extract_partial_block() {
+        check_assist(
+            extract_function,
+            r#"
+fn foo() {
+    let m = 2;
+    let n = 1;
+    let mut v = m $0* n;
+    let mut w = 3;$0
+    v += 1;
+    w += 1;
+}"#,
+            r#"
+fn foo() {
+    let m = 2;
+    let n = 1;
+    let (mut v, mut w) = fun_name(m, n);
+    v += 1;
+    w += 1;
+}
+
+fn $0fun_name(m: i32, n: i32) -> (i32, i32) {
+    let mut v = m * n;
+    let mut w = 3;
+    (v, w)
+}"#,
+        );
+    }
+
+    #[test]
     fn argument_form_expr() {
         check_assist(
             extract_function,
diff --git a/crates/ide_assists/src/handlers/introduce_named_lifetime.rs b/crates/ide_assists/src/handlers/introduce_named_lifetime.rs
index 02782eb6d..9f4f71d6c 100644
--- a/crates/ide_assists/src/handlers/introduce_named_lifetime.rs
+++ b/crates/ide_assists/src/handlers/introduce_named_lifetime.rs
@@ -1,7 +1,8 @@
 use rustc_hash::FxHashSet;
 use syntax::{
-    ast::{self, GenericParamsOwner, NameOwner},
-    AstNode, TextRange, TextSize,
+    ast::{self, edit_in_place::GenericParamsOwnerEdit, make, GenericParamsOwner},
+    ted::{self, Position},
+    AstNode, TextRange,
 };
 
 use crate::{assist_context::AssistBuilder, AssistContext, AssistId, AssistKind, Assists};
@@ -37,10 +38,12 @@ static ASSIST_LABEL: &str = "Introduce named lifetime";
 pub(crate) fn introduce_named_lifetime(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let lifetime =
         ctx.find_node_at_offset::<ast::Lifetime>().filter(|lifetime| lifetime.text() == "'_")?;
+    let lifetime_loc = lifetime.lifetime_ident_token()?.text_range();
+
     if let Some(fn_def) = lifetime.syntax().ancestors().find_map(ast::Fn::cast) {
-        generate_fn_def_assist(acc, &fn_def, lifetime.lifetime_ident_token()?.text_range())
+        generate_fn_def_assist(acc, fn_def, lifetime_loc, lifetime)
     } else if let Some(impl_def) = lifetime.syntax().ancestors().find_map(ast::Impl::cast) {
-        generate_impl_def_assist(acc, &impl_def, lifetime.lifetime_ident_token()?.text_range())
+        generate_impl_def_assist(acc, impl_def, lifetime_loc, lifetime)
     } else {
         None
     }
@@ -49,26 +52,26 @@ pub(crate) fn introduce_named_lifetime(acc: &mut Assists, ctx: &AssistContext) -
 /// Generate the assist for the fn def case
 fn generate_fn_def_assist(
     acc: &mut Assists,
-    fn_def: &ast::Fn,
+    fn_def: ast::Fn,
     lifetime_loc: TextRange,
+    lifetime: ast::Lifetime,
 ) -> Option<()> {
     let param_list: ast::ParamList = fn_def.param_list()?;
-    let new_lifetime_param = generate_unique_lifetime_param_name(&fn_def.generic_param_list())?;
-    let end_of_fn_ident = fn_def.name()?.ident_token()?.text_range().end();
+    let new_lifetime_param = generate_unique_lifetime_param_name(fn_def.generic_param_list())?;
     let self_param =
         // use the self if it's a reference and has no explicit lifetime
         param_list.self_param().filter(|p| p.lifetime().is_none() && p.amp_token().is_some());
     // compute the location which implicitly has the same lifetime as the anonymous lifetime
     let loc_needing_lifetime = if let Some(self_param) = self_param {
         // if we have a self reference, use that
-        Some(self_param.name()?.syntax().text_range().start())
+        Some(NeedsLifetime::SelfParam(self_param))
     } else {
         // otherwise, if there's a single reference parameter without a named liftime, use that
         let fn_params_without_lifetime: Vec<_> = param_list
             .params()
             .filter_map(|param| match param.ty() {
                 Some(ast::Type::RefType(ascribed_type)) if ascribed_type.lifetime().is_none() => {
-                    Some(ascribed_type.amp_token()?.text_range().end())
+                    Some(NeedsLifetime::RefType(ascribed_type))
                 }
                 _ => None,
             })
@@ -81,30 +84,46 @@ fn generate_fn_def_assist(
         }
     };
     acc.add(AssistId(ASSIST_NAME, AssistKind::Refactor), ASSIST_LABEL, lifetime_loc, |builder| {
-        add_lifetime_param(fn_def, builder, end_of_fn_ident, new_lifetime_param);
-        builder.replace(lifetime_loc, format!("'{}", new_lifetime_param));
-        loc_needing_lifetime.map(|loc| builder.insert(loc, format!("'{} ", new_lifetime_param)));
+        let fn_def = builder.make_ast_mut(fn_def);
+        let lifetime = builder.make_ast_mut(lifetime);
+        let loc_needing_lifetime =
+            loc_needing_lifetime.and_then(|it| it.make_mut(builder).to_position());
+
+        add_lifetime_param(fn_def.get_or_create_generic_param_list(), new_lifetime_param);
+        ted::replace(
+            lifetime.syntax(),
+            make_ast_lifetime(new_lifetime_param).clone_for_update().syntax(),
+        );
+        loc_needing_lifetime.map(|position| {
+            ted::insert(position, make_ast_lifetime(new_lifetime_param).clone_for_update().syntax())
+        });
     })
 }
 
 /// Generate the assist for the impl def case
 fn generate_impl_def_assist(
     acc: &mut Assists,
-    impl_def: &ast::Impl,
+    impl_def: ast::Impl,
     lifetime_loc: TextRange,
+    lifetime: ast::Lifetime,
 ) -> Option<()> {
-    let new_lifetime_param = generate_unique_lifetime_param_name(&impl_def.generic_param_list())?;
-    let end_of_impl_kw = impl_def.impl_token()?.text_range().end();
+    let new_lifetime_param = generate_unique_lifetime_param_name(impl_def.generic_param_list())?;
     acc.add(AssistId(ASSIST_NAME, AssistKind::Refactor), ASSIST_LABEL, lifetime_loc, |builder| {
-        add_lifetime_param(impl_def, builder, end_of_impl_kw, new_lifetime_param);
-        builder.replace(lifetime_loc, format!("'{}", new_lifetime_param));
+        let impl_def = builder.make_ast_mut(impl_def);
+        let lifetime = builder.make_ast_mut(lifetime);
+
+        add_lifetime_param(impl_def.get_or_create_generic_param_list(), new_lifetime_param);
+        ted::replace(
+            lifetime.syntax(),
+            make_ast_lifetime(new_lifetime_param).clone_for_update().syntax(),
+        );
     })
 }
 
 /// Given a type parameter list, generate a unique lifetime parameter name
 /// which is not in the list
 fn generate_unique_lifetime_param_name(
-    existing_type_param_list: &Option<ast::GenericParamList>,
+    existing_type_param_list: Option<ast::GenericParamList>,
 ) -> Option<char> {
     match existing_type_param_list {
         Some(type_params) => {
@@ -118,25 +137,37 @@ fn generate_unique_lifetime_param_name(
     }
 }
 
-/// Add the lifetime param to `builder`. If there are type parameters in `type_params_owner`, add it to the end. Otherwise
-/// add new type params brackets with the lifetime parameter at `new_type_params_loc`.
-fn add_lifetime_param<TypeParamsOwner: ast::GenericParamsOwner>(
-    type_params_owner: &TypeParamsOwner,
-    builder: &mut AssistBuilder,
-    new_type_params_loc: TextSize,
-    new_lifetime_param: char,
-) {
-    match type_params_owner.generic_param_list() {
-        // add the new lifetime parameter to an existing type param list
-        Some(type_params) => {
-            builder.insert(
-                (u32::from(type_params.syntax().text_range().end()) - 1).into(),
-                format!(", '{}", new_lifetime_param),
-            );
+fn add_lifetime_param(type_params: ast::GenericParamList, new_lifetime_param: char) {
+    let generic_param =
+        make::generic_param(format!("'{}", new_lifetime_param), None).clone_for_update();
+    type_params.add_generic_param(generic_param);
+}
+
+fn make_ast_lifetime(new_lifetime_param: char) -> ast::Lifetime {
+    make::generic_param(format!("'{}", new_lifetime_param), None)
+        .syntax()
+        .descendants()
+        .find_map(ast::Lifetime::cast)
+        .unwrap()
+}
+
+enum NeedsLifetime {
+    SelfParam(ast::SelfParam),
+    RefType(ast::RefType),
+}
+
+impl NeedsLifetime {
+    fn make_mut(self, builder: &mut AssistBuilder) -> Self {
+        match self {
+            Self::SelfParam(it) => Self::SelfParam(builder.make_ast_mut(it)),
+            Self::RefType(it) => Self::RefType(builder.make_ast_mut(it)),
         }
-        // create a new type param list containing only the new lifetime parameter
-        None => {
-            builder.insert(new_type_params_loc, format!("<'{}>", new_lifetime_param));
+    }
+
+    fn to_position(self) -> Option<Position> {
+        match self {
+            Self::SelfParam(it) => Some(Position::after(it.amp_token()?)),
+            Self::RefType(it) => Some(Position::after(it.amp_token()?)),
         }
     }
 }
@@ -312,4 +343,13 @@ mod tests {
             r#"fn my_fun<'other, 'a>(self, f: &'a Foo, b: &'other Bar) -> X<'a>"#,
         );
     }
+
+    #[test]
+    fn test_function_add_lifetime_to_self_ref_mut() {
+        check_assist(
+            introduce_named_lifetime,
+            r#"fn foo(&mut self) -> &'_$0 ()"#,
+            r#"fn foo<'a>(&'a mut self) -> &'a ()"#,
+        );
+    }
 }
diff --git a/crates/syntax/src/ast/edit_in_place.rs b/crates/syntax/src/ast/edit_in_place.rs
index 529bd0eb1..04f97f368 100644
--- a/crates/syntax/src/ast/edit_in_place.rs
+++ b/crates/syntax/src/ast/edit_in_place.rs
@@ -14,10 +14,29 @@ use crate::{
 use super::NameOwner;
 
 pub trait GenericParamsOwnerEdit: ast::GenericParamsOwner + AstNodeEdit {
+    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList;
     fn get_or_create_where_clause(&self) -> ast::WhereClause;
 }
 
 impl GenericParamsOwnerEdit for ast::Fn {
+    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
+        match self.generic_param_list() {
+            Some(it) => it,
+            None => {
+                let position = if let Some(name) = self.name() {
+                    Position::after(name.syntax)
+                } else if let Some(fn_token) = self.fn_token() {
+                    Position::after(fn_token)
+                } else if let Some(param_list) = self.param_list() {
+                    Position::before(param_list.syntax)
+                } else {
+                    Position::last_child_of(self.syntax())
+                };
+                create_generic_param_list(position)
+            }
+        }
+    }
+
     fn get_or_create_where_clause(&self) -> WhereClause {
         if self.where_clause().is_none() {
             let position = if let Some(ty) = self.ret_type() {
@@ -34,6 +53,20 @@ impl GenericParamsOwnerEdit for ast::Fn {
 }
 
 impl GenericParamsOwnerEdit for ast::Impl {
+    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
+        match self.generic_param_list() {
+            Some(it) => it,
+            None => {
+                let position = if let Some(imp_token) = self.impl_token() {
+                    Position::after(imp_token)
+                } else {
+                    Position::last_child_of(self.syntax())
+                };
+                create_generic_param_list(position)
+            }
+        }
+    }
+
     fn get_or_create_where_clause(&self) -> WhereClause {
         if self.where_clause().is_none() {
             let position = if let Some(items) = self.assoc_item_list() {
@@ -48,6 +81,22 @@ impl GenericParamsOwnerEdit for ast::Impl {
 }
 
 impl GenericParamsOwnerEdit for ast::Trait {
+    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
+        match self.generic_param_list() {
+            Some(it) => it,
+            None => {
+                let position = if let Some(name) = self.name() {
+                    Position::after(name.syntax)
+                } else if let Some(trait_token) = self.trait_token() {
+                    Position::after(trait_token)
+                } else {
+                    Position::last_child_of(self.syntax())
+                };
+                create_generic_param_list(position)
+            }
+        }
+    }
+
     fn get_or_create_where_clause(&self) -> WhereClause {
         if self.where_clause().is_none() {
             let position = if let Some(items) = self.assoc_item_list() {
@@ -62,6 +111,22 @@ impl GenericParamsOwnerEdit for ast::Trait {
 }
 
 impl GenericParamsOwnerEdit for ast::Struct {
+    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
+        match self.generic_param_list() {
+            Some(it) => it,
+            None => {
+                let position = if let Some(name) = self.name() {
+                    Position::after(name.syntax)
+                } else if let Some(struct_token) = self.struct_token() {
+                    Position::after(struct_token)
+                } else {
+                    Position::last_child_of(self.syntax())
+                };
+                create_generic_param_list(position)
+            }
+        }
+    }
+
     fn get_or_create_where_clause(&self) -> WhereClause {
         if self.where_clause().is_none() {
             let tfl = self.field_list().and_then(|fl| match fl {
@@ -84,6 +149,22 @@ impl GenericParamsOwnerEdit for ast::Struct {
 }
 
 impl GenericParamsOwnerEdit for ast::Enum {
+    fn get_or_create_generic_param_list(&self) -> ast::GenericParamList {
+        match self.generic_param_list() {
+            Some(it) => it,
+            None => {
+                let position = if let Some(name) = self.name() {
+                    Position::after(name.syntax)
+                } else if let Some(enum_token) = self.enum_token() {
+                    Position::after(enum_token)
+                } else {
+                    Position::last_child_of(self.syntax())
+                };
+                create_generic_param_list(position)
+            }
+        }
+    }
+
     fn get_or_create_where_clause(&self) -> WhereClause {
         if self.where_clause().is_none() {
             let position = if let Some(gpl) = self.generic_param_list() {
@@ -104,6 +185,37 @@ fn create_where_clause(position: Position) {
     ted::insert(position, where_clause.syntax());
 }
 
+fn create_generic_param_list(position: Position) -> ast::GenericParamList {
+    let gpl = make::generic_param_list(empty()).clone_for_update();
+    ted::insert_raw(position, gpl.syntax());
+    gpl
+}
+
+impl ast::GenericParamList {
+    pub fn add_generic_param(&self, generic_param: ast::GenericParam) {
+        match self.generic_params().last() {
+            Some(last_param) => {
+                let mut elems = Vec::new();
+                if !last_param
+                    .syntax()
+                    .siblings_with_tokens(Direction::Next)
+                    .any(|it| it.kind() == T![,])
+                {
+                    elems.push(make::token(T![,]).into());
+                    elems.push(make::tokens::single_space().into());
+                };
+                elems.push(generic_param.syntax().clone().into());
+                let after_last_param = Position::after(last_param.syntax());
+                ted::insert_all(after_last_param, elems);
+            }
+            None => {
+                let after_l_angle = Position::after(self.l_angle_token().unwrap());
+                ted::insert(after_l_angle, generic_param.syntax())
+            }
+        }
+    }
+}
+
 impl ast::WhereClause {
     pub fn add_predicate(&self, predicate: ast::WherePred) {
         if let Some(pred) = self.predicates().last() {
@@ -164,3 +276,44 @@ impl ast::Use {
         ted::remove(self.syntax())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::fmt;
+
+    use crate::SourceFile;
+
+    use super::*;
+
+    fn ast_mut_from_text<N: AstNode>(text: &str) -> N {
+        let parse = SourceFile::parse(text);
+        parse.tree().syntax().descendants().find_map(N::cast).unwrap().clone_for_update()
+    }
+
+    #[test]
+    fn test_create_generic_param_list() {
+        fn check_create_gpl<N: GenericParamsOwnerEdit + fmt::Display>(before: &str, after: &str) {
+            let gpl_owner = ast_mut_from_text::<N>(before);
+            gpl_owner.get_or_create_generic_param_list();
+            assert_eq!(gpl_owner.to_string(), after);
+        }
+
+        check_create_gpl::<ast::Fn>("fn foo", "fn foo<>");
+        check_create_gpl::<ast::Fn>("fn foo() {}", "fn foo<>() {}");
+
+        check_create_gpl::<ast::Impl>("impl", "impl<>");
+        check_create_gpl::<ast::Impl>("impl Struct {}", "impl<> Struct {}");
+        check_create_gpl::<ast::Impl>("impl Trait for Struct {}", "impl<> Trait for Struct {}");
+
+        check_create_gpl::<ast::Trait>("trait Trait<>", "trait Trait<>");
+        check_create_gpl::<ast::Trait>("trait Trait<> {}", "trait Trait<> {}");
+
+        check_create_gpl::<ast::Struct>("struct A", "struct A<>");
+        check_create_gpl::<ast::Struct>("struct A;", "struct A<>;");
+        check_create_gpl::<ast::Struct>("struct A();", "struct A<>();");
+        check_create_gpl::<ast::Struct>("struct A {}", "struct A<> {}");
+
+        check_create_gpl::<ast::Enum>("enum E", "enum E<>");
+        check_create_gpl::<ast::Enum>("enum E {", "enum E<> {");
+    }
+}
diff --git a/crates/syntax/src/ted.rs b/crates/syntax/src/ted.rs
index 177d4ff67..450f2e447 100644
--- a/crates/syntax/src/ted.rs
+++ b/crates/syntax/src/ted.rs
@@ -165,6 +165,13 @@ fn ws_between(left: &SyntaxElement, right: &SyntaxElement) -> Option<SyntaxToken
     if right.kind() == T![;] || right.kind() == T![,] {
         return None;
     }
+    if left.kind() == T![<] || right.kind() == T![>] {
+        return None;
+    }
+    if left.kind() == T![&] && right.kind() == SyntaxKind::LIFETIME {
+        return None;
+    }
+
     if right.kind() == SyntaxKind::USE {
         let indent = IndentLevel::from_element(left);
         return Some(make::tokens::whitespace(&format!("\n{}", indent)));