11use clippy_utils:: diagnostics:: span_lint_hir_and_then;
22use clippy_utils:: ty:: is_type_diagnostic_item;
3- use clippy_utils:: usage:: is_potentially_mutated ;
3+ use clippy_utils:: usage:: is_potentially_local_place ;
44use clippy_utils:: { higher, path_to_local} ;
55use if_chain:: if_chain;
66use rustc_errors:: Applicability ;
77use rustc_hir:: intravisit:: { walk_expr, walk_fn, FnKind , Visitor } ;
8- use rustc_hir:: { BinOpKind , Body , Expr , ExprKind , FnDecl , HirId , PathSegment , UnOp } ;
8+ use rustc_hir:: { BinOpKind , Body , Expr , ExprKind , FnDecl , HirId , Node , PathSegment , UnOp } ;
9+ use rustc_hir_typeck:: expr_use_visitor:: { Delegate , ExprUseVisitor , PlaceWithHirId } ;
10+ use rustc_infer:: infer:: TyCtxtInferExt ;
911use rustc_lint:: { LateContext , LateLintPass } ;
1012use rustc_middle:: hir:: nested_filter;
1113use rustc_middle:: lint:: in_external_macro;
12- use rustc_middle:: ty:: Ty ;
14+ use rustc_middle:: mir:: FakeReadCause ;
15+ use rustc_middle:: ty:: { self , Ty , TyCtxt } ;
1316use rustc_session:: { declare_lint_pass, declare_tool_lint} ;
1417use rustc_span:: def_id:: LocalDefId ;
1518use rustc_span:: source_map:: Span ;
@@ -192,6 +195,55 @@ fn collect_unwrap_info<'tcx>(
192195 Vec :: new ( )
193196}
194197
198+ /// A HIR visitor delegate that checks if a local variable of type `Option<_>` is mutated,
199+ /// *except* for if `Option::as_mut` is called.
200+ /// The reason for why we allow that one specifically is that `.as_mut()` cannot change
201+ /// the option to `None`, and that is important because this lint relies on the fact that
202+ /// `is_some` + `unwrap` is equivalent to `if let Some(..) = ..`, which it would not be if
203+ /// the option is changed to None between `is_some` and `unwrap`.
204+ /// (And also `.as_mut()` is a somewhat common method that is still worth linting on.)
205+ struct MutationVisitor < ' tcx > {
206+ is_mutated : bool ,
207+ local_id : HirId ,
208+ tcx : TyCtxt < ' tcx > ,
209+ }
210+
211+ /// Checks if the parent of the expression pointed at by the given `HirId` is a call to
212+ /// `Option::as_mut`.
213+ ///
214+ /// Used by the mutation visitor to specifically allow `.as_mut()` calls.
215+ /// In particular, the `HirId` that the visitor receives is the id of the local expression
216+ /// (i.e. the `x` in `x.as_mut()`), and that is the reason for why we care about its parent
217+ /// expression: that will be where the actual method call is.
218+ fn is_option_as_mut_use ( tcx : TyCtxt < ' _ > , expr_id : HirId ) -> bool {
219+ if let Node :: Expr ( mutating_expr) = tcx. hir ( ) . get_parent ( expr_id)
220+ && let ExprKind :: MethodCall ( path, ..) = mutating_expr. kind
221+ {
222+ path. ident . name . as_str ( ) == "as_mut"
223+ } else {
224+ false
225+ }
226+ }
227+
228+ impl < ' tcx > Delegate < ' tcx > for MutationVisitor < ' tcx > {
229+ fn borrow ( & mut self , cat : & PlaceWithHirId < ' tcx > , diag_expr_id : HirId , bk : ty:: BorrowKind ) {
230+ if let ty:: BorrowKind :: MutBorrow = bk
231+ && is_potentially_local_place ( self . local_id , & cat. place )
232+ && !is_option_as_mut_use ( self . tcx , diag_expr_id)
233+ {
234+ self . is_mutated = true ;
235+ }
236+ }
237+
238+ fn mutate ( & mut self , _: & PlaceWithHirId < ' tcx > , _: HirId ) {
239+ self . is_mutated = true ;
240+ }
241+
242+ fn consume ( & mut self , _: & PlaceWithHirId < ' tcx > , _: HirId ) { }
243+
244+ fn fake_read ( & mut self , _: & PlaceWithHirId < ' tcx > , _: FakeReadCause , _: HirId ) { }
245+ }
246+
195247impl < ' a , ' tcx > UnwrappableVariablesVisitor < ' a , ' tcx > {
196248 fn visit_branch (
197249 & mut self ,
@@ -202,10 +254,26 @@ impl<'a, 'tcx> UnwrappableVariablesVisitor<'a, 'tcx> {
202254 ) {
203255 let prev_len = self . unwrappables . len ( ) ;
204256 for unwrap_info in collect_unwrap_info ( self . cx , if_expr, cond, branch, else_branch, true ) {
205- if is_potentially_mutated ( unwrap_info. local_id , cond, self . cx )
206- || is_potentially_mutated ( unwrap_info. local_id , branch, self . cx )
207- {
208- // if the variable is mutated, we don't know whether it can be unwrapped:
257+ let mut delegate = MutationVisitor {
258+ tcx : self . cx . tcx ,
259+ is_mutated : false ,
260+ local_id : unwrap_info. local_id ,
261+ } ;
262+
263+ let infcx = self . cx . tcx . infer_ctxt ( ) . build ( ) ;
264+ let mut vis = ExprUseVisitor :: new (
265+ & mut delegate,
266+ & infcx,
267+ cond. hir_id . owner . def_id ,
268+ self . cx . param_env ,
269+ self . cx . typeck_results ( ) ,
270+ ) ;
271+ vis. walk_expr ( cond) ;
272+ vis. walk_expr ( branch) ;
273+
274+ if delegate. is_mutated {
275+ // if the variable is mutated, we don't know whether it can be unwrapped.
276+ // it might have been changed to `None` in between `is_some` + `unwrap`.
209277 continue ;
210278 }
211279 self . unwrappables . push ( unwrap_info) ;
@@ -215,6 +283,27 @@ impl<'a, 'tcx> UnwrappableVariablesVisitor<'a, 'tcx> {
215283 }
216284}
217285
286+ enum AsRefKind {
287+ AsRef ,
288+ AsMut ,
289+ }
290+
291+ /// Checks if the expression is a method call to `as_{ref,mut}` and returns the receiver of it.
292+ /// If it isn't, the expression itself is returned.
293+ fn consume_option_as_ref < ' tcx > ( expr : & ' tcx Expr < ' tcx > ) -> ( & ' tcx Expr < ' tcx > , Option < AsRefKind > ) {
294+ if let ExprKind :: MethodCall ( path, recv, ..) = expr. kind {
295+ if path. ident . name == sym:: as_ref {
296+ ( recv, Some ( AsRefKind :: AsRef ) )
297+ } else if path. ident . name . as_str ( ) == "as_mut" {
298+ ( recv, Some ( AsRefKind :: AsMut ) )
299+ } else {
300+ ( expr, None )
301+ }
302+ } else {
303+ ( expr, None )
304+ }
305+ }
306+
218307impl < ' a , ' tcx > Visitor < ' tcx > for UnwrappableVariablesVisitor < ' a , ' tcx > {
219308 type NestedFilter = nested_filter:: OnlyBodies ;
220309
@@ -233,6 +322,7 @@ impl<'a, 'tcx> Visitor<'tcx> for UnwrappableVariablesVisitor<'a, 'tcx> {
233322 // find `unwrap[_err]()` calls:
234323 if_chain ! {
235324 if let ExprKind :: MethodCall ( method_name, self_arg, ..) = expr. kind;
325+ let ( self_arg, as_ref_kind) = consume_option_as_ref( self_arg) ;
236326 if let Some ( id) = path_to_local( self_arg) ;
237327 if [ sym:: unwrap, sym:: expect, sym!( unwrap_err) ] . contains( & method_name. ident. name) ;
238328 let call_to_unwrap = [ sym:: unwrap, sym:: expect] . contains( & method_name. ident. name) ;
@@ -268,7 +358,12 @@ impl<'a, 'tcx> Visitor<'tcx> for UnwrappableVariablesVisitor<'a, 'tcx> {
268358 unwrappable. check. span. with_lo( unwrappable. if_expr. span. lo( ) ) ,
269359 "try" ,
270360 format!(
271- "if let {suggested_pattern} = {unwrappable_variable_name}" ,
361+ "if let {suggested_pattern} = {borrow_prefix}{unwrappable_variable_name}" ,
362+ borrow_prefix = match as_ref_kind {
363+ Some ( AsRefKind :: AsRef ) => "&" ,
364+ Some ( AsRefKind :: AsMut ) => "&mut " ,
365+ None => "" ,
366+ } ,
272367 ) ,
273368 // We don't track how the unwrapped value is used inside the
274369 // block or suggest deleting the unwrap, so we can't offer a
0 commit comments