From 73501daa39de2a16fe1de3021ff41d0fd5dcf2cd Mon Sep 17 00:00:00 2001
From: Cameron Steffen <cam.steffen94@gmail.com>
Date: Wed, 20 Oct 2021 13:16:50 -0500
Subject: [PATCH] Add expr_visitor util

---
 clippy_utils/src/visitors.rs | 45 ++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/clippy_utils/src/visitors.rs b/clippy_utils/src/visitors.rs
index 503effbdad57..4a3181dd0276 100644
--- a/clippy_utils/src/visitors.rs
+++ b/clippy_utils/src/visitors.rs
@@ -6,6 +6,51 @@ use rustc_lint::LateContext;
 use rustc_middle::hir::map::Map;
 use std::ops::ControlFlow;
 
+/// Convenience method for creating a `Visitor` with just `visit_expr` overridden and nested
+/// bodies (i.e. closures) are visited.
+/// If the callback returns `true`, the expr just provided to the callback is walked.
+#[must_use]
+pub fn expr_visitor<'tcx>(cx: &LateContext<'tcx>, f: impl FnMut(&'tcx Expr<'tcx>) -> bool) -> impl Visitor<'tcx> {
+    struct V<'tcx, F> {
+        hir: Map<'tcx>,
+        f: F,
+    }
+    impl<'tcx, F: FnMut(&'tcx Expr<'tcx>) -> bool> Visitor<'tcx> for V<'tcx, F> {
+        type Map = Map<'tcx>;
+        fn nested_visit_map(&mut self) -> NestedVisitorMap<Self::Map> {
+            NestedVisitorMap::OnlyBodies(self.hir)
+        }
+
+        fn visit_expr(&mut self, expr: &'tcx Expr<'tcx>) {
+            if (self.f)(expr) {
+                walk_expr(self, expr);
+            }
+        }
+    }
+    V { hir: cx.tcx.hir(), f }
+}
+
+/// Convenience method for creating a `Visitor` with just `visit_expr` overridden and nested
+/// bodies (i.e. closures) are not visited.
+/// If the callback returns `true`, the expr just provided to the callback is walked.
+#[must_use]
+pub fn expr_visitor_no_bodies<'tcx>(f: impl FnMut(&'tcx Expr<'tcx>) -> bool) -> impl Visitor<'tcx> {
+    struct V<F>(F);
+    impl<'tcx, F: FnMut(&'tcx Expr<'tcx>) -> bool> Visitor<'tcx> for V<F> {
+        type Map = intravisit::ErasedMap<'tcx>;
+        fn nested_visit_map(&mut self) -> NestedVisitorMap<Self::Map> {
+            NestedVisitorMap::None
+        }
+
+        fn visit_expr(&mut self, e: &'tcx Expr<'_>) {
+            if (self.0)(e) {
+                walk_expr(self, e);
+            }
+        }
+    }
+    V(f)
+}
+
 /// returns `true` if expr contains match expr desugared from try
 fn contains_try(expr: &hir::Expr<'_>) -> bool {
     struct TryFinder {