Extract logic to decide how to complete semicolon for unit-returning function into `CompletionContext`

So that we don't recompute it for every item.

2 files changed, 73 insertions, 39 deletions

diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/context.rs b/src/tools/rust-analyzer/crates/ide-completion/src/context.rs
index d457ba32bf0..72a57915169 100644
--- a/src/tools/rust-analyzer/crates/ide-completion/src/context.rs
+++ b/src/tools/rust-analyzer/crates/ide-completion/src/context.rs
@@ -4,7 +4,7 @@ mod analysis;
 #[cfg(test)]
 mod tests;
 
-use std::iter;
+use std::{iter, ops::ControlFlow};
 
 use hir::{
     HasAttrs, Local, Name, PathResolution, ScopeDef, Semantics, SemanticsScope, Type, TypeInfo,
@@ -15,7 +15,7 @@ use ide_db::{
 };
 use syntax::{
     ast::{self, AttrKind, NameOrNameRef},
-    AstNode, Edition, SmolStr,
+    match_ast, AstNode, Edition, SmolStr,
     SyntaxKind::{self, *},
     SyntaxToken, TextRange, TextSize, T,
 };
@@ -457,6 +457,16 @@ pub(crate) struct CompletionContext<'a> {
     ///
     /// Here depth will be 2
     pub(crate) depth_from_crate_root: usize,
+
+    /// Whether and how to complete semicolon for unit-returning functions.
+    pub(crate) complete_semicolon: CompleteSemicolon,
+}
+
+#[derive(Debug)]
+pub(crate) enum CompleteSemicolon {
+    DoNotComplete,
+    CompleteSemi,
+    CompleteComma,
 }
 
 impl CompletionContext<'_> {
@@ -735,6 +745,53 @@ impl<'a> CompletionContext<'a> {
 
         let depth_from_crate_root = iter::successors(module.parent(db), |m| m.parent(db)).count();
 
+        let complete_semicolon = if config.add_semicolon_to_unit {
+            let inside_closure_ret = token.parent_ancestors().try_for_each(|ancestor| {
+                match_ast! {
+                    match ancestor {
+                        ast::BlockExpr(_) => ControlFlow::Break(false),
+                        ast::ClosureExpr(_) => ControlFlow::Break(true),
+                        _ => ControlFlow::Continue(())
+                    }
+                }
+            });
+
+            if inside_closure_ret == ControlFlow::Break(true) {
+                CompleteSemicolon::DoNotComplete
+            } else {
+                let next_non_trivia_token =
+                    std::iter::successors(token.next_token(), |it| it.next_token())
+                        .find(|it| !it.kind().is_trivia());
+                let in_match_arm = token.parent_ancestors().try_for_each(|ancestor| {
+                    if ast::MatchArm::can_cast(ancestor.kind()) {
+                        ControlFlow::Break(true)
+                    } else if matches!(
+                        ancestor.kind(),
+                        SyntaxKind::EXPR_STMT | SyntaxKind::BLOCK_EXPR
+                    ) {
+                        ControlFlow::Break(false)
+                    } else {
+                        ControlFlow::Continue(())
+                    }
+                });
+                // FIXME: This will assume expr macros are not inside match, we need to somehow go to the "parent" of the root node.
+                let in_match_arm = match in_match_arm {
+                    ControlFlow::Continue(()) => false,
+                    ControlFlow::Break(it) => it,
+                };
+                let complete_token = if in_match_arm { T![,] } else { T![;] };
+                if next_non_trivia_token.map(|it| it.kind()) == Some(complete_token) {
+                    CompleteSemicolon::DoNotComplete
+                } else if in_match_arm {
+                    CompleteSemicolon::CompleteComma
+                } else {
+                    CompleteSemicolon::CompleteSemi
+                }
+            }
+        } else {
+            CompleteSemicolon::DoNotComplete
+        };
+
         let ctx = CompletionContext {
             sema,
             scope,
@@ -752,6 +809,7 @@ impl<'a> CompletionContext<'a> {
             qualifier_ctx,
             locals,
             depth_from_crate_root,
+            complete_semicolon,
         };
         Some((ctx, analysis))
     }
diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/render/function.rs b/src/tools/rust-analyzer/crates/ide-completion/src/render/function.rs
index 332a4ebd8e3..a859d79e243 100644
--- a/src/tools/rust-analyzer/crates/ide-completion/src/render/function.rs
+++ b/src/tools/rust-analyzer/crates/ide-completion/src/render/function.rs
@@ -1,15 +1,15 @@
 //! Renderer for function calls.
 
-use std::ops::ControlFlow;
-
 use hir::{db::HirDatabase, AsAssocItem, HirDisplay};
 use ide_db::{SnippetCap, SymbolKind};
 use itertools::Itertools;
 use stdx::{format_to, to_lower_snake_case};
-use syntax::{ast, format_smolstr, match_ast, AstNode, Edition, SmolStr, SyntaxKind, ToSmolStr, T};
+use syntax::{format_smolstr, AstNode, Edition, SmolStr, ToSmolStr};
 
 use crate::{
-    context::{CompletionContext, DotAccess, DotAccessKind, PathCompletionCtx, PathKind},
+    context::{
+        CompleteSemicolon, CompletionContext, DotAccess, DotAccessKind, PathCompletionCtx, PathKind,
+    },
     item::{
         Builder, CompletionItem, CompletionItemKind, CompletionRelevance, CompletionRelevanceFn,
         CompletionRelevanceReturnType, CompletionRelevanceTraitInfo,
@@ -277,40 +277,16 @@ pub(super) fn add_call_parens<'b>(
 
         (snippet, "(…)")
     };
-    if ret_type.is_unit() && ctx.config.add_semicolon_to_unit {
-        let inside_closure_ret = ctx.token.parent_ancestors().try_for_each(|ancestor| {
-            match_ast! {
-                match ancestor {
-                    ast::BlockExpr(_) => ControlFlow::Break(false),
-                    ast::ClosureExpr(_) => ControlFlow::Break(true),
-                    _ => ControlFlow::Continue(())
-                }
-            }
-        });
-
-        if inside_closure_ret != ControlFlow::Break(true) {
-            let next_non_trivia_token =
-                std::iter::successors(ctx.token.next_token(), |it| it.next_token())
-                    .find(|it| !it.kind().is_trivia());
-            let in_match_arm = ctx.token.parent_ancestors().try_for_each(|ancestor| {
-                if ast::MatchArm::can_cast(ancestor.kind()) {
-                    ControlFlow::Break(true)
-                } else if matches!(ancestor.kind(), SyntaxKind::EXPR_STMT | SyntaxKind::BLOCK_EXPR)
-                {
-                    ControlFlow::Break(false)
-                } else {
-                    ControlFlow::Continue(())
-                }
-            });
-            // FIXME: This will assume expr macros are not inside match, we need to somehow go to the "parent" of the root node.
-            let in_match_arm = match in_match_arm {
-                ControlFlow::Continue(()) => false,
-                ControlFlow::Break(it) => it,
-            };
-            let complete_token = if in_match_arm { T![,] } else { T![;] };
-            if next_non_trivia_token.map(|it| it.kind()) != Some(complete_token) {
+    if ret_type.is_unit() {
+        match ctx.complete_semicolon {
+            CompleteSemicolon::DoNotComplete => {}
+            CompleteSemicolon::CompleteSemi | CompleteSemicolon::CompleteComma => {
                 cov_mark::hit!(complete_semicolon);
-                let ch = if in_match_arm { ',' } else { ';' };
+                let ch = if matches!(ctx.complete_semicolon, CompleteSemicolon::CompleteComma) {
+                    ','
+                } else {
+                    ';'
+                };
                 if snippet.ends_with("$0") {
                     snippet.insert(snippet.len() - "$0".len(), ch);
                 } else {