Merge pull request #20381 from A4-Tacks/fix-assign-sug

Add assignment type analysis for ide-completion

2 files changed, 82 insertions, 0 deletions

diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/context/analysis.rs b/src/tools/rust-analyzer/crates/ide-completion/src/context/analysis.rs
index 2eabf99fc69..9d246017131 100644
--- a/src/tools/rust-analyzer/crates/ide-completion/src/context/analysis.rs
+++ b/src/tools/rust-analyzer/crates/ide-completion/src/context/analysis.rs
@@ -559,6 +559,7 @@ fn expected_type_and_name<'db>(
     token: &SyntaxToken,
     name_like: &ast::NameLike,
 ) -> (Option<Type<'db>>, Option<NameOrNameRef>) {
+    let token = prev_assign_token_at_trivia(token.clone());
     let mut node = match token.parent() {
         Some(it) => it,
         None => return (None, None),
@@ -629,6 +630,17 @@ fn expected_type_and_name<'db>(
                         .map(TypeInfo::original);
                     (ty, None)
                 },
+                ast::BinExpr(it) => {
+                    if let Some(ast::BinaryOp::Assignment { op: None }) = it.op_kind() {
+                        let ty = it.lhs()
+                            .and_then(|lhs| sema.type_of_expr(&lhs))
+                            .or_else(|| it.rhs().and_then(|rhs| sema.type_of_expr(&rhs)))
+                            .map(TypeInfo::original);
+                        (ty, None)
+                    } else {
+                        (None, None)
+                    }
+                },
                 ast::ArgList(_) => {
                     cov_mark::hit!(expected_type_fn_param);
                     ActiveParameter::at_token(
@@ -1856,3 +1868,23 @@ fn next_non_trivia_sibling(ele: SyntaxElement) -> Option<SyntaxElement> {
     }
     None
 }
+
+fn prev_assign_token_at_trivia(mut token: SyntaxToken) -> SyntaxToken {
+    while token.kind().is_trivia()
+        && let Some(prev) = token.prev_token()
+        && let T![=]
+        | T![+=]
+        | T![/=]
+        | T![*=]
+        | T![%=]
+        | T![>>=]
+        | T![<<=]
+        | T![-=]
+        | T![|=]
+        | T![&=]
+        | T![^=] = prev.kind()
+    {
+        token = prev
+    }
+    token
+}
diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/context/tests.rs b/src/tools/rust-analyzer/crates/ide-completion/src/context/tests.rs
index 75c20968e1e..7a8c70f190e 100644
--- a/src/tools/rust-analyzer/crates/ide-completion/src/context/tests.rs
+++ b/src/tools/rust-analyzer/crates/ide-completion/src/context/tests.rs
@@ -434,3 +434,53 @@ fn f(thing: u32) -> &u32 {
         expect!["ty: u32, name: ?"],
     );
 }
+
+#[test]
+fn expected_type_assign() {
+    check_expected_type_and_name(
+        r#"
+enum State { Stop }
+fn foo() {
+    let x: &mut State = &mut State::Stop;
+    x = $0;
+}
+"#,
+        expect![[r#"ty: &'_ mut State, name: ?"#]],
+    );
+}
+
+#[test]
+fn expected_type_deref_assign() {
+    check_expected_type_and_name(
+        r#"
+enum State { Stop }
+fn foo() {
+    let x: &mut State = &mut State::Stop;
+    match x {
+        State::Stop => {
+            *x = $0;
+        },
+    }
+}
+"#,
+        expect![[r#"ty: State, name: ?"#]],
+    );
+}
+
+#[test]
+fn expected_type_deref_assign_at_block_end() {
+    check_expected_type_and_name(
+        r#"
+enum State { Stop }
+fn foo() {
+    let x: &mut State = &mut State::Stop;
+    match x {
+        State::Stop => {
+            *x = $0
+        },
+    }
+}
+"#,
+        expect![[r#"ty: State, name: ?"#]],
+    );
+}