@@ -67,19 +67,20 @@ use rustc_data_structures::unhash::UnhashMap;
6767use rustc_hir as hir;
6868use rustc_hir:: def:: { DefKind , Res } ;
6969use rustc_hir:: def_id:: DefId ;
70- use rustc_hir:: hir_id:: HirIdSet ;
70+ use rustc_hir:: hir_id:: { HirIdMap , HirIdSet } ;
7171use rustc_hir:: intravisit:: { self , walk_expr, ErasedMap , FnKind , NestedVisitorMap , Visitor } ;
7272use rustc_hir:: LangItem :: { ResultErr , ResultOk } ;
7373use rustc_hir:: {
7474 def, Arm , BindingAnnotation , Block , Body , Constness , Destination , Expr , ExprKind , FnDecl , GenericArgs , HirId , Impl ,
75- ImplItem , ImplItemKind , IsAsync , Item , ItemKind , LangItem , Local , MatchSource , Node , Param , Pat , PatKind , Path ,
76- PathSegment , PrimTy , QPath , Stmt , StmtKind , TraitItem , TraitItemKind , TraitRef , TyKind , UnOp ,
75+ ImplItem , ImplItemKind , IsAsync , Item , ItemKind , LangItem , Local , MatchSource , Mutability , Node , Param , Pat ,
76+ PatKind , Path , PathSegment , PrimTy , QPath , Stmt , StmtKind , TraitItem , TraitItemKind , TraitRef , TyKind , UnOp ,
7777} ;
7878use rustc_lint:: { LateContext , Level , Lint , LintContext } ;
7979use rustc_middle:: hir:: exports:: Export ;
8080use rustc_middle:: hir:: map:: Map ;
8181use rustc_middle:: ty as rustc_ty;
82- use rustc_middle:: ty:: { layout:: IntegerExt , DefIdTree , Ty , TyCtxt , TypeFoldable } ;
82+ use rustc_middle:: ty:: adjustment:: { Adjust , Adjustment , AutoBorrow } ;
83+ use rustc_middle:: ty:: { layout:: IntegerExt , DefIdTree , Ty , TyCtxt , TypeAndMut , TypeFoldable } ;
8384use rustc_semver:: RustcVersion ;
8485use rustc_session:: Session ;
8586use rustc_span:: hygiene:: { ExpnKind , MacroKind } ;
@@ -670,8 +671,82 @@ pub fn can_move_expr_to_closure_no_visit(
670671 }
671672}
672673
673- /// Checks if the expression can be moved into a closure as is.
674- pub fn can_move_expr_to_closure ( cx : & LateContext < ' tcx > , expr : & ' tcx Expr < ' _ > ) -> bool {
674+ /// How a local is captured by a closure
675+ #[ derive( Debug , Clone , Copy , PartialEq , Eq ) ]
676+ pub enum CaptureKind {
677+ Value ,
678+ Ref ( Mutability ) ,
679+ }
680+ impl std:: ops:: BitOr for CaptureKind {
681+ type Output = Self ;
682+ fn bitor ( self , rhs : Self ) -> Self :: Output {
683+ match ( self , rhs) {
684+ ( CaptureKind :: Value , _) | ( _, CaptureKind :: Value ) => CaptureKind :: Value ,
685+ ( CaptureKind :: Ref ( Mutability :: Mut ) , CaptureKind :: Ref ( _) )
686+ | ( CaptureKind :: Ref ( _) , CaptureKind :: Ref ( Mutability :: Mut ) ) => CaptureKind :: Ref ( Mutability :: Mut ) ,
687+ ( CaptureKind :: Ref ( Mutability :: Not ) , CaptureKind :: Ref ( Mutability :: Not ) ) => CaptureKind :: Ref ( Mutability :: Not ) ,
688+ }
689+ }
690+ }
691+ impl std:: ops:: BitOrAssign for CaptureKind {
692+ fn bitor_assign ( & mut self , rhs : Self ) {
693+ * self = * self | rhs;
694+ }
695+ }
696+
697+ /// Given an expression referencing a local, determines how it would be captured in a closure.
698+ /// Note as this will walk up to parent expressions until the capture can be determined it should
699+ /// only be used while making a closure somewhere a value is consumed. e.g. a block, match arm, or
700+ /// function argument (other than a receiver).
701+ pub fn capture_local_usage ( cx : & LateContext < ' tcx > , e : & Expr < ' _ > ) -> CaptureKind {
702+ debug_assert ! ( matches!(
703+ e. kind,
704+ ExprKind :: Path ( QPath :: Resolved ( None , Path { res: Res :: Local ( _) , .. } ) )
705+ ) ) ;
706+
707+ let map = cx. tcx . hir ( ) ;
708+ let mut child_id = e. hir_id ;
709+ let mut capture = CaptureKind :: Value ;
710+
711+ for ( parent_id, parent) in map. parent_iter ( e. hir_id ) {
712+ if let [ Adjustment {
713+ kind : Adjust :: Deref ( _) | Adjust :: Borrow ( AutoBorrow :: Ref ( ..) ) ,
714+ target,
715+ } , ref adjust @ ..] = * cx
716+ . typeck_results ( )
717+ . adjustments ( )
718+ . get ( child_id)
719+ . map_or ( & [ ] [ ..] , |x| & * * x)
720+ {
721+ if let rustc_ty:: RawPtr ( TypeAndMut { mutbl : mutability, .. } ) | rustc_ty:: Ref ( _, _, mutability) =
722+ * adjust. last ( ) . map_or ( target, |a| a. target ) . kind ( )
723+ {
724+ return CaptureKind :: Ref ( mutability) ;
725+ }
726+ }
727+
728+ if let Node :: Expr ( e) = parent {
729+ match e. kind {
730+ ExprKind :: AddrOf ( _, mutability, _) => return CaptureKind :: Ref ( mutability) ,
731+ ExprKind :: Index ( ..) | ExprKind :: Unary ( UnOp :: Deref , _) => capture = CaptureKind :: Ref ( Mutability :: Not ) ,
732+ ExprKind :: Assign ( lhs, ..) | ExprKind :: Assign ( _, lhs, _) if lhs. hir_id == child_id => {
733+ return CaptureKind :: Ref ( Mutability :: Mut ) ;
734+ } ,
735+ _ => break ,
736+ }
737+ } else {
738+ break ;
739+ }
740+
741+ child_id = parent_id;
742+ }
743+
744+ capture
745+ }
746+
747+ /// Checks if the expression can be moved into a closure as is. This will return a list of captures
748+ /// if so, otherwise, `None`.
749+ pub fn can_move_expr_to_closure ( cx : & LateContext < ' tcx > , expr : & ' tcx Expr < ' _ > ) -> Option < HirIdMap < CaptureKind > > {
675750 struct V < ' cx , ' tcx > {
676751 cx : & ' cx LateContext < ' tcx > ,
677752 // Stack of potential break targets contained in the expression.
@@ -680,6 +755,9 @@ pub fn can_move_expr_to_closure(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) ->
680755 locals : HirIdSet ,
681756 /// Whether this expression can be turned into a closure.
682757 allow_closure : bool ,
758+ /// Locals which need to be captured, and whether they need to be by value, reference, or
759+ /// mutable reference.
760+ captures : HirIdMap < CaptureKind > ,
683761 }
684762 impl Visitor < ' tcx > for V < ' _ , ' tcx > {
685763 type Map = ErasedMap < ' tcx > ;
@@ -691,13 +769,23 @@ pub fn can_move_expr_to_closure(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) ->
691769 if !self . allow_closure {
692770 return ;
693771 }
694- if let ExprKind :: Loop ( b, ..) = e. kind {
695- self . loops . push ( e. hir_id ) ;
696- self . visit_block ( b) ;
697- self . loops . pop ( ) ;
698- } else {
699- self . allow_closure &= can_move_expr_to_closure_no_visit ( self . cx , e, & self . loops , & self . locals ) ;
700- walk_expr ( self , e) ;
772+
773+ match e. kind {
774+ ExprKind :: Path ( QPath :: Resolved ( None , & Path { res : Res :: Local ( l) , .. } ) ) => {
775+ if !self . locals . contains ( & l) {
776+ let cap = capture_local_usage ( self . cx , e) ;
777+ self . captures . entry ( l) . and_modify ( |e| * e |= cap) . or_insert ( cap) ;
778+ }
779+ } ,
780+ ExprKind :: Loop ( b, ..) => {
781+ self . loops . push ( e. hir_id ) ;
782+ self . visit_block ( b) ;
783+ self . loops . pop ( ) ;
784+ } ,
785+ _ => {
786+ self . allow_closure &= can_move_expr_to_closure_no_visit ( self . cx , e, & self . loops , & self . locals ) ;
787+ walk_expr ( self , e) ;
788+ } ,
701789 }
702790 }
703791
@@ -713,9 +801,10 @@ pub fn can_move_expr_to_closure(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) ->
713801 allow_closure : true ,
714802 loops : Vec :: new ( ) ,
715803 locals : HirIdSet :: default ( ) ,
804+ captures : HirIdMap :: default ( ) ,
716805 } ;
717806 v. visit_expr ( expr) ;
718- v. allow_closure
807+ v. allow_closure . then ( || v . captures )
719808}
720809
721810/// Returns the method names and argument list of nested method call expressions that make up
0 commit comments