From c84beefd83df7bb39a02a341f5615019e1ffaa63 Mon Sep 17 00:00:00 2001
From: hyd-dev <yd-huang@outlook.com>
Date: Sun, 20 Jun 2021 21:49:08 +0800
Subject: [PATCH] Add associated functions that have custom linkage to
 `reachable_set`

---
 compiler/rustc_passes/src/reachable.rs        | 34 ++++++++++++-------
 .../ui/auxiliary/no-mangle-associated-fn.rs   | 10 ++++++
 src/test/ui/no-mangle-associated-fn.rs        | 22 ++++++++++++
 3 files changed, 53 insertions(+), 13 deletions(-)
 create mode 100644 src/test/ui/auxiliary/no-mangle-associated-fn.rs
 create mode 100644 src/test/ui/no-mangle-associated-fn.rs

diff --git a/compiler/rustc_passes/src/reachable.rs b/compiler/rustc_passes/src/reachable.rs
index ad02f4f8269..963153a01a0 100644
--- a/compiler/rustc_passes/src/reachable.rs
+++ b/compiler/rustc_passes/src/reachable.rs
@@ -211,13 +211,15 @@ fn propagate_node(&mut self, node: &Node<'tcx>, search_item: LocalDefId) {
         if !self.any_library {
             // If we are building an executable, only explicitly extern
             // types need to be exported.
-            if let Node::Item(item) = *node {
-                let reachable = if let hir::ItemKind::Fn(ref sig, ..) = item.kind {
-                    sig.header.abi != Abi::Rust
-                } else {
-                    false
-                };
-                let codegen_attrs = self.tcx.codegen_fn_attrs(item.def_id);
+            if let Node::Item(hir::Item { kind: hir::ItemKind::Fn(sig, ..), def_id, .. })
+            | Node::ImplItem(hir::ImplItem {
+                kind: hir::ImplItemKind::Fn(sig, ..),
+                def_id,
+                ..
+            }) = *node
+            {
+                let reachable = sig.header.abi != Abi::Rust;
+                let codegen_attrs = self.tcx.codegen_fn_attrs(*def_id);
                 let is_extern = codegen_attrs.contains_extern_indicator();
                 let std_internal =
                     codegen_attrs.flags.contains(CodegenFnAttrFlags::RUSTC_STD_INTERNAL_SYMBOL);
@@ -335,17 +337,23 @@ struct CollectPrivateImplItemsVisitor<'a, 'tcx> {
     worklist: &'a mut Vec<LocalDefId>,
 }
 
-impl<'a, 'tcx> ItemLikeVisitor<'tcx> for CollectPrivateImplItemsVisitor<'a, 'tcx> {
-    fn visit_item(&mut self, item: &hir::Item<'_>) {
+impl CollectPrivateImplItemsVisitor<'_, '_> {
+    fn push_to_worklist_if_has_custom_linkage(&mut self, def_id: LocalDefId) {
         // Anything which has custom linkage gets thrown on the worklist no
         // matter where it is in the crate, along with "special std symbols"
         // which are currently akin to allocator symbols.
-        let codegen_attrs = self.tcx.codegen_fn_attrs(item.def_id);
+        let codegen_attrs = self.tcx.codegen_fn_attrs(def_id);
         if codegen_attrs.contains_extern_indicator()
             || codegen_attrs.flags.contains(CodegenFnAttrFlags::RUSTC_STD_INTERNAL_SYMBOL)
         {
-            self.worklist.push(item.def_id);
+            self.worklist.push(def_id);
         }
+    }
+}
+
+impl<'a, 'tcx> ItemLikeVisitor<'tcx> for CollectPrivateImplItemsVisitor<'a, 'tcx> {
+    fn visit_item(&mut self, item: &hir::Item<'_>) {
+        self.push_to_worklist_if_has_custom_linkage(item.def_id);
 
         // We need only trait impls here, not inherent impls, and only non-exported ones
         if let hir::ItemKind::Impl(hir::Impl { of_trait: Some(ref trait_ref), ref items, .. }) =
@@ -375,8 +383,8 @@ fn visit_item(&mut self, item: &hir::Item<'_>) {
 
     fn visit_trait_item(&mut self, _trait_item: &hir::TraitItem<'_>) {}
 
-    fn visit_impl_item(&mut self, _impl_item: &hir::ImplItem<'_>) {
-        // processed in visit_item above
+    fn visit_impl_item(&mut self, impl_item: &hir::ImplItem<'_>) {
+        self.push_to_worklist_if_has_custom_linkage(impl_item.def_id);
     }
 
     fn visit_foreign_item(&mut self, _foreign_item: &hir::ForeignItem<'_>) {
diff --git a/src/test/ui/auxiliary/no-mangle-associated-fn.rs b/src/test/ui/auxiliary/no-mangle-associated-fn.rs
new file mode 100644
index 00000000000..2ef161edf29
--- /dev/null
+++ b/src/test/ui/auxiliary/no-mangle-associated-fn.rs
@@ -0,0 +1,10 @@
+#![crate_type = "lib"]
+
+struct Bar;
+
+impl Bar {
+    #[no_mangle]
+    fn bar() -> u8 {
+        2
+    }
+}
diff --git a/src/test/ui/no-mangle-associated-fn.rs b/src/test/ui/no-mangle-associated-fn.rs
new file mode 100644
index 00000000000..e9c621914b2
--- /dev/null
+++ b/src/test/ui/no-mangle-associated-fn.rs
@@ -0,0 +1,22 @@
+// aux-build: no-mangle-associated-fn.rs
+// run-pass
+
+extern crate no_mangle_associated_fn;
+
+struct Foo;
+
+impl Foo {
+    #[no_mangle]
+    fn foo() -> u8 {
+        1
+    }
+}
+
+fn main() {
+    extern "Rust" {
+        fn foo() -> u8;
+        fn bar() -> u8;
+    }
+    assert_eq!(unsafe { foo() }, 1);
+    assert_eq!(unsafe { bar() }, 2);
+}
-- 
2.44.0