Extend implicit_saturation_sub lint

clippy_lints/src/implicit_saturating_sub.rs

@@ -1,11 +1,13 @@
-use clippy_utils::diagnostics::span_lint_and_sugg;
-use clippy_utils::{higher, is_integer_literal, peel_blocks_with_stmt, SpanlessEq};
+use clippy_utils::diagnostics::{span_lint_and_sugg, span_lint_and_then};
+use clippy_utils::source::snippet_opt;
+use clippy_utils::{higher, is_integer_literal, path_to_local, peel_blocks, peel_blocks_with_stmt, SpanlessEq};
 use rustc_ast::ast::LitKind;
 use rustc_data_structures::packed::Pu128;
 use rustc_errors::Applicability;
-use rustc_hir::{BinOpKind, Expr, ExprKind, QPath};
+use rustc_hir::{BinOp, BinOpKind, Expr, ExprKind, HirId, QPath};
 use rustc_lint::{LateContext, LateLintPass};
 use rustc_session::declare_lint_pass;
+use rustc_span::Span;
 
 declare_clippy_lint! {
     /// ### What it does
@@ -50,9 +52,166 @@ impl<'tcx> LateLintPass<'tcx> for ImplicitSaturatingSub {
 
             // Check if the conditional expression is a binary operation
             && let ExprKind::Binary(ref cond_op, cond_left, cond_right) = cond.kind
+        {
+            check_with_condition(cx, expr, cond_op.node, cond_left, cond_right, then);
+        } else if let Some(higher::If {
+            cond,
+            then: if_block,
+            r#else: Some(else_block),
+        }) = higher::If::hir(expr)
+            && let ExprKind::Binary(ref cond_op, cond_left, cond_right) = cond.kind
+        {
+            check_manual_check(cx, expr, cond_op, cond_left, cond_right, if_block, else_block);
+        }
+    }
+}
 
+fn check_manual_check<'tcx>(
+    cx: &LateContext<'tcx>,
+    expr: &Expr<'tcx>,
+    condition: &BinOp,
+    left_hand: &Expr<'tcx>,
+    right_hand: &Expr<'tcx>,
+    if_block: &Expr<'tcx>,
+    else_block: &Expr<'tcx>,
+) {
+    let ty = cx.typeck_results().expr_ty(left_hand);
+    if ty.is_numeric() && !ty.is_signed() {
+        match condition.node {
+            BinOpKind::Gt | BinOpKind::Ge => check_gt(
+                cx,
+                condition.span,
+                expr.span,
+                left_hand,
+                right_hand,
+                if_block,
+                else_block,
+            ),
+            BinOpKind::Lt | BinOpKind::Le => check_gt(
+                cx,
+                condition.span,
+                expr.span,
+                right_hand,
+                left_hand,
+                if_block,
+                else_block,
+            ),
+            _ => {},
+        }
+    }
+}
+
+fn check_gt(
+    cx: &LateContext<'_>,
+    condition_span: Span,
+    expr_span: Span,
+    big_var: &Expr<'_>,
+    little_var: &Expr<'_>,
+    if_block: &Expr<'_>,
+    else_block: &Expr<'_>,
+) {
+    if let Some(big_var) = Var::new(big_var)
+        && let Some(little_var) = Var::new(little_var)
+    {
+        check_subtraction(cx, condition_span, expr_span, big_var, little_var, if_block, else_block);
+    }
+}
+
+struct Var {
+    span: Span,
+    hir_id: HirId,
+}
+
+impl Var {
+    fn new(expr: &Expr<'_>) -> Option<Self> {
+        path_to_local(expr).map(|hir_id| Self {
+            span: expr.span,
+            hir_id,
+        })
+    }
+}
+
+fn check_subtraction(
+    cx: &LateContext<'_>,
+    condition_span: Span,
+    expr_span: Span,
+    big_var: Var,
+    little_var: Var,
+    if_block: &Expr<'_>,
+    else_block: &Expr<'_>,
+) {
+    let if_block = peel_blocks(if_block);
+    let else_block = peel_blocks(else_block);
+    if is_integer_literal(if_block, 0) {
+        // We need to check this case as well to prevent infinite recursion.
+        if is_integer_literal(else_block, 0) {
+            // Well, seems weird but who knows?
+            return;
+        }
+        // If the subtraction is done in the `else` block, then we need to also revert the two
+        // variables as it means that the check was reverted too.
+        check_subtraction(cx, condition_span, expr_span, little_var, big_var, else_block, if_block);
+        return;
+    }
+    if is_integer_literal(else_block, 0)
+        && let ExprKind::Binary(op, left, right) = if_block.kind
+        && let BinOpKind::Sub = op.node
+    {
+        let local_left = path_to_local(left);
+        let local_right = path_to_local(right);
+        if Some(big_var.hir_id) == local_left && Some(little_var.hir_id) == local_right {
+            // This part of the condition is voluntarily split from the one before to ensure that
+            // if `snippet_opt` fails, it won't try the next conditions.
+            if let Some(big_var_snippet) = snippet_opt(cx, big_var.span)
+                && let Some(little_var_snippet) = snippet_opt(cx, little_var.span)
+            {
+                span_lint_and_sugg(
+                    cx,
+                    IMPLICIT_SATURATING_SUB,
+                    expr_span,
+                    "manual arithmetic check found",
+                    "replace it with",
+                    format!("{big_var_snippet}.saturating_sub({little_var_snippet})"),
+                    Applicability::MachineApplicable,
+                );
+            }
+        } else if Some(little_var.hir_id) == local_left
+            && Some(big_var.hir_id) == local_right
+            && let Some(big_var_snippet) = snippet_opt(cx, big_var.span)
+            && let Some(little_var_snippet) = snippet_opt(cx, little_var.span)
+        {
+            span_lint_and_then(
+                cx,
+                IMPLICIT_SATURATING_SUB,
+                condition_span,
+                "inverted arithmetic check before subtraction",
+                |diag| {
+                    diag.span_note(
+                        if_block.span,
+                        format!("this subtraction underflows when `{little_var_snippet} < {big_var_snippet}`"),
+                    );
+                    diag.span_suggestion(
+                        if_block.span,
+                        "try replacing it with",
+                        format!("{big_var_snippet} - {little_var_snippet}"),
+                        Applicability::MaybeIncorrect,
+                    );
+                },
+            );
+        }
+    }
+}
+
+fn check_with_condition<'tcx>(
+    cx: &LateContext<'tcx>,
+    expr: &Expr<'tcx>,
+    cond_op: BinOpKind,
+    cond_left: &Expr<'tcx>,
+    cond_right: &Expr<'tcx>,
+    then: &Expr<'tcx>,
+) {
     // Ensure that the binary operator is >, !=, or <
-            && (BinOpKind::Ne == cond_op.node || BinOpKind::Gt == cond_op.node || BinOpKind::Lt == cond_op.node)
+    if (BinOpKind::Ne == cond_op || BinOpKind::Gt == cond_op || BinOpKind::Lt == cond_op)
 
         // Check if assign operation is done
         && let Some(target) = subtracts_one(cx, then)
@@ -62,13 +221,13 @@ impl<'tcx> LateLintPass<'tcx> for ImplicitSaturatingSub {
     {
         // Handle symmetric conditions in the if statement
         let (cond_var, cond_num_val) = if SpanlessEq::new(cx).eq_expr(cond_left, target) {
-                if BinOpKind::Gt == cond_op.node || BinOpKind::Ne == cond_op.node {
+            if BinOpKind::Gt == cond_op || BinOpKind::Ne == cond_op {
                 (cond_left, cond_right)
             } else {
                 return;
             }
         } else if SpanlessEq::new(cx).eq_expr(cond_right, target) {
-                if BinOpKind::Lt == cond_op.node || BinOpKind::Ne == cond_op.node {
+            if BinOpKind::Lt == cond_op || BinOpKind::Ne == cond_op {
                 (cond_right, cond_left)
             } else {
                 return;
@@ -118,7 +277,6 @@ impl<'tcx> LateLintPass<'tcx> for ImplicitSaturatingSub {
             _ => (),
         }
     }
-    }
 }
 
 fn subtracts_one<'a>(cx: &LateContext<'_>, expr: &'a Expr<'a>) -> Option<&'a Expr<'a>> {