@@ -67,19 +67,20 @@ use rustc_data_structures::unhash::UnhashMap;
67
67
use rustc_hir as hir;
68
68
use rustc_hir:: def:: { DefKind , Res } ;
69
69
use rustc_hir:: def_id:: DefId ;
70
- use rustc_hir:: hir_id:: HirIdSet ;
70
+ use rustc_hir:: hir_id:: { HirIdMap , HirIdSet } ;
71
71
use rustc_hir:: intravisit:: { self , walk_expr, ErasedMap , FnKind , NestedVisitorMap , Visitor } ;
72
72
use rustc_hir:: LangItem :: { ResultErr , ResultOk } ;
73
73
use rustc_hir:: {
74
74
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 ,
77
77
} ;
78
78
use rustc_lint:: { LateContext , Level , Lint , LintContext } ;
79
79
use rustc_middle:: hir:: exports:: Export ;
80
80
use rustc_middle:: hir:: map:: Map ;
81
81
use 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 } ;
83
84
use rustc_semver:: RustcVersion ;
84
85
use rustc_session:: Session ;
85
86
use rustc_span:: hygiene:: { ExpnKind , MacroKind } ;
@@ -670,8 +671,82 @@ pub fn can_move_expr_to_closure_no_visit(
670
671
}
671
672
}
672
673
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 > > {
675
750
struct V < ' cx , ' tcx > {
676
751
cx : & ' cx LateContext < ' tcx > ,
677
752
// 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<'_>) ->
680
755
locals : HirIdSet ,
681
756
/// Whether this expression can be turned into a closure.
682
757
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 > ,
683
761
}
684
762
impl Visitor < ' tcx > for V < ' _ , ' tcx > {
685
763
type Map = ErasedMap < ' tcx > ;
@@ -691,13 +769,23 @@ pub fn can_move_expr_to_closure(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) ->
691
769
if !self . allow_closure {
692
770
return ;
693
771
}
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
+ } ,
701
789
}
702
790
}
703
791
@@ -713,9 +801,10 @@ pub fn can_move_expr_to_closure(cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) ->
713
801
allow_closure : true ,
714
802
loops : Vec :: new ( ) ,
715
803
locals : HirIdSet :: default ( ) ,
804
+ captures : HirIdMap :: default ( ) ,
716
805
} ;
717
806
v. visit_expr ( expr) ;
718
- v. allow_closure
807
+ v. allow_closure . then ( || v . captures )
719
808
}
720
809
721
810
/// Returns the method names and argument list of nested method call expressions that make up
0 commit comments