NLL Diagnostic Review 3: Unions not reinitialized after assignment into field

src/librustc_mir/borrow_check/mod.rs

@@ -1719,12 +1719,13 @@ impl<'cx, 'gcx, 'tcx> MirBorrowckCtxt<'cx, 'gcx, 'tcx> {
             }
         }
 
-        fn check_parent_of_field<'cx, 'gcx, 'tcx>(this: &mut MirBorrowckCtxt<'cx, 'gcx, 'tcx>,
-                                                  context: Context,
-                                                  base: &Place<'tcx>,
-                                                  span: Span,
-                                                  flow_state: &Flows<'cx, 'gcx, 'tcx>)
-        {
+        fn check_parent_of_field<'cx, 'gcx, 'tcx>(
+            this: &mut MirBorrowckCtxt<'cx, 'gcx, 'tcx>,
+            context: Context,
+            base: &Place<'tcx>,
+            span: Span,
+            flow_state: &Flows<'cx, 'gcx, 'tcx>,
+        ) {
             // rust-lang/rust#21232: Until Rust allows reads from the
             // initialized parts of partially initialized structs, we
             // will, starting with the 2018 edition, reject attempts
@@ -1776,6 +1777,22 @@ impl<'cx, 'gcx, 'tcx> MirBorrowckCtxt<'cx, 'gcx, 'tcx> {
             }
 
             if let Some((prefix, mpi)) = shortest_uninit_seen {
+                // Check for a reassignment into a uninitialized field of a union (for example,
+                // after a move out). In this case, do not report a error here. There is an
+                // exception, if this is the first assignment into the union (that is, there is
+                // no move out from an earlier location) then this is an attempt at initialization
+                // of the union - we should error in that case.
+                let tcx = this.infcx.tcx;
+                if let ty::TyKind::Adt(def, _) = base.ty(this.mir, tcx).to_ty(tcx).sty {
+                    let moved_before_this = this.move_data.path_map[mpi].iter().any(|moi| {
+                        this.move_data.moves[*moi].source < context.loc
+                    });
+
+                    if def.is_union() && moved_before_this {
+                        return;
+                    }
+                }
+
                 this.report_use_of_moved_or_uninitialized(
                     context,
                     InitializationRequiringAction::PartialAssignment,

src/librustc_mir/dataflow/move_paths/builder.rs

@@ -430,6 +430,20 @@ impl<'b, 'a, 'gcx, 'tcx> Gatherer<'b, 'a, 'gcx, 'tcx> {
     fn gather_init(&mut self, place: &Place<'tcx>, kind: InitKind) {
         debug!("gather_init({:?}, {:?})", self.loc, place);
 
+        let place = match place {
+            // Check if we are assigning into a field of a union, if so, lookup the place
+            // of the union so it is marked as initialized again.
+            Place::Projection(box Projection {
+                base,
+                elem: ProjectionElem::Field(_, _),
+            }) if match base.ty(self.builder.mir, self.builder.tcx).to_ty(self.builder.tcx).sty {
+                    ty::TyKind::Adt(def, _) if def.is_union() => true,
+                    _ => false,
+            } => base,
+            // Otherwise, lookup the place.
+            _ => place,
+        };
+
         if let LookupResult::Exact(path) = self.builder.data.rev_lookup.find(place) {
             let init = self.builder.data.inits.push(Init {
                 location: InitLocation::Statement(self.loc),

src/test/ui/borrowck/borrowck-union-move-assign.nll.stderr

@@ -8,28 +8,6 @@ LL |             let a = u.a; //~ ERROR use of moved value: `u.a`
    |
    = note: move occurs because `u` has type `U`, which does not implement the `Copy` trait
 
-error[E0382]: use of moved value: `u`
-  --> $DIR/borrowck-union-move-assign.rs:33:21
-   |
-LL |             let a = u.a;
-   |                     --- value moved here
-LL |             u.a = A;
-LL |             let a = u.a; // OK
-   |                     ^^^ value used here after move
-   |
-   = note: move occurs because `u` has type `U`, which does not implement the `Copy` trait
-
-error[E0382]: use of moved value: `u`
-  --> $DIR/borrowck-union-move-assign.rs:39:21
-   |
-LL |             let a = u.a;
-   |                     --- value moved here
-LL |             u.b = B;
-LL |             let a = u.a; // OK
-   |                     ^^^ value used here after move
-   |
-   = note: move occurs because `u` has type `U`, which does not implement the `Copy` trait
-
-error: aborting due to 3 previous errors
+error: aborting due to previous error
 
 For more information about this error, try `rustc --explain E0382`.

src/test/ui/borrowck/reassignment_immutable_fields_overlapping.nll.stderr

@@ -4,6 +4,16 @@ error[E0381]: assign to part of possibly uninitialized variable: `x`
 LL |     x.a = 1;  //~ ERROR
    |     ^^^^^^^ use of possibly uninitialized `x`
 
-error: aborting due to previous error
+error[E0594]: cannot assign to `x.b`, as `x` is not declared as mutable
+  --> $DIR/reassignment_immutable_fields_overlapping.rs:23:5
+   |
+LL |     let x: Foo;
+   |         - help: consider changing this to be mutable: `mut x`
+LL |     x.a = 1;  //~ ERROR
+LL |     x.b = 22; //~ ERROR
+   |     ^^^^^^^^ cannot assign
+
+error: aborting due to 2 previous errors
 
-For more information about this error, try `rustc --explain E0381`.
+Some errors occurred: E0381, E0594.
+For more information about an error, try `rustc --explain E0381`.