Swap order of methods in needless_range_loop suggestion in some cases

Author: pengowen123

clippy_lints/src/loops.rs

@@ -27,6 +27,7 @@ use crate::rustc::ty::subst::Subst;
 use crate::rustc_errors::Applicability;
 use crate::rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use std::iter::{once, Iterator};
+use std::mem;
 use crate::syntax::ast;
 use crate::syntax::source_map::Span;
 use crate::syntax_pos::BytePos;
@@ -1082,16 +1083,35 @@ fn check_for_loop_range<'a, 'tcx>(
                     format!(".skip({})", snippet(cx, start.span, ".."))
                 };
 
+                let mut end_is_start_plus_val = false;
+
                 let take = if let Some(end) = *end {
+                    let mut take_expr = end;
+
+                    if let ExprKind::Binary(ref op, ref left, ref right) = end.node {
+                        if let BinOpKind::Add = op.node {
+                            let start_equal_left = SpanlessEq::new(cx).eq_expr(start, left);
+                            let start_equal_right = SpanlessEq::new(cx).eq_expr(start, right);
+
+                            if start_equal_left {
+                                take_expr = right;
+                            } else if start_equal_right {
+                                take_expr = left;
+                            }
+
+                            end_is_start_plus_val = start_equal_left | start_equal_right;
+                        }
+                    }
+
                     if is_len_call(end, indexed) {
                         String::new()
                     } else {
                         match limits {
                             ast::RangeLimits::Closed => {
-                                let end = sugg::Sugg::hir(cx, end, "<count>");
-                                format!(".take({})", end + sugg::ONE)
+                                let take_expr = sugg::Sugg::hir(cx, take_expr, "<count>");
+                                format!(".take({})", take_expr + sugg::ONE)
                             },
-                            ast::RangeLimits::HalfOpen => format!(".take({})", snippet(cx, end.span, "..")),
+                            ast::RangeLimits::HalfOpen => format!(".take({})", snippet(cx, take_expr.span, "..")),
                         }
                     }
                 } else {
@@ -1104,6 +1124,14 @@ fn check_for_loop_range<'a, 'tcx>(
                     ("", "iter")
                 };
 
+                let take_is_empty = take.is_empty();
+                let mut method_1 = take;
+                let mut method_2 = skip;
+
+                if end_is_start_plus_val {
+                    mem::swap(&mut method_1, &mut method_2);
+                }
+
                 if visitor.nonindex {
                     span_lint_and_then(
                         cx,
@@ -1116,16 +1144,16 @@ fn check_for_loop_range<'a, 'tcx>(
                                 "consider using an iterator".to_string(),
                                 vec![
                                     (pat.span, format!("({}, <item>)", ident.name)),
-                                    (arg.span, format!("{}.{}().enumerate(){}{}", indexed, method, take, skip)),
+                                    (arg.span, format!("{}.{}().enumerate(){}{}", indexed, method, method_1, method_2)),
                                 ],
                             );
                         },
                     );
                 } else {
-                    let repl = if starts_at_zero && take.is_empty() {
+                    let repl = if starts_at_zero && take_is_empty {
                         format!("&{}{}", ref_mut, indexed)
                     } else {
-                        format!("{}.{}(){}{}", indexed, method, take, skip)
+                        format!("{}.{}(){}{}", indexed, method, method_1, method_2)
                     };
 
                     span_lint_and_then(

tests/ui/author/for_loop.stderr

@@ -62,4 +62,18 @@ fn main() {
         g[i] = g[i+1..].iter().sum();
     }
     assert_eq!(g, vec![20, 18, 15, 11, 6, 0]);
+
+    let x = 5;
+    let mut vec = vec![0; 9];
+
+    for i in x..x + 4 {
+        vec[i] += 1;
+    }
+
+    let x = 5;
+    let mut vec = vec![0; 10];
+
+    for i in x..=x + 4 {
+        vec[i] += 1;
+    }
 }

tests/ui/needless_range_loop.rs

@@ -30,5 +30,25 @@ help: consider using an iterator
 45 |     for <item> in &mut ms {
    |         ^^^^^^    ^^^^^^^
 
-error: aborting due to 3 previous errors
+error: the loop variable `i` is only used to index `vec`.
+  --> $DIR/needless_range_loop.rs:69:14
+   |
+69 |     for i in x..x + 4 {
+   |              ^^^^^^^^
+help: consider using an iterator
+   |
+69 |     for <item> in vec.iter_mut().skip(x).take(4) {
+   |         ^^^^^^    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: the loop variable `i` is only used to index `vec`.
+  --> $DIR/needless_range_loop.rs:76:14
+   |
+76 |     for i in x..=x + 4 {
+   |              ^^^^^^^^^
+help: consider using an iterator
+   |
+76 |     for <item> in vec.iter_mut().skip(x).take(4 + 1) {
+   |         ^^^^^^    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: aborting due to 5 previous errors