modified lint to work with MIR

Working with MIR let's us exclude expressions like `&fn_name as &dyn Something`
and `(&fn_name)()`. Also added ABI, unsafety and whether a function is variadic
in the lint suggestion, included the `&` in the span of the lint and updated the
test.

7 files changed, 322 insertions, 59 deletions

diff --git a/compiler/rustc_lint/src/builtin.rs b/compiler/rustc_lint/src/builtin.rs
index 86cf4ebb187..e5f66611d0f 100644
--- a/compiler/rustc_lint/src/builtin.rs
+++ b/compiler/rustc_lint/src/builtin.rs
@@ -2942,37 +2942,3 @@ impl<'tcx> LateLintPass<'tcx> for ClashingExternDeclarations {
         }
     }
 }
-
-declare_lint! {
-    FUNCTION_REFERENCES,
-    Warn,
-    "suggest casting functions to pointers when attempting to take references"
-}
-
-declare_lint_pass!(FunctionReferences => [FUNCTION_REFERENCES]);
-
-impl<'tcx> LateLintPass<'tcx> for FunctionReferences {
-    fn check_expr(&mut self, cx: &LateContext<'_>, e: &hir::Expr<'_>) {
-        if let hir::ExprKind::AddrOf(hir::BorrowKind::Ref, _, referent) = e.kind {
-            if let hir::ExprKind::Path(qpath) = &referent.kind {
-                if let Some(def_id) = cx.qpath_res(qpath, referent.hir_id).opt_def_id() {
-                    cx.tcx.hir().get_if_local(def_id).map(|node| {
-                        node.fn_decl().map(|decl| {
-                            if let Some(ident) = node.ident() {
-                                cx.struct_span_lint(FUNCTION_REFERENCES, referent.span, |lint| {
-                                    let num_args = decl.inputs.len();
-                                    lint.build(&format!(
-                                        "cast `{}` with `as *const fn({}) -> _` to use it as a pointer",
-                                        ident.to_string(),
-                                        vec!["_"; num_args].join(", ")
-                                    ))
-                                    .emit()
-                                });
-                            }
-                        });
-                    });
-                }
-            }
-        }
-    }
-}
diff --git a/compiler/rustc_lint/src/lib.rs b/compiler/rustc_lint/src/lib.rs
index b574a7da0aa..1db59bfc39d 100644
--- a/compiler/rustc_lint/src/lib.rs
+++ b/compiler/rustc_lint/src/lib.rs
@@ -194,7 +194,6 @@ macro_rules! late_lint_mod_passes {
                 UnreachablePub: UnreachablePub,
                 ExplicitOutlivesRequirements: ExplicitOutlivesRequirements,
                 InvalidValue: InvalidValue,
-                FunctionReferences: FunctionReferences,
             ]
         );
     };
diff --git a/compiler/rustc_mir/src/transform/function_references.rs b/compiler/rustc_mir/src/transform/function_references.rs
new file mode 100644
index 00000000000..2daa468136f
--- /dev/null
+++ b/compiler/rustc_mir/src/transform/function_references.rs
@@ -0,0 +1,158 @@
+use rustc_hir::def_id::DefId;
+use rustc_middle::mir::visit::Visitor;
+use rustc_middle::mir::*;
+use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_session::lint::builtin::FUNCTION_REFERENCES;
+use rustc_span::Span;
+use rustc_target::spec::abi::Abi;
+
+use crate::transform::{MirPass, MirSource};
+
+pub struct FunctionReferences;
+
+impl<'tcx> MirPass<'tcx> for FunctionReferences {
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, _src: MirSource<'tcx>, body: &mut Body<'tcx>) {
+        let source_info = SourceInfo::outermost(body.span);
+        let mut checker = FunctionRefChecker {
+            tcx,
+            body,
+            potential_lints: Vec::new(),
+            casts: Vec::new(),
+            calls: Vec::new(),
+            source_info,
+        };
+        checker.visit_body(&body);
+    }
+}
+
+struct FunctionRefChecker<'a, 'tcx> {
+    tcx: TyCtxt<'tcx>,
+    body: &'a Body<'tcx>,
+    potential_lints: Vec<FunctionRefLint>,
+    casts: Vec<Span>,
+    calls: Vec<Span>,
+    source_info: SourceInfo,
+}
+
+impl<'a, 'tcx> Visitor<'tcx> for FunctionRefChecker<'a, 'tcx> {
+    fn visit_basic_block_data(&mut self, block: BasicBlock, data: &BasicBlockData<'tcx>) {
+        self.super_basic_block_data(block, data);
+        for cast_span in self.casts.drain(..) {
+            self.potential_lints.retain(|lint| lint.source_info.span != cast_span);
+        }
+        for call_span in self.calls.drain(..) {
+            self.potential_lints.retain(|lint| lint.source_info.span != call_span);
+        }
+        for lint in self.potential_lints.drain(..) {
+            lint.emit(self.tcx, self.body);
+        }
+    }
+    fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
+        self.source_info = statement.source_info;
+        self.super_statement(statement, location);
+    }
+    fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
+        self.source_info = terminator.source_info;
+        if let TerminatorKind::Call {
+            func,
+            args: _,
+            destination: _,
+            cleanup: _,
+            from_hir_call: _,
+            fn_span: _,
+        } = &terminator.kind
+        {
+            let span = match func {
+                Operand::Copy(place) | Operand::Move(place) => {
+                    self.body.local_decls[place.local].source_info.span
+                }
+                Operand::Constant(constant) => constant.span,
+            };
+            self.calls.push(span);
+        };
+        self.super_terminator(terminator, location);
+    }
+    fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, location: Location) {
+        match rvalue {
+            Rvalue::Ref(_, _, place) | Rvalue::AddressOf(_, place) => {
+                let decl = &self.body.local_decls[place.local];
+                if let ty::FnDef(def_id, _) = decl.ty.kind {
+                    let ident = self
+                        .body
+                        .var_debug_info
+                        .iter()
+                        .find(|info| info.source_info.span == decl.source_info.span)
+                        .map(|info| info.name.to_ident_string())
+                        .unwrap_or(self.tcx.def_path_str(def_id));
+                    let lint = FunctionRefLint { ident, def_id, source_info: self.source_info };
+                    self.potential_lints.push(lint);
+                }
+            }
+            Rvalue::Cast(_, op, _) => {
+                let op_ty = op.ty(self.body, self.tcx);
+                if self.is_fn_ref(op_ty) {
+                    self.casts.push(self.source_info.span);
+                }
+            }
+            _ => {}
+        }
+        self.super_rvalue(rvalue, location);
+    }
+}
+
+impl<'a, 'tcx> FunctionRefChecker<'a, 'tcx> {
+    fn is_fn_ref(&self, ty: Ty<'tcx>) -> bool {
+        let referent_ty = match ty.kind {
+            ty::Ref(_, referent_ty, _) => Some(referent_ty),
+            ty::RawPtr(ty_and_mut) => Some(ty_and_mut.ty),
+            _ => None,
+        };
+        referent_ty
+            .map(|ref_ty| if let ty::FnDef(..) = ref_ty.kind { true } else { false })
+            .unwrap_or(false)
+    }
+}
+
+struct FunctionRefLint {
+    ident: String,
+    def_id: DefId,
+    source_info: SourceInfo,
+}
+
+impl<'tcx> FunctionRefLint {
+    fn emit(&self, tcx: TyCtxt<'tcx>, body: &Body<'tcx>) {
+        let def_id = self.def_id;
+        let source_info = self.source_info;
+        let lint_root = body.source_scopes[source_info.scope]
+            .local_data
+            .as_ref()
+            .assert_crate_local()
+            .lint_root;
+        let fn_sig = tcx.fn_sig(def_id);
+        let unsafety = fn_sig.unsafety().prefix_str();
+        let abi = match fn_sig.abi() {
+            Abi::Rust => String::from(""),
+            other_abi => {
+                let mut s = String::from("extern \"");
+                s.push_str(other_abi.name());
+                s.push_str("\" ");
+                s
+            }
+        };
+        let num_args = fn_sig.inputs().map_bound(|inputs| inputs.len()).skip_binder();
+        let variadic = if fn_sig.c_variadic() { ", ..." } else { "" };
+        let ret = if fn_sig.output().skip_binder().is_unit() { "" } else { " -> _" };
+        tcx.struct_span_lint_hir(FUNCTION_REFERENCES, lint_root, source_info.span, |lint| {
+            lint.build(&format!(
+                "cast `{}` with `as {}{}fn({}{}){}` to use it as a pointer",
+                self.ident,
+                unsafety,
+                abi,
+                vec!["_"; num_args].join(", "),
+                variadic,
+                ret,
+            ))
+            .emit()
+        });
+    }
+}
diff --git a/compiler/rustc_mir/src/transform/mod.rs b/compiler/rustc_mir/src/transform/mod.rs
index 20b8c90a9dc..3f50420b86b 100644
--- a/compiler/rustc_mir/src/transform/mod.rs
+++ b/compiler/rustc_mir/src/transform/mod.rs
@@ -27,6 +27,7 @@ pub mod dest_prop;
 pub mod dump_mir;
 pub mod early_otherwise_branch;
 pub mod elaborate_drops;
+pub mod function_references;
 pub mod generator;
 pub mod inline;
 pub mod instcombine;
@@ -266,6 +267,7 @@ fn mir_const<'tcx>(
             // MIR-level lints.
             &check_packed_ref::CheckPackedRef,
             &check_const_item_mutation::CheckConstItemMutation,
+            &function_references::FunctionReferences,
             // What we need to do constant evaluation.
             &simplify::SimplifyCfg::new("initial"),
             &rustc_peek::SanityCheck,
diff --git a/compiler/rustc_session/src/lint/builtin.rs b/compiler/rustc_session/src/lint/builtin.rs
index ab56a0a5667..4cfd3277202 100644
--- a/compiler/rustc_session/src/lint/builtin.rs
+++ b/compiler/rustc_session/src/lint/builtin.rs
@@ -2645,6 +2645,11 @@ declare_lint! {
         reference: "issue #76200 <https://github.com/rust-lang/rust/issues/76200>",
         edition: None,
     };
+
+declare_lint! {
+    pub FUNCTION_REFERENCES,
+    Warn,
+    "suggest casting functions to pointers when attempting to take references",
 }
 
 declare_lint! {
@@ -2762,6 +2767,7 @@ declare_lint_pass! {
         CONST_EVALUATABLE_UNCHECKED,
         INEFFECTIVE_UNSTABLE_TRAIT_IMPL,
         UNINHABITED_STATIC,
+        FUNCTION_REFERENCES,
     ]
 }
 
diff --git a/src/test/ui/lint/function-references.rs b/src/test/ui/lint/function-references.rs
index 748f2560caa..ad0380375d6 100644
--- a/src/test/ui/lint/function-references.rs
+++ b/src/test/ui/lint/function-references.rs
@@ -1,22 +1,71 @@
 // check-pass
-fn foo() -> usize { 42 }
-fn bar(x: usize) -> usize { x }
-fn baz(x: usize, y: usize) -> usize { x + y }
+#![feature(c_variadic)]
+#![allow(dead_code)]
+
+fn foo() -> u32 { 42 }
+fn bar(x: u32) -> u32 { x }
+fn baz(x: u32, y: u32) -> u32 { x + y }
+unsafe fn unsafe_fn() { }
+extern "C" fn c_fn() { }
+unsafe extern "C" fn unsafe_c_fn() { }
+unsafe extern fn variadic_fn(_x: u32, _args: ...) { }
+fn call_fn(f: &dyn Fn(u32) -> u32, x: u32) { f(x); }
+fn parameterized_call_fn<F: Fn(u32) -> u32>(f: &F, x: u32) { f(x); }
 
 fn main() {
+    let _zst_ref = &foo;
+    //~^ WARN cast `foo` with `as fn() -> _` to use it as a pointer
+    let fn_item = foo;
+    let _indirect_ref = &fn_item;
+    //~^ WARN cast `fn_item` with `as fn() -> _` to use it as a pointer
+    let _cast_zst_ptr = &foo as *const _;
+    //~^ WARN cast `foo` with `as fn() -> _` to use it as a pointer
+    let _coerced_zst_ptr: *const _ = &foo;
+    //~^ WARN cast `foo` with `as fn() -> _` to use it as a pointer
+
+    let _zst_ref = &mut foo;
+    //~^ WARN cast `foo` with `as fn() -> _` to use it as a pointer
+    let mut mut_fn_item = foo;
+    let _indirect_ref = &mut mut_fn_item;
+    //~^ WARN cast `fn_item` with `as fn() -> _` to use it as a pointer
+    let _cast_zst_ptr = &mut foo as *mut _;
+    //~^ WARN cast `foo` with `as fn() -> _` to use it as a pointer
+    let _coerced_zst_ptr: *mut _ = &mut foo;
+    //~^ WARN cast `foo` with `as fn() -> _` to use it as a pointer
+
+    let _cast_zst_ref = &foo as &dyn Fn() -> u32;
+    let _coerced_zst_ref: &dyn Fn() -> u32 = &foo;
+
+    let _cast_zst_ref = &mut foo as &mut dyn Fn() -> u32;
+    let _coerced_zst_ref: &mut dyn Fn() -> u32 = &mut foo;
+    let _fn_ptr = foo as fn() -> u32;
+
     println!("{:p}", &foo);
-    //~^ WARN cast `foo` with `as *const fn() -> _` to use it as a pointer
+    //~^ WARN cast `foo` with as fn() -> _` to use it as a pointer
     println!("{:p}", &bar);
-    //~^ WARN cast `bar` with `as *const fn(_) -> _` to use it as a pointer
+    //~^ WARN cast `bar` with as fn(_) -> _` to use it as a pointer
     println!("{:p}", &baz);
-    //~^ WARN cast `baz` with `as *const fn(_, _) -> _` to use it as a pointer
+    //~^ WARN cast `baz` with as fn(_, _) -> _` to use it as a pointer
+    println!("{:p}", &unsafe_fn);
+    //~^ WARN cast `baz` with as unsafe fn()` to use it as a pointer
+    println!("{:p}", &c_fn);
+    //~^ WARN cast `baz` with as extern "C" fn()` to use it as a pointer
+    println!("{:p}", &unsafe_c_fn);
+    //~^ WARN cast `baz` with as unsafe extern "C" fn()` to use it as a pointer
+    println!("{:p}", &variadic_fn);
+    //~^ WARN cast `baz` with as unsafe extern "C" fn(_, ...) -> _` to use it as a pointer
+    println!("{:p}", &std::env::var::<String>);
+    //~^ WARN cast `std::env::var` with as fn(_) -> _` to use it as a pointer
+
+    println!("{:p}", foo as fn() -> u32);
 
-    //should not produce any warnings
-    println!("{:p}", foo as *const fn() -> usize);
-    println!("{:p}", bar as *const fn(usize) -> usize);
-    println!("{:p}", baz as *const fn(usize, usize) -> usize);
+    unsafe {
+        std::mem::transmute::<_, usize>(&foo);
+        //~^ WARN cast `foo` with as fn() -> _` to use it as a pointer
+        std::mem::transmute::<_, usize>(foo as fn() -> u32);
+    }
 
-    //should not produce any warnings
-    let fn_thing = foo;
-    println!("{:p}", &fn_thing);
+    (&bar)(1);
+    call_fn(&bar, 1);
+    parameterized_call_fn(&bar, 1);
 }
diff --git a/src/test/ui/lint/function-references.stderr b/src/test/ui/lint/function-references.stderr
index bd2da6e1167..62dbf7f835d 100644
--- a/src/test/ui/lint/function-references.stderr
+++ b/src/test/ui/lint/function-references.stderr
@@ -1,22 +1,105 @@
-warning: cast `foo` with `as *const fn() -> _` to use it as a pointer
-  --> $DIR/function-references.rs:7:23
+warning: cast `foo` with `as fn() -> _` to use it as a pointer
+  --> $DIR/function-references.rs:16:20
    |
-LL |     println!("{:p}", &foo);
-   |                       ^^^
+LL |     let _zst_ref = &foo;
+   |                    ^^^^
    |
    = note: `#[warn(function_references)]` on by default
 
-warning: cast `bar` with `as *const fn(_) -> _` to use it as a pointer
-  --> $DIR/function-references.rs:9:23
+warning: cast `fn_item` with `as fn() -> _` to use it as a pointer
+  --> $DIR/function-references.rs:19:25
+   |
+LL |     let _indirect_ref = &fn_item;
+   |                         ^^^^^^^^
+
+warning: cast `foo` with `as fn() -> _` to use it as a pointer
+  --> $DIR/function-references.rs:21:25
+   |
+LL |     let _cast_zst_ptr = &foo as *const _;
+   |                         ^^^^
+
+warning: cast `foo` with `as fn() -> _` to use it as a pointer
+  --> $DIR/function-references.rs:23:38
+   |
+LL |     let _coerced_zst_ptr: *const _ = &foo;
+   |                                      ^^^^
+
+warning: cast `foo` with `as fn() -> _` to use it as a pointer
+  --> $DIR/function-references.rs:26:20
+   |
+LL |     let _zst_ref = &mut foo;
+   |                    ^^^^^^^^
+
+warning: cast `mut_fn_item` with `as fn() -> _` to use it as a pointer
+  --> $DIR/function-references.rs:29:25
+   |
+LL |     let _indirect_ref = &mut mut_fn_item;
+   |                         ^^^^^^^^^^^^^^^^
+
+warning: cast `foo` with `as fn() -> _` to use it as a pointer
+  --> $DIR/function-references.rs:31:25
+   |
+LL |     let _cast_zst_ptr = &mut foo as *mut _;
+   |                         ^^^^^^^^
+
+warning: cast `foo` with `as fn() -> _` to use it as a pointer
+  --> $DIR/function-references.rs:33:36
+   |
+LL |     let _coerced_zst_ptr: *mut _ = &mut foo;
+   |                                    ^^^^^^^^
+
+warning: cast `foo` with `as fn() -> _` to use it as a pointer
+  --> $DIR/function-references.rs:43:22
+   |
+LL |     println!("{:p}", &foo);
+   |                      ^^^^
+
+warning: cast `bar` with `as fn(_) -> _` to use it as a pointer
+  --> $DIR/function-references.rs:45:22
    |
 LL |     println!("{:p}", &bar);
-   |                       ^^^
+   |                      ^^^^
 
-warning: cast `baz` with `as *const fn(_, _) -> _` to use it as a pointer
-  --> $DIR/function-references.rs:11:23
+warning: cast `baz` with `as fn(_, _) -> _` to use it as a pointer
+  --> $DIR/function-references.rs:47:22
    |
 LL |     println!("{:p}", &baz);
-   |                       ^^^
+   |                      ^^^^
+
+warning: cast `unsafe_fn` with `as unsafe fn()` to use it as a pointer
+  --> $DIR/function-references.rs:49:22
+   |
+LL |     println!("{:p}", &unsafe_fn);
+   |                      ^^^^^^^^^^
+
+warning: cast `c_fn` with `as extern "C" fn()` to use it as a pointer
+  --> $DIR/function-references.rs:51:22
+   |
+LL |     println!("{:p}", &c_fn);
+   |                      ^^^^^
 
-warning: 3 warnings emitted
+warning: cast `unsafe_c_fn` with `as unsafe extern "C" fn()` to use it as a pointer
+  --> $DIR/function-references.rs:53:22
+   |
+LL |     println!("{:p}", &unsafe_c_fn);
+   |                      ^^^^^^^^^^^^
+
+warning: cast `variadic_fn` with `as unsafe extern "C" fn(_, ...)` to use it as a pointer
+  --> $DIR/function-references.rs:55:22
+   |
+LL |     println!("{:p}", &variadic_fn);
+   |                      ^^^^^^^^^^^^
+
+warning: cast `std::env::var` with `as fn(_) -> _` to use it as a pointer
+  --> $DIR/function-references.rs:57:22
+   |
+LL |     println!("{:p}", &std::env::var::<String>);
+   |                      ^^^^^^^^^^^^^^^^^^^^^^^^
+
+warning: cast `foo` with `as fn() -> _` to use it as a pointer
+  --> $DIR/function-references.rs:63:41
+   |
+LL |         std::mem::transmute::<_, usize>(&foo);
+   |                                         ^^^^
 
+warning: 17 warnings emitted