Skip to content

Commit b86c5db

Browse files
arora-amanarora-aman
committed
Implement reborrow for closure captures
1 parent d1206f9 commit b86c5db

File tree

6 files changed

+121
-53
lines changed

6 files changed

+121
-53
lines changed

‎compiler/rustc_typeck/src/check/upvar.rs

Lines changed: 62 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
180180
debug!("seed place {:?}", place);
181181

182182
let upvar_id = ty::UpvarId::new(*var_hir_id, local_def_id);
183-
let capture_kind = self.init_capture_kind(capture_clause, upvar_id, span);
183+
let capture_kind =
184+
self.init_capture_kind_for_place(&place, capture_clause, upvar_id, span);
184185
let fake_info = ty::CaptureInfo {
185186
capture_kind_expr_id: None,
186187
path_expr_id: None,
@@ -449,7 +450,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
449450
base => bug!("Expected upvar, found={:?}", base),
450451
};
451452

452-
let place = restrict_capture_precision(place, capture_info.capture_kind);
453+
let place = restrict_capture_precision(place);
453454

454455
let min_cap_list = match root_var_min_capture_list.get_mut(&var_hir_id) {
455456
None => {
@@ -897,15 +898,24 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
897898
}
898899
}
899900

900-
fn init_capture_kind(
901+
fn init_capture_kind_for_place(
901902
&self,
903+
place: &Place<'tcx>,
902904
capture_clause: hir::CaptureBy,
903905
upvar_id: ty::UpvarId,
904906
closure_span: Span,
905907
) -> ty::UpvarCapture<'tcx> {
906908
match capture_clause {
907-
hir::CaptureBy::Value => ty::UpvarCapture::ByValue(None),
908-
hir::CaptureBy::Ref => {
909+
// In case of a move closure if the data is accessed through a reference we
910+
// want to capture by ref to allow precise capture using reborrows.
911+
//
912+
// If the data will be moved out of this place, then the place will be truncated
913+
// at the first Deref in `adjust_upvar_borrow_kind_for_consume` and then moved into
914+
// the closure.
915+
hir::CaptureBy::Value if !place.deref_tys().any(ty::TyS::is_ref) => {
916+
ty::UpvarCapture::ByValue(None)
917+
}
918+
hir::CaptureBy::Value | hir::CaptureBy::Ref => {
909919
let origin = UpvarRegion(upvar_id, closure_span);
910920
let upvar_region = self.next_region_var(origin);
911921
let upvar_borrow = ty::UpvarBorrow { kind: ty::ImmBorrow, region: upvar_region };
@@ -1109,12 +1119,18 @@ impl<'a, 'tcx> InferBorrowKind<'a, 'tcx> {
11091119
place_with_id, diag_expr_id, mode
11101120
);
11111121

1112-
// we only care about moves
1113-
match mode {
1114-
euv::Copy => {
1115-
return;
1116-
}
1117-
euv::Move => {}
1122+
let place = truncate_capture_for_move(place_with_id.place.clone());
1123+
match (self.capture_clause, mode) {
1124+
// In non-move closures, we only care about moves
1125+
(hir::CaptureBy::Ref, euv::Copy) => return,
1126+
1127+
(hir::CaptureBy::Ref, euv::Move) | (hir::CaptureBy::Value, euv::Move | euv::Copy) => {}
1128+
};
1129+
1130+
let place_with_id = PlaceWithHirId { place: place.clone(), hir_id: place_with_id.hir_id };
1131+
1132+
if !self.capture_information.contains_key(&place) {
1133+
self.init_capture_info_for_place(&place_with_id, diag_expr_id);
11181134
}
11191135

11201136
let tcx = self.fcx.tcx;
@@ -1128,13 +1144,15 @@ impl<'a, 'tcx> InferBorrowKind<'a, 'tcx> {
11281144

11291145
let usage_span = tcx.hir().span(diag_expr_id);
11301146

1131-
// To move out of an upvar, this must be a FnOnce closure
1132-
self.adjust_closure_kind(
1133-
upvar_id.closure_expr_id,
1134-
ty::ClosureKind::FnOnce,
1135-
usage_span,
1136-
place_with_id.place.clone(),
1137-
);
1147+
if matches!(mode, euv::Move) {
1148+
// To move out of an upvar, this must be a FnOnce closure
1149+
self.adjust_closure_kind(
1150+
upvar_id.closure_expr_id,
1151+
ty::ClosureKind::FnOnce,
1152+
usage_span,
1153+
place.clone(),
1154+
);
1155+
}
11381156

11391157
let capture_info = ty::CaptureInfo {
11401158
capture_kind_expr_id: Some(diag_expr_id),
@@ -1317,8 +1335,12 @@ impl<'a, 'tcx> InferBorrowKind<'a, 'tcx> {
13171335
if let PlaceBase::Upvar(upvar_id) = place_with_id.place.base {
13181336
assert_eq!(self.closure_def_id.expect_local(), upvar_id.closure_expr_id);
13191337

1320-
let capture_kind =
1321-
self.fcx.init_capture_kind(self.capture_clause, upvar_id, self.closure_span);
1338+
let capture_kind = self.fcx.init_capture_kind_for_place(
1339+
&place_with_id.place,
1340+
self.capture_clause,
1341+
upvar_id,
1342+
self.closure_span,
1343+
);
13221344

13231345
let expr_id = Some(diag_expr_id);
13241346
let capture_info = ty::CaptureInfo {
@@ -1392,15 +1414,10 @@ impl<'a, 'tcx> euv::Delegate<'tcx> for InferBorrowKind<'a, 'tcx> {
13921414
}
13931415

13941416
/// Truncate projections so that following rules are obeyed by the captured `place`:
1395-
///
1396-
/// - No Derefs in move closure, this will result in value behind a reference getting moved.
13971417
/// - No projections are applied to raw pointers, since these require unsafe blocks. We capture
13981418
/// them completely.
13991419
/// - No Index projections are captured, since arrays are captured completely.
1400-
fn restrict_capture_precision<'tcx>(
1401-
mut place: Place<'tcx>,
1402-
capture_kind: ty::UpvarCapture<'tcx>,
1403-
) -> Place<'tcx> {
1420+
fn restrict_capture_precision<'tcx>(mut place: Place<'tcx>) -> Place<'tcx> {
14041421
if place.projections.is_empty() {
14051422
// Nothing to do here
14061423
return place;
@@ -1412,7 +1429,6 @@ fn restrict_capture_precision<'tcx>(
14121429
}
14131430

14141431
let mut truncated_length = usize::MAX;
1415-
let mut first_deref_projection = usize::MAX;
14161432

14171433
for (i, proj) in place.projections.iter().enumerate() {
14181434
if proj.ty.is_unsafe_ptr() {
@@ -1426,31 +1442,36 @@ fn restrict_capture_precision<'tcx>(
14261442
truncated_length = truncated_length.min(i);
14271443
break;
14281444
}
1429-
ProjectionKind::Deref => {
1430-
// We only drop Derefs in case of move closures
1431-
// There might be an index projection or raw ptr ahead, so we don't stop here.
1432-
first_deref_projection = first_deref_projection.min(i);
1433-
}
1445+
ProjectionKind::Deref => {}
14341446
ProjectionKind::Field(..) => {} // ignore
14351447
ProjectionKind::Subslice => {} // We never capture this
14361448
}
14371449
}
14381450

1439-
let length = place
1440-
.projections
1441-
.len()
1442-
.min(truncated_length)
1443-
// In case of capture `ByValue` we want to not capture derefs
1444-
.min(match capture_kind {
1445-
ty::UpvarCapture::ByValue(..) => first_deref_projection,
1446-
ty::UpvarCapture::ByRef(..) => usize::MAX,
1447-
});
1451+
let length = place.projections.len().min(truncated_length);
14481452

14491453
place.projections.truncate(length);
14501454

14511455
place
14521456
}
14531457

1458+
/// Truncates a place so that the resultant capture doesn't move data out of a reference
1459+
fn truncate_capture_for_move(mut place: Place<'tcx>) -> Place<'tcx> {
1460+
for (i, proj) in place.projections.iter().enumerate() {
1461+
match proj.kind {
1462+
ProjectionKind::Deref => {
1463+
// We only drop Derefs in case of move closures
1464+
// There might be an index projection or raw ptr ahead, so we don't stop here.
1465+
place.projections.truncate(i);
1466+
return place;
1467+
}
1468+
_ => {}
1469+
}
1470+
}
1471+
1472+
place
1473+
}
1474+
14541475
fn construct_place_string(tcx: TyCtxt<'_>, place: &Place<'tcx>) -> String {
14551476
let variable_name = match place.base {
14561477
PlaceBase::Upvar(upvar_id) => var_name(tcx, upvar_id.var_path.hir_id).to_string(),

‎src/test/ui/closures/2229_closure_analysis/by_value.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ fn big_box() {
2626
//~^ First Pass analysis includes:
2727
//~| Min Capture analysis includes:
2828
let p = t.0.0;
29-
//~^ NOTE: Capturing t[(0, 0),Deref,(0, 0)] -> ByValue
29+
//~^ NOTE: Capturing t[(0, 0),Deref,(0, 0)] -> ImmBorrow
30+
//~| NOTE: Capturing t[(0, 0)] -> ByValue
3031
//~| NOTE: Min Capture t[(0, 0)] -> ByValue
3132
println!("{} {:?}", t.1, p);
3233
//~^ NOTE: Capturing t[(1, 0)] -> ImmBorrow

‎src/test/ui/closures/2229_closure_analysis/by_value.stderr

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,18 @@ LL | |
2828
LL | | };
2929
| |_____^
3030
|
31-
note: Capturing t[(0, 0),Deref,(0, 0)] -> ByValue
31+
note: Capturing t[(0, 0),Deref,(0, 0)] -> ImmBorrow
32+
--> $DIR/by_value.rs:28:17
33+
|
34+
LL | let p = t.0.0;
35+
| ^^^^^
36+
note: Capturing t[(0, 0)] -> ByValue
3237
--> $DIR/by_value.rs:28:17
3338
|
3439
LL | let p = t.0.0;
3540
| ^^^^^
3641
note: Capturing t[(1, 0)] -> ImmBorrow
37-
--> $DIR/by_value.rs:31:29
42+
--> $DIR/by_value.rs:32:29
3843
|
3944
LL | println!("{} {:?}", t.1, p);
4045
| ^^^
@@ -57,7 +62,7 @@ note: Min Capture t[(0, 0)] -> ByValue
5762
LL | let p = t.0.0;
5863
| ^^^^^
5964
note: Min Capture t[(1, 0)] -> ImmBorrow
60-
--> $DIR/by_value.rs:31:29
65+
--> $DIR/by_value.rs:32:29
6166
|
6267
LL | println!("{} {:?}", t.1, p);
6368
| ^^^

‎src/test/ui/closures/2229_closure_analysis/move_closure.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@ fn simple_ref() {
1818
//~^ ERROR: First Pass analysis includes:
1919
//~| ERROR: Min Capture analysis includes:
2020
*ref_s += 10;
21-
//~^ NOTE: Capturing ref_s[Deref] -> ByValue
22-
//~| NOTE: Min Capture ref_s[] -> ByValue
21+
//~^ NOTE: Capturing ref_s[Deref] -> UniqueImmBorrow
22+
//~| NOTE: Min Capture ref_s[Deref] -> UniqueImmBorrow
2323
};
2424
c();
2525
}
@@ -39,8 +39,8 @@ fn struct_contains_ref_to_another_struct() {
3939
//~^ ERROR: First Pass analysis includes:
4040
//~| ERROR: Min Capture analysis includes:
4141
t.0.0 = "new s".into();
42-
//~^ NOTE: Capturing t[(0, 0),Deref,(0, 0)] -> ByValue
43-
//~| NOTE: Min Capture t[(0, 0)] -> ByValue
42+
//~^ NOTE: Capturing t[(0, 0),Deref,(0, 0)] -> UniqueImmBorrow
43+
//~| NOTE: Min Capture t[(0, 0),Deref,(0, 0)] -> UniqueImmBorrow
4444
};
4545

4646
c();

‎src/test/ui/closures/2229_closure_analysis/move_closure.stderr

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ LL | |
4646
LL | | };
4747
| |_____^
4848
|
49-
note: Capturing ref_s[Deref] -> ByValue
49+
note: Capturing ref_s[Deref] -> UniqueImmBorrow
5050
--> $DIR/move_closure.rs:20:9
5151
|
5252
LL | *ref_s += 10;
@@ -64,7 +64,7 @@ LL | |
6464
LL | | };
6565
| |_____^
6666
|
67-
note: Min Capture ref_s[] -> ByValue
67+
note: Min Capture ref_s[Deref] -> UniqueImmBorrow
6868
--> $DIR/move_closure.rs:20:9
6969
|
7070
LL | *ref_s += 10;
@@ -82,7 +82,7 @@ LL | |
8282
LL | | };
8383
| |_____^
8484
|
85-
note: Capturing t[(0, 0),Deref,(0, 0)] -> ByValue
85+
note: Capturing t[(0, 0),Deref,(0, 0)] -> UniqueImmBorrow
8686
--> $DIR/move_closure.rs:41:9
8787
|
8888
LL | t.0.0 = "new s".into();
@@ -100,7 +100,7 @@ LL | |
100100
LL | | };
101101
| |_____^
102102
|
103-
note: Min Capture t[(0, 0)] -> ByValue
103+
note: Min Capture t[(0, 0),Deref,(0, 0)] -> UniqueImmBorrow
104104
--> $DIR/move_closure.rs:41:9
105105
|
106106
LL | t.0.0 = "new s".into();

‎src/test/ui/closures/2229_closure_analysis/run_pass/move_closure.rs

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,50 @@ fn no_ref_nested() {
5656
c();
5757
}
5858

59+
struct A<'a>(&'a mut String, &'a mut String);
60+
// Test that reborrowing works as expected for move closures
61+
// by attempting a disjoint capture through a reference.
62+
fn disjoint_via_ref() {
63+
let mut x = String::new();
64+
let mut y = String::new();
65+
66+
let mut a = A(&mut x, &mut y);
67+
let a = &mut a;
68+
69+
let mut c1 = move || {
70+
a.0.truncate(0);
71+
};
72+
73+
let mut c2 = move || {
74+
a.1.truncate(0);
75+
};
76+
77+
c1();
78+
c2();
79+
}
80+
81+
// Test that even if a path is moved into the closure, the closure is not FnOnce
82+
// if the path is not moved by the closure call.
83+
fn data_moved_but_not_fn_once() {
84+
let x = Box::new(10i32);
85+
86+
let c = move || {
87+
// *x has type i32 which is Copy. So even though the box `x` will be moved
88+
// into the closure, `x` is never moved when the closure is called, i.e. the
89+
// ownership stays with the closure and therefore we can call the function multiple times.
90+
let _x = *x;
91+
};
92+
93+
c();
94+
c();
95+
}
96+
5997
fn main() {
6098
simple_ref();
6199
struct_contains_ref_to_another_struct();
62100
no_ref();
63101
no_ref_nested();
102+
103+
disjoint_via_ref();
104+
data_moved_but_not_fn_once();
64105
}

0 commit comments

Comments
 (0)