Update detector

CoinFabrik · Aug 29, 2024 · a653f08 · a653f08
1 parent 53157fd
commit a653f08
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/detectors/unsafe-map-get/src/lib.rs b/detectors/unsafe-map-get/src/lib.rs
@@ -4,7 +4,6 @@ extern crate rustc_errors;
 extern crate rustc_hir;
 extern crate rustc_span;
 
-use clippy_utils::diagnostics::span_lint_and_sugg;
 use if_chain::if_chain;
 use rustc_errors::Applicability;
 use rustc_hir::{
@@ -13,7 +12,7 @@ use rustc_hir::{
 };
 use rustc_lint::{LateContext, LateLintPass, LintContext};
 use rustc_span::{def_id::LocalDefId, Span};
-use utils::{get_receiver_ident_name, is_soroban_map};
+use utils::is_soroban_map;
 
 const LINT_MESSAGE: &str = "Unsafe access on Map, method could panic.";
 const UNSAFE_GET_METHODS: [&str; 3] = ["get", "get_unchecked", "try_get_unchecked"];
@@ -36,8 +35,11 @@ struct UnsafeMapGetVisitor<'a, 'tcx> {
 }
 
 impl UnsafeMapGetVisitor<'_, '_> {
-    fn get_first_arg_str(&self, arg: Option<&Expr<'_>>) -> String {
-        arg.and_then(|arg| self.cx.sess().source_map().span_to_snippet(arg.span).ok())
+    fn get_receiver_ident_name<'tcx>(&self, receiver: &'tcx Expr<'tcx>) -> String {
+        self.cx
+            .sess()
+            .source_map()
+            .span_to_snippet(receiver.span)
             .unwrap_or_default()
     }
 }
@@ -49,15 +51,15 @@ impl<'a, 'tcx> Visitor<'tcx> for UnsafeMapGetVisitor<'a, 'tcx> {
             if UNSAFE_GET_METHODS.contains(&path_segment.ident.as_str());
             if is_soroban_map(self.cx, self.cx.typeck_results().node_type(receiver.hir_id));
             then {
-                let receiver_ident_name = get_receiver_ident_name(receiver);
-                let first_arg_str = self.get_first_arg_str(args.first());
+                let receiver_ident_name = self.get_receiver_ident_name(receiver);
+                let first_arg_str = self.get_receiver_ident_name(&args[0]);
                 span_lint_and_sugg(
                     self.cx,
                     UNSAFE_MAP_GET,
                     expr.span,
                     LINT_MESSAGE,
                     format!("Using `{}` on a Map is unsafe as it could panic, please use", path_segment.ident),
-                    format!("{}.try_get({})", receiver_ident_name, first_arg_str),
+                    format!("{}.try_get({}).unwrap_or_default()", receiver_ident_name, first_arg_str),
                     Applicability::MaybeIncorrect,
                 );
             }