Rewrite unwrap_in_result lint

clippy_lints/src/lib.rs

@@ -654,7 +654,7 @@ pub fn register_lint_passes(store: &mut rustc_lint::LintStore, conf: &'static Co
     store.register_early_pass(move || Box::new(nonstandard_macro_braces::MacroBraces::new(conf)));
     store.register_late_pass(|_| Box::<macro_use::MacroUseImports>::default());
     store.register_late_pass(|_| Box::new(pattern_type_mismatch::PatternTypeMismatch));
-    store.register_late_pass(|_| Box::new(unwrap_in_result::UnwrapInResult));
+    store.register_late_pass(|_| Box::<unwrap_in_result::UnwrapInResult>::default());
     store.register_late_pass(|_| Box::new(semicolon_if_nothing_returned::SemicolonIfNothingReturned));
     store.register_late_pass(|_| Box::new(async_yields_async::AsyncYieldsAsync));
     let attrs = attr_storage.clone();

clippy_lints/src/unwrap_in_result.rs

@@ -1,13 +1,12 @@
 use clippy_utils::diagnostics::span_lint_and_then;
-use clippy_utils::ty::is_type_diagnostic_item;
-use clippy_utils::visitors::for_each_expr;
-use clippy_utils::{method_chain_args, return_ty};
-use core::ops::ControlFlow;
-use rustc_hir as hir;
-use rustc_hir::ImplItemKind;
+use clippy_utils::{return_ty, sym};
+use rustc_hir::{
+    Body, BodyOwnerKind, Expr, ExprKind, FnSig, ImplItem, ImplItemKind, Item, ItemKind, OwnerId, PathSegment, QPath,
+};
 use rustc_lint::{LateContext, LateLintPass};
-use rustc_session::declare_lint_pass;
-use rustc_span::{Span, sym};
+use rustc_middle::ty::Ty;
+use rustc_session::impl_lint_pass;
+use rustc_span::{Ident, Span, Symbol};
 
 declare_clippy_lint! {
     /// ### What it does
@@ -57,62 +56,166 @@ declare_clippy_lint! {
     "functions of type `Result<..>` or `Option`<...> that contain `expect()` or `unwrap()`"
 }
 
-declare_lint_pass!(UnwrapInResult=> [UNWRAP_IN_RESULT]);
+impl_lint_pass!(UnwrapInResult=> [UNWRAP_IN_RESULT]);
 
-impl<'tcx> LateLintPass<'tcx> for UnwrapInResult {
-    fn check_impl_item(&mut self, cx: &LateContext<'tcx>, impl_item: &'tcx hir::ImplItem<'_>) {
-        if let ImplItemKind::Fn(ref _signature, _) = impl_item.kind
-            // first check if it's a method or function
-            // checking if its return type is `result` or `option`
-            && (is_type_diagnostic_item(cx, return_ty(cx, impl_item.owner_id), sym::Result)
-                || is_type_diagnostic_item(cx, return_ty(cx, impl_item.owner_id), sym::Option))
-        {
-            lint_impl_body(cx, impl_item.span, impl_item);
+#[derive(Clone, Copy, Eq, PartialEq)]
+enum OptionOrResult {
+    Option,
+    Result,
+}
+
+impl OptionOrResult {
+    fn with_article(self) -> &'static str {
+        if matches!(self, Self::Option) {
+            "an `Option`"
+        } else {
+            "a `Result`"
         }
     }
 }
 
-fn lint_impl_body<'tcx>(cx: &LateContext<'tcx>, impl_span: Span, impl_item: &'tcx hir::ImplItem<'_>) {
-    if let ImplItemKind::Fn(_, body_id) = impl_item.kind {
-        let body = cx.tcx.hir_body(body_id);
-        let typeck = cx.tcx.typeck(impl_item.owner_id.def_id);
-        let mut result = Vec::new();
-        let _: Option<!> = for_each_expr(cx, body.value, |e| {
-            // check for `expect`
-            if let Some(arglists) = method_chain_args(e, &[sym::expect]) {
-                let receiver_ty = typeck.expr_ty(arglists[0].0).peel_refs();
-                if is_type_diagnostic_item(cx, receiver_ty, sym::Option)
-                    || is_type_diagnostic_item(cx, receiver_ty, sym::Result)
-                {
-                    result.push(e.span);
-                }
-            }
-
-            // check for `unwrap`
-            if let Some(arglists) = method_chain_args(e, &[sym::unwrap]) {
-                let receiver_ty = typeck.expr_ty(arglists[0].0).peel_refs();
-                if is_type_diagnostic_item(cx, receiver_ty, sym::Option)
-                    || is_type_diagnostic_item(cx, receiver_ty, sym::Result)
-                {
-                    result.push(e.span);
-                }
-            }
-
-            ControlFlow::Continue(())
+struct OptionOrResultFn {
+    kind: OptionOrResult,
+    return_ty_span: Option<Span>,
+}
+
+#[derive(Default)]
+pub struct UnwrapInResult {
+    fn_stack: Vec<Option<OptionOrResultFn>>,
+    current_fn: Option<OptionOrResultFn>,
+}
+
+impl UnwrapInResult {
+    fn enter_item(&mut self, cx: &LateContext<'_>, fn_def_id: OwnerId, sig: &FnSig<'_>) {
+        self.fn_stack.push(self.current_fn.take());
+        self.current_fn = is_option_or_result(cx, return_ty(cx, fn_def_id)).map(|kind| OptionOrResultFn {
+            kind,
+            return_ty_span: Some(sig.decl.output.span()),
         });
+    }
 
-        // if we've found one, lint
-        if !result.is_empty() {
+    fn leave_item(&mut self) {
+        self.current_fn = self.fn_stack.pop().unwrap();
+    }
+}
+
+impl<'tcx> LateLintPass<'tcx> for UnwrapInResult {
+    fn check_impl_item(&mut self, cx: &LateContext<'tcx>, impl_item: &'tcx ImplItem<'_>) {
+        if let ImplItemKind::Fn(sig, _) = &impl_item.kind {
+            self.enter_item(cx, impl_item.owner_id, sig);
+        }
+    }
+
+    fn check_impl_item_post(&mut self, _: &LateContext<'tcx>, impl_item: &'tcx ImplItem<'tcx>) {
+        if let ImplItemKind::Fn(..) = impl_item.kind {
+            self.leave_item();
+        }
+    }
+
+    fn check_item(&mut self, cx: &LateContext<'tcx>, item: &'tcx Item<'tcx>) {
+        if let ItemKind::Fn {
+            has_body: true, sig, ..
+        } = &item.kind
+        {
+            self.enter_item(cx, item.owner_id, sig);
+        }
+    }
+
+    fn check_item_post(&mut self, _: &LateContext<'tcx>, item: &'tcx Item<'tcx>) {
+        if let ItemKind::Fn { has_body: true, .. } = item.kind {
+            self.leave_item();
+        }
+    }
+
+    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &Expr<'_>) {
+        if expr.span.from_expansion() {
+            return;
+        }
+
+        if let Some(OptionOrResultFn {
+            kind,
+            ref mut return_ty_span,
+        }) = self.current_fn
+            && let Some((oor, fn_name)) = is_unwrap_or_expect_call(cx, expr)
+            && oor == kind
+        {
             span_lint_and_then(
                 cx,
                 UNWRAP_IN_RESULT,
-                impl_span,
-                "used unwrap or expect in a function that returns result or option",
-                move |diag| {
-                    diag.help("unwrap and expect should not be used in a function that returns result or option");
-                    diag.span_note(result, "potential non-recoverable error(s)");
+                expr.span,
+                format!("`{fn_name}` used in a function that returns {}", kind.with_article()),
+                |diag| {
+                    // Issue the note and help only once per function
+                    if let Some(span) = return_ty_span.take() {
+                        diag.span_note(span, "in this function signature");
+                        let complement = if kind == OptionOrResult::Result {
+                            " or calling the `.map_err()` method"
+                        } else {
+                            ""
+                        };
+                        diag.help(format!("consider using the `?` operator{complement}"));
+                    }
                 },
             );
         }
     }
+
+    fn check_body(&mut self, cx: &LateContext<'tcx>, body: &Body<'tcx>) {
+        let body_def_id = cx.tcx.hir_body_owner_def_id(body.id());
+        if !matches!(cx.tcx.hir_body_owner_kind(body_def_id), BodyOwnerKind::Fn) {
+            // When entering a body which is not a function, mask the potential surrounding
+            // function to not apply the lint.
+            self.fn_stack.push(self.current_fn.take());
+        }
+    }
+
+    fn check_body_post(&mut self, cx: &LateContext<'tcx>, body: &Body<'tcx>) {
+        let body_def_id = cx.tcx.hir_body_owner_def_id(body.id());
+        if !matches!(cx.tcx.hir_body_owner_kind(body_def_id), BodyOwnerKind::Fn) {
+            // Unmask the potential surrounding function.
+            self.current_fn = self.fn_stack.pop().unwrap();
+        }
+    }
+}
+
+fn is_option_or_result(cx: &LateContext<'_>, ty: Ty<'_>) -> Option<OptionOrResult> {
+    match ty.ty_adt_def().and_then(|def| cx.tcx.get_diagnostic_name(def.did())) {
+        Some(sym::Option) => Some(OptionOrResult::Option),
+        Some(sym::Result) => Some(OptionOrResult::Result),
+        _ => None,
+    }
+}
+
+fn is_unwrap_or_expect_call(cx: &LateContext<'_>, expr: &Expr<'_>) -> Option<(OptionOrResult, Symbol)> {
+    if let ExprKind::Call(func, _) = expr.kind
+        && let ExprKind::Path(QPath::TypeRelative(
+            hir_ty,
+            PathSegment {
+                ident:
+                    Ident {
+                        name: name @ (sym::unwrap | sym::expect),
+                        ..
+                    },
+                ..
+            },
+        )) = func.kind
+    {
+        is_option_or_result(cx, cx.typeck_results().node_type(hir_ty.hir_id)).map(|oor| (oor, *name))
+    } else if let ExprKind::MethodCall(
+        PathSegment {
+            ident: Ident {
+                name: name @ (sym::unwrap | sym::expect),
+                ..
+            },
+            ..
+        },
+        recv,
+        _,
+        _,
+    ) = expr.kind
+    {
+        is_option_or_result(cx, cx.typeck_results().expr_ty_adjusted(recv)).map(|oor| (oor, *name))
+    } else {
+        None
+    }
 }

tests/ui/unwrap_in_result.rs

@@ -1,4 +1,5 @@
 #![warn(clippy::unwrap_in_result)]
+#![allow(clippy::ok_expect)]
 
 struct A;
 
@@ -20,10 +21,9 @@ impl A {
 
     // should be detected
     fn bad_divisible_by_3(i_str: String) -> Result<bool, String> {
-        //~^ unwrap_in_result
-
         // checks whether a string represents a number divisible by 3
         let i = i_str.parse::<i32>().unwrap();
+        //~^ unwrap_in_result
         if i % 3 == 0 {
             Ok(true)
         } else {
@@ -32,23 +32,75 @@ impl A {
     }
 
     fn example_option_expect(i_str: String) -> Option<bool> {
+        let i = i_str.parse::<i32>().ok().expect("not a number");
         //~^ unwrap_in_result
-
-        let i = i_str.parse::<i32>().expect("not a number");
         if i % 3 == 0 {
             return Some(true);
         }
         None
     }
 
     fn in_closure(a: Option<bool>) -> Option<bool> {
-        //~^ unwrap_in_result
+        // No lint inside a closure
         let c = || a.unwrap();
-        Some(c())
+
+        // But lint outside
+        let a = c().then_some(true);
+        let _ = a.unwrap();
+        //~^ unwrap_in_result
+
+        None
     }
+
+    const fn in_const_inside_fn() -> bool {
+        const A: bool = {
+            const fn inner(b: Option<bool>) -> Option<bool> {
+                Some(b.unwrap())
+                //~^ unwrap_in_result
+            }
+
+            // No lint inside `const`
+            inner(Some(true)).unwrap()
+        };
+        A
+    }
+
+    fn in_static_inside_fn() -> bool {
+        static A: bool = {
+            const fn inner(b: Option<bool>) -> Option<bool> {
+                Some(b.unwrap())
+                //~^ unwrap_in_result
+            }
+
+            // No lint inside `static`
+            inner(Some(true)).unwrap()
+        };
+        A
+    }
+}
+
+macro_rules! mac {
+    () => {
+        Option::unwrap(Some(3))
+    };
+}
+
+fn type_relative_unwrap() -> Option<()> {
+    _ = Option::unwrap(Some(3));
+    //~^ unwrap_in_result
+
+    // Do not lint macro output
+    _ = mac!();
+
+    None
 }
 
-fn main() {
-    A::bad_divisible_by_3("3".to_string());
-    A::good_divisible_by_3("3".to_string());
+fn main() -> Result<(), ()> {
+    A::bad_divisible_by_3("3".to_string()).unwrap();
+    //~^ unwrap_in_result
+    A::good_divisible_by_3("3".to_string()).unwrap();
+    //~^ unwrap_in_result
+    Result::unwrap(A::good_divisible_by_3("3".to_string()));
+    //~^ unwrap_in_result
+    Ok(())
 }