Auto merge of #55657 - davidtwco:issue-55651, r=pnkfelix

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

Fixes #55651, #55652.

This PR makes two changes:

First, it updates the dataflow builder to add an init for the place
containing a union if there is an assignment into the field of
that union.

Second, it stops a "use of uninitialized" error occuring when there is an
assignment into the field of an uninitialized union that was previously
initialized. Making this assignment would re-initialize the union, as
tested in `src/test/ui/borrowck/borrowck-union-move-assign.nll.stderr`.
The check for previous initialization ensures that we do not start
supporting partial initialization yet (cc #21232, #54499, #54986).

This PR also fixes #55652 which was marked as requiring investigation
as the changes in this PR add an error that was previously missing
(and mentioned in the review comments) and confirms that the error
that was present is correct and a result of earlier partial
initialization changes in NLL.

r? @pnkfelix (due to earlier work with partial initialization)
cc @nikomatsakis

6 files changed, 121 insertions, 31 deletions

diff --git a/src/librustc/mir/mod.rs b/src/librustc/mir/mod.rs
index a84226bf665..8a32e8f46c2 100644
--- a/src/librustc/mir/mod.rs
+++ b/src/librustc/mir/mod.rs
@@ -20,6 +20,7 @@ use mir::interpret::{ConstValue, EvalErrorKind, Scalar};
 use mir::visit::MirVisitable;
 use rustc_apfloat::ieee::{Double, Single};
 use rustc_apfloat::Float;
+use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::graph::dominators::{dominators, Dominators};
 use rustc_data_structures::graph::{self, GraphPredecessors, GraphSuccessors};
 use rustc_data_structures::indexed_vec::{Idx, IndexVec};
@@ -2720,6 +2721,36 @@ impl Location {
         }
     }
 
+    /// Returns `true` if `other` is earlier in the control flow graph than `self`.
+    pub fn is_predecessor_of<'tcx>(&self, other: Location, mir: &Mir<'tcx>) -> bool {
+        // If we are in the same block as the other location and are an earlier statement
+        // then we are a predecessor of `other`.
+        if self.block == other.block && self.statement_index < other.statement_index {
+            return true;
+        }
+
+        // If we're in another block, then we want to check that block is a predecessor of `other`.
+        let mut queue: Vec<BasicBlock> = mir.predecessors_for(other.block).clone();
+        let mut visited = FxHashSet::default();
+
+        while let Some(block) = queue.pop() {
+            // If we haven't visited this block before, then make sure we visit it's predecessors.
+            if visited.insert(block) {
+                queue.append(&mut mir.predecessors_for(block).clone());
+            } else {
+                continue;
+            }
+
+            // If we found the block that `self` is in, then we are a predecessor of `other` (since
+            // we found that block by looking at the predecessors of `other`).
+            if self.block == block {
+                return true;
+            }
+        }
+
+        false
+    }
+
     pub fn dominates(&self, other: Location, dominators: &Dominators<BasicBlock>) -> bool {
         if self.block == other.block {
             self.statement_index <= other.statement_index
diff --git a/src/librustc_mir/borrow_check/mod.rs b/src/librustc_mir/borrow_check/mod.rs
index 4e03f6f7f5e..eeb111dcd9b 100644
--- a/src/librustc_mir/borrow_check/mod.rs
+++ b/src/librustc_mir/borrow_check/mod.rs
@@ -1717,12 +1717,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
@@ -1774,6 +1775,24 @@ 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 {
+                    if def.is_union() {
+                        if this.move_data.path_map[mpi].iter().any(|moi| {
+                            this.move_data.moves[*moi].source.is_predecessor_of(
+                                context.loc, this.mir,
+                            )
+                        }) {
+                            return;
+                        }
+                    }
+                }
+
                 this.report_use_of_moved_or_uninitialized(
                     context,
                     InitializationRequiringAction::PartialAssignment,
diff --git a/src/librustc_mir/dataflow/move_paths/builder.rs b/src/librustc_mir/dataflow/move_paths/builder.rs
index 874e862de23..2399f0bc1a0 100644
--- a/src/librustc_mir/dataflow/move_paths/builder.rs
+++ b/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),
diff --git a/src/test/ui/borrowck/borrowck-union-move-assign.nll.stderr b/src/test/ui/borrowck/borrowck-union-move-assign.nll.stderr
index 423a44514db..a7125450e1c 100644
--- a/src/test/ui/borrowck/borrowck-union-move-assign.nll.stderr
+++ b/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`.
diff --git a/src/test/ui/borrowck/reassignment_immutable_fields_overlapping.nll.stderr b/src/test/ui/borrowck/reassignment_immutable_fields_overlapping.nll.stderr
index 7da9dbfc088..fb1fe277740 100644
--- a/src/test/ui/borrowck/reassignment_immutable_fields_overlapping.nll.stderr
+++ b/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`.
diff --git a/src/test/ui/nll/issue-55651.rs b/src/test/ui/nll/issue-55651.rs
new file mode 100644
index 00000000000..b9777831bf6
--- /dev/null
+++ b/src/test/ui/nll/issue-55651.rs
@@ -0,0 +1,38 @@
+// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// compile-pass
+
+#![feature(untagged_unions)]
+
+struct A;
+struct B;
+
+union U {
+    a: A,
+    b: B,
+}
+
+fn main() {
+    unsafe {
+        {
+            let mut u = U { a: A };
+            let a = u.a;
+            u.a = A;
+            let a = u.a; // OK
+        }
+        {
+            let mut u = U { a: A };
+            let a = u.a;
+            u.b = B;
+            let a = u.a; // OK
+        }
+    }
+}