From 18feb726be47e534e8e903cfd9080592e055253b Mon Sep 17 00:00:00 2001
From: Yunfei <i.heyunfei@gmail.com>
Date: Mon, 29 Jul 2024 22:42:31 +0800
Subject: [PATCH] feat(ide-completion): extra sugar auto-completion `async fn
 ...` in `impl trait` for `async fn in trait` that's defined in desugar form

---
 crates/hir/src/lib.rs                         | 47 +++++++++++++++++++
 .../src/completions/item_list/trait_impl.rs   | 38 ++++++++++++++-
 2 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 96a6e6f1f1..74ec5227b0 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -2207,6 +2207,53 @@ impl Function {
         db.function_data(self.id).is_async()
     }
 
+    /// Whether this function is a `fn` that returns `impl Future`.
+    pub fn is_desugar_async(self, db: &dyn HirDatabase) -> bool {
+        if self.is_async(db) || self.is_const(db) {
+            return false;
+        }
+
+        let Some(impl_traits) = self.ret_type(db).as_impl_traits(db) else { return false };
+
+        let Some(future_trait_id) =
+            db.lang_item(self.ty(db).env.krate, LangItem::Future).and_then(|t| t.as_trait())
+        else {
+            return false;
+        };
+
+        let Some(size_trait_id) =
+            db.lang_item(self.ty(db).env.krate, LangItem::Sized).and_then(|t| t.as_trait())
+        else {
+            return false;
+        };
+
+        let Some(sync_trait_id) =
+            db.lang_item(self.ty(db).env.krate, LangItem::Sync).and_then(|t| t.as_trait())
+        else {
+            return false;
+        };
+
+        // TODO: There's no `LangItem::Send`. How do we get the id of `Send` trait?
+        // let Some(send_trait_id) = db.lang_item(self.ty(db).env.krate, LangItem::Send).and_then(|t| t.as_trait()) else {
+        //     eprint!("no future_trait_id\n");
+        //     return false
+        // };
+
+        let allowed_to_leaked_types = vec![size_trait_id, sync_trait_id];
+
+        let mut has_impl_future = false;
+        let mut has_types_not_allow_to_leaked = false;
+        for impl_trait in impl_traits {
+            if impl_trait.id == future_trait_id {
+                has_impl_future = true;
+            } else if !allowed_to_leaked_types.contains(&impl_trait.id) {
+                has_types_not_allow_to_leaked = true;
+            }
+        }
+
+        has_impl_future && !has_types_not_allow_to_leaked
+    }
+
     /// Does this function have `#[test]` attribute?
     pub fn is_test(self, db: &dyn HirDatabase) -> bool {
         db.function_data(self.id).attrs.is_test()
diff --git a/crates/ide-completion/src/completions/item_list/trait_impl.rs b/crates/ide-completion/src/completions/item_list/trait_impl.rs
index fc6e1ebf05..7887a87448 100644
--- a/crates/ide-completion/src/completions/item_list/trait_impl.rs
+++ b/crates/ide-completion/src/completions/item_list/trait_impl.rs
@@ -210,7 +210,7 @@ fn add_function_impl(
                 ast::AssocItem::Fn(func) => func,
                 _ => unreachable!(),
             };
-
+            // TODO: need `function_decl` that unwraps future in the return type
             let function_decl = function_declaration(&transformed_fn, source.file_id.is_macro());
             match ctx.config.snippet_cap {
                 Some(cap) => {
@@ -225,6 +225,42 @@ fn add_function_impl(
             item.add_to(acc, ctx.db);
         }
     }
+
+    eprint!("is_desugar_async: {}", func.is_desugar_async(ctx.db));
+    if func.is_desugar_async(ctx.db) {
+        let label = format_smolstr!(
+            "async fn {}({})",
+            fn_name.display(ctx.db),
+            if func.assoc_fn_params(ctx.db).is_empty() { "" } else { ".." }
+        );
+        let mut item = CompletionItem::new(completion_kind, replacement_range, label);
+        item.lookup_by(format!("async fn {}", fn_name.display(ctx.db)))
+            .set_documentation(func.docs(ctx.db))
+            .set_relevance(CompletionRelevance { is_item_from_trait: true, ..Default::default() });
+        if let Some(source) = ctx.sema.source(func) {
+            let assoc_item = ast::AssocItem::Fn(source.value);
+            if let Some(transformed_item) = get_transformed_assoc_item(ctx, assoc_item, impl_def) {
+                let transformed_fn = match transformed_item {
+                    ast::AssocItem::Fn(func) => func,
+                    _ => unreachable!(),
+                };
+
+                let function_decl =
+                    function_declaration(&transformed_fn, source.file_id.is_macro());
+                match ctx.config.snippet_cap {
+                    Some(cap) => {
+                        let snippet = format!("{function_decl} {{\n    $0\n}}");
+                        item.snippet_edit(cap, TextEdit::replace(replacement_range, snippet));
+                    }
+                    None => {
+                        let header = format!("{function_decl} {{");
+                        item.text_edit(TextEdit::replace(replacement_range, header));
+                    }
+                };
+                item.add_to(acc, ctx.db);
+            }
+        }
+    }
 }
 
 /// Transform a relevant associated item to inline generics from the impl, remove attrs and docs, etc.