Merge pull request #19271 from vishruth-thimmaiah/for_loop_to_while_let

feat: for loop to while let assist

5 files changed, 467 insertions, 2 deletions

diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/convert_for_to_while_let.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/convert_for_to_while_let.rs
new file mode 100644
index 00000000000..2d2751de99d
--- /dev/null
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/convert_for_to_while_let.rs
@@ -0,0 +1,422 @@
+use hir::{
+    sym::{self},
+    Name,
+};
+use ide_db::{famous_defs::FamousDefs, syntax_helpers::suggest_name};
+use syntax::{
+    ast::{self, edit::IndentLevel, make, syntax_factory::SyntaxFactory, HasLoopBody},
+    syntax_editor::Position,
+    AstNode,
+};
+
+use crate::{AssistContext, AssistId, AssistKind, Assists};
+
+// Assist: convert_for_loop_to_while_let
+//
+// Converts a for loop into a while let on the Iterator.
+//
+// ```
+// fn main() {
+//     let x = vec![1, 2, 3];
+//     for$0 v in x {
+//         let y = v * 2;
+//     };
+// }
+// ```
+// ->
+// ```
+// fn main() {
+//     let x = vec![1, 2, 3];
+//     let mut tmp = x.into_iter();
+//     while let Some(v) = tmp.next() {
+//         let y = v * 2;
+//     };
+// }
+// ```
+pub(crate) fn convert_for_loop_to_while_let(
+    acc: &mut Assists,
+    ctx: &AssistContext<'_>,
+) -> Option<()> {
+    let for_loop = ctx.find_node_at_offset::<ast::ForExpr>()?;
+    let iterable = for_loop.iterable()?;
+    let pat = for_loop.pat()?;
+    let body = for_loop.loop_body()?;
+    if body.syntax().text_range().start() < ctx.offset() {
+        cov_mark::hit!(not_available_in_body);
+        return None;
+    }
+
+    acc.add(
+        AssistId("convert_for_loop_to_while_let", AssistKind::RefactorRewrite),
+        "Replace this for loop with `while let`",
+        for_loop.syntax().text_range(),
+        |builder| {
+            let make = SyntaxFactory::new();
+            let mut editor = builder.make_editor(for_loop.syntax());
+
+            let (iterable, method) = if impls_core_iter(&ctx.sema, &iterable) {
+                (iterable, None)
+            } else if let Some((expr, method)) = is_ref_and_impls_iter_method(&ctx.sema, &iterable)
+            {
+                (expr, Some(make.name_ref(method.as_str())))
+            } else if let ast::Expr::RefExpr(_) = iterable {
+                (make::expr_paren(iterable), Some(make.name_ref("into_iter")))
+            } else {
+                (iterable, Some(make.name_ref("into_iter")))
+            };
+
+            let iterable = if let Some(method) = method {
+                make::expr_method_call(iterable, method, make::arg_list([]))
+            } else {
+                iterable
+            };
+
+            let mut new_name = suggest_name::NameGenerator::new_from_scope_locals(
+                ctx.sema.scope(for_loop.syntax()),
+            );
+            let tmp_var = new_name.suggest_name("tmp");
+
+            let mut_expr = make.let_stmt(
+                make.ident_pat(false, true, make.name(&tmp_var)).into(),
+                None,
+                Some(iterable),
+            );
+            let indent = IndentLevel::from_node(for_loop.syntax());
+            editor.insert(
+                Position::before(for_loop.syntax()),
+                make::tokens::whitespace(format!("\n{indent}").as_str()),
+            );
+            editor.insert(Position::before(for_loop.syntax()), mut_expr.syntax());
+
+            let opt_pat = make.tuple_struct_pat(make::ext::ident_path("Some"), [pat]);
+            let iter_next_expr = make.expr_method_call(
+                make.expr_path(make::ext::ident_path(&tmp_var)),
+                make.name_ref("next"),
+                make.arg_list([]),
+            );
+            let cond = make.expr_let(opt_pat.into(), iter_next_expr.into());
+
+            let while_loop = make.expr_while_loop(cond.into(), body);
+
+            editor.replace(for_loop.syntax(), while_loop.syntax());
+
+            editor.add_mappings(make.finish_with_mappings());
+            builder.add_file_edits(ctx.file_id(), editor);
+        },
+    )
+}
+
+/// If iterable is a reference where the expression behind the reference implements a method
+/// returning an Iterator called iter or iter_mut (depending on the type of reference) then return
+/// the expression behind the reference and the method name
+fn is_ref_and_impls_iter_method(
+    sema: &hir::Semantics<'_, ide_db::RootDatabase>,
+    iterable: &ast::Expr,
+) -> Option<(ast::Expr, hir::Name)> {
+    let ref_expr = match iterable {
+        ast::Expr::RefExpr(r) => r,
+        _ => return None,
+    };
+    let wanted_method = Name::new_symbol_root(if ref_expr.mut_token().is_some() {
+        sym::iter_mut.clone()
+    } else {
+        sym::iter.clone()
+    });
+    let expr_behind_ref = ref_expr.expr()?;
+    let ty = sema.type_of_expr(&expr_behind_ref)?.adjusted();
+    let scope = sema.scope(iterable.syntax())?;
+    let krate = scope.krate();
+    let iter_trait = FamousDefs(sema, krate).core_iter_Iterator()?;
+
+    let has_wanted_method = ty
+        .iterate_method_candidates(sema.db, &scope, None, Some(&wanted_method), |func| {
+            if func.ret_type(sema.db).impls_trait(sema.db, iter_trait, &[]) {
+                return Some(());
+            }
+            None
+        })
+        .is_some();
+    if !has_wanted_method {
+        return None;
+    }
+
+    Some((expr_behind_ref, wanted_method))
+}
+
+/// Whether iterable implements core::Iterator
+fn impls_core_iter(sema: &hir::Semantics<'_, ide_db::RootDatabase>, iterable: &ast::Expr) -> bool {
+    (|| {
+        let it_typ = sema.type_of_expr(iterable)?.adjusted();
+
+        let module = sema.scope(iterable.syntax())?.module();
+
+        let krate = module.krate();
+        let iter_trait = FamousDefs(sema, krate).core_iter_Iterator()?;
+        cov_mark::hit!(test_already_impls_iterator);
+        Some(it_typ.impls_trait(sema.db, iter_trait, &[]))
+    })()
+    .unwrap_or(false)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    #[test]
+    fn each_to_for_simple_for() {
+        check_assist(
+            convert_for_loop_to_while_let,
+            r"
+fn main() {
+    let mut x = vec![1, 2, 3];
+    for $0v in x {
+        v *= 2;
+    };
+}",
+            r"
+fn main() {
+    let mut x = vec![1, 2, 3];
+    let mut tmp = x.into_iter();
+    while let Some(v) = tmp.next() {
+        v *= 2;
+    };
+}",
+        )
+    }
+
+    #[test]
+    fn each_to_for_for_in_range() {
+        check_assist(
+            convert_for_loop_to_while_let,
+            r#"
+//- minicore: range, iterators
+impl<T> core::iter::Iterator for core::ops::Range<T> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        None
+    }
+}
+
+fn main() {
+    for $0x in 0..92 {
+        print!("{}", x);
+    }
+}"#,
+            r#"
+impl<T> core::iter::Iterator for core::ops::Range<T> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        None
+    }
+}
+
+fn main() {
+    let mut tmp = 0..92;
+    while let Some(x) = tmp.next() {
+        print!("{}", x);
+    }
+}"#,
+        )
+    }
+
+    #[test]
+    fn each_to_for_not_available_in_body() {
+        cov_mark::check!(not_available_in_body);
+        check_assist_not_applicable(
+            convert_for_loop_to_while_let,
+            r"
+fn main() {
+    let mut x = vec![1, 2, 3];
+    for v in x {
+        $0v *= 2;
+    }
+}",
+        )
+    }
+
+    #[test]
+    fn each_to_for_for_borrowed() {
+        check_assist(
+            convert_for_loop_to_while_let,
+            r#"
+//- minicore: iterators
+use core::iter::{Repeat, repeat};
+
+struct S;
+impl S {
+    fn iter(&self) -> Repeat<i32> { repeat(92) }
+    fn iter_mut(&mut self) -> Repeat<i32> { repeat(92) }
+}
+
+fn main() {
+    let x = S;
+    for $0v in &x {
+        let a = v * 2;
+    }
+}
+"#,
+            r#"
+use core::iter::{Repeat, repeat};
+
+struct S;
+impl S {
+    fn iter(&self) -> Repeat<i32> { repeat(92) }
+    fn iter_mut(&mut self) -> Repeat<i32> { repeat(92) }
+}
+
+fn main() {
+    let x = S;
+    let mut tmp = x.iter();
+    while let Some(v) = tmp.next() {
+        let a = v * 2;
+    }
+}
+"#,
+        )
+    }
+
+    #[test]
+    fn each_to_for_for_borrowed_no_iter_method() {
+        check_assist(
+            convert_for_loop_to_while_let,
+            r"
+struct NoIterMethod;
+fn main() {
+    let x = NoIterMethod;
+    for $0v in &x {
+        let a = v * 2;
+    }
+}
+",
+            r"
+struct NoIterMethod;
+fn main() {
+    let x = NoIterMethod;
+    let mut tmp = (&x).into_iter();
+    while let Some(v) = tmp.next() {
+        let a = v * 2;
+    }
+}
+",
+        )
+    }
+
+    #[test]
+    fn each_to_for_for_borrowed_no_iter_method_mut() {
+        check_assist(
+            convert_for_loop_to_while_let,
+            r"
+struct NoIterMethod;
+fn main() {
+    let x = NoIterMethod;
+    for $0v in &mut x {
+        let a = v * 2;
+    }
+}
+",
+            r"
+struct NoIterMethod;
+fn main() {
+    let x = NoIterMethod;
+    let mut tmp = (&mut x).into_iter();
+    while let Some(v) = tmp.next() {
+        let a = v * 2;
+    }
+}
+",
+        )
+    }
+
+    #[test]
+    fn each_to_for_for_borrowed_mut() {
+        check_assist(
+            convert_for_loop_to_while_let,
+            r#"
+//- minicore: iterators
+use core::iter::{Repeat, repeat};
+
+struct S;
+impl S {
+    fn iter(&self) -> Repeat<i32> { repeat(92) }
+    fn iter_mut(&mut self) -> Repeat<i32> { repeat(92) }
+}
+
+fn main() {
+    let x = S;
+    for $0v in &mut x {
+        let a = v * 2;
+    }
+}
+"#,
+            r#"
+use core::iter::{Repeat, repeat};
+
+struct S;
+impl S {
+    fn iter(&self) -> Repeat<i32> { repeat(92) }
+    fn iter_mut(&mut self) -> Repeat<i32> { repeat(92) }
+}
+
+fn main() {
+    let x = S;
+    let mut tmp = x.iter_mut();
+    while let Some(v) = tmp.next() {
+        let a = v * 2;
+    }
+}
+"#,
+        )
+    }
+
+    #[test]
+    fn each_to_for_for_borrowed_mut_behind_var() {
+        check_assist(
+            convert_for_loop_to_while_let,
+            r"
+fn main() {
+    let mut x = vec![1, 2, 3];
+    let y = &mut x;
+    for $0v in y {
+        *v *= 2;
+    }
+}",
+            r"
+fn main() {
+    let mut x = vec![1, 2, 3];
+    let y = &mut x;
+    let mut tmp = y.into_iter();
+    while let Some(v) = tmp.next() {
+        *v *= 2;
+    }
+}",
+        )
+    }
+
+    #[test]
+    fn each_to_for_already_impls_iterator() {
+        cov_mark::check!(test_already_impls_iterator);
+        check_assist(
+            convert_for_loop_to_while_let,
+            r#"
+//- minicore: iterators
+fn main() {
+    for$0 a in core::iter::repeat(92).take(1) {
+        println!("{}", a);
+    }
+}
+"#,
+            r#"
+fn main() {
+    let mut tmp = core::iter::repeat(92).take(1);
+    while let Some(a) = tmp.next() {
+        println!("{}", a);
+    }
+}
+"#,
+        );
+    }
+}
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs b/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
index e8480b0de19..7e9d5966148 100644
--- a/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
@@ -122,6 +122,7 @@ mod handlers {
     mod convert_closure_to_fn;
     mod convert_comment_block;
     mod convert_comment_from_or_to_doc;
+    mod convert_for_to_while_let;
     mod convert_from_to_tryfrom;
     mod convert_integer_literal;
     mod convert_into_to_from;
@@ -252,6 +253,7 @@ mod handlers {
             convert_closure_to_fn::convert_closure_to_fn,
             convert_comment_block::convert_comment_block,
             convert_comment_from_or_to_doc::convert_comment_from_or_to_doc,
+            convert_for_to_while_let::convert_for_loop_to_while_let,
             convert_from_to_tryfrom::convert_from_to_tryfrom,
             convert_integer_literal::convert_integer_literal,
             convert_into_to_from::convert_into_to_from,
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs b/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs
index 4234124d670..359de438ed6 100644
--- a/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs
@@ -440,6 +440,30 @@ fn main() {
 }
 
 #[test]
+fn doctest_convert_for_loop_to_while_let() {
+    check_doc_test(
+        "convert_for_loop_to_while_let",
+        r#####"
+fn main() {
+    let x = vec![1, 2, 3];
+    for$0 v in x {
+        let y = v * 2;
+    };
+}
+"#####,
+        r#####"
+fn main() {
+    let x = vec![1, 2, 3];
+    let mut tmp = x.into_iter();
+    while let Some(v) = tmp.next() {
+        let y = v * 2;
+    };
+}
+"#####,
+    )
+}
+
+#[test]
 fn doctest_convert_for_loop_with_for_each() {
     check_doc_test(
         "convert_for_loop_with_for_each",
diff --git a/src/tools/rust-analyzer/crates/syntax/src/ast/make.rs b/src/tools/rust-analyzer/crates/syntax/src/ast/make.rs
index 231c21c38f8..234e3d1efe1 100644
--- a/src/tools/rust-analyzer/crates/syntax/src/ast/make.rs
+++ b/src/tools/rust-analyzer/crates/syntax/src/ast/make.rs
@@ -623,6 +623,10 @@ pub fn expr_for_loop(pat: ast::Pat, expr: ast::Expr, block: ast::BlockExpr) -> a
     expr_from_text(&format!("for {pat} in {expr} {block}"))
 }
 
+pub fn expr_while_loop(condition: ast::Expr, block: ast::BlockExpr) -> ast::WhileExpr {
+    expr_from_text(&format!("while {condition} {block}"))
+}
+
 pub fn expr_loop(block: ast::BlockExpr) -> ast::Expr {
     expr_from_text(&format!("loop {block}"))
 }
diff --git a/src/tools/rust-analyzer/crates/syntax/src/ast/syntax_factory/constructors.rs b/src/tools/rust-analyzer/crates/syntax/src/ast/syntax_factory/constructors.rs
index 44f13041c24..bec5d2a39a3 100644
--- a/src/tools/rust-analyzer/crates/syntax/src/ast/syntax_factory/constructors.rs
+++ b/src/tools/rust-analyzer/crates/syntax/src/ast/syntax_factory/constructors.rs
@@ -1,8 +1,8 @@
 //! Wrappers over [`make`] constructors
 use crate::{
     ast::{
-        self, make, HasArgList, HasGenericArgs, HasGenericParams, HasName, HasTypeBounds,
-        HasVisibility,
+        self, make, HasArgList, HasGenericArgs, HasGenericParams, HasLoopBody, HasName,
+        HasTypeBounds, HasVisibility,
     },
     syntax_editor::SyntaxMappingBuilder,
     AstNode, NodeOrToken, SyntaxKind, SyntaxNode, SyntaxToken,
@@ -543,6 +543,19 @@ impl SyntaxFactory {
         ast
     }
 
+    pub fn expr_while_loop(&self, condition: ast::Expr, body: ast::BlockExpr) -> ast::WhileExpr {
+        let ast = make::expr_while_loop(condition.clone(), body.clone()).clone_for_update();
+
+        if let Some(mut mapping) = self.mappings() {
+            let mut builder = SyntaxMappingBuilder::new(ast.syntax().clone());
+            builder.map_node(condition.syntax().clone(), ast.condition().unwrap().syntax().clone());
+            builder.map_node(body.syntax().clone(), ast.loop_body().unwrap().syntax().clone());
+            builder.finish(&mut mapping);
+        }
+
+        ast
+    }
+
     pub fn expr_let(&self, pattern: ast::Pat, expr: ast::Expr) -> ast::LetExpr {
         let ast = make::expr_let(pattern.clone(), expr.clone()).clone_for_update();