Improve predecessor detection.

It is necessary to detect whether we are making the first
assignment into a union. This is checked by looking at the moves and
checking if there are any from locations earlier in the control flow
graph.

This commit improves the detection of this by switching from a naive
method that compared only the statement and basic block indices with
a more robust method that looks at the predecessors of a location.

2 files changed, 39 insertions, 6 deletions

diff --git a/src/librustc/mir/mod.rs b/src/librustc/mir/mod.rs
index 636fe115746..21459c97304 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};
@@ -2706,6 +2707,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 b440cf87a24..625423baece 100644
--- a/src/librustc_mir/borrow_check/mod.rs
+++ b/src/librustc_mir/borrow_check/mod.rs
@@ -1784,12 +1784,14 @@ impl<'cx, 'gcx, 'tcx> MirBorrowckCtxt<'cx, 'gcx, 'tcx> {
                 // 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;
+                    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;
+                        }
                     }
                 }