Fix NRVO pass.

Author: cjgillot

compiler/rustc_mir_transform/src/nrvo.rs

@@ -1,10 +1,11 @@
 //! See the docs for [`RenameReturnPlace`].
 
 use rustc_hir::Mutability;
-use rustc_index::bit_set::HybridBitSet;
-use rustc_middle::mir::visit::{MutVisitor, NonUseContext, PlaceContext, Visitor};
-use rustc_middle::mir::{self, BasicBlock, Local, Location};
+use rustc_index::bit_set::BitSet;
+use rustc_middle::mir::visit::{MutVisitor, NonMutatingUseContext, PlaceContext, Visitor};
+use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
+use rustc_mir_dataflow::impls::borrowed_locals;
 
 use crate::MirPass;
 
@@ -38,17 +39,17 @@ impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
         sess.mir_opt_level() > 0 && sess.opts.unstable_opts.unsound_mir_opts
     }
 
-    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         let def_id = body.source.def_id();
+        if !tcx.consider_optimizing(|| format!("RenameReturnPlace {def_id:?}")) {
+            return;
+        }
+
         let Some(returned_local) = local_eligible_for_nrvo(body) else {
             debug!("`{:?}` was ineligible for NRVO", def_id);
             return;
         };
 
-        if !tcx.consider_optimizing(|| format!("RenameReturnPlace {def_id:?}")) {
-            return;
-        }
-
         debug!(
             "`{:?}` was eligible for NRVO, making {:?} the return place",
             def_id, returned_local
@@ -58,12 +59,11 @@ impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
 
         // Clean up the `NOP`s we inserted for statements made useless by our renaming.
         for block_data in body.basic_blocks.as_mut_preserves_cfg() {
-            block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop);
+            block_data.statements.retain(|stmt| stmt.kind != StatementKind::Nop);
         }
 
         // Overwrite the debuginfo of `_0` with that of the renamed local.
-        let (renamed_decl, ret_decl) =
-            body.local_decls.pick2_mut(returned_local, mir::RETURN_PLACE);
+        let (renamed_decl, ret_decl) = body.local_decls.pick2_mut(returned_local, RETURN_PLACE);
 
         // Sometimes, the return place is assigned a local of a different but coercible type, for
         // example `&mut T` instead of `&T`. Overwriting the `LocalInfo` for the return place means
@@ -84,26 +84,26 @@ impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
 ///
 /// If the MIR fulfills both these conditions, this function returns the `Local` that is assigned
 /// to the return place along all possible paths through the control-flow graph.
-fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> {
+fn local_eligible_for_nrvo(body: &mut Body<'_>) -> Option<Local> {
     if IsReturnPlaceRead::run(body) {
         return None;
     }
 
     let mut copied_to_return_place = None;
     for block in body.basic_blocks.indices() {
         // Look for blocks with a `Return` terminator.
-        if !matches!(body[block].terminator().kind, mir::TerminatorKind::Return) {
+        if !matches!(body[block].terminator().kind, TerminatorKind::Return) {
             continue;
         }
 
         // Look for an assignment of a single local to the return place prior to the `Return`.
         let returned_local = find_local_assigned_to_return_place(block, body)?;
         match body.local_kind(returned_local) {
             // FIXME: Can we do this for arguments as well?
-            mir::LocalKind::Arg => return None,
+            LocalKind::Arg => return None,
 
-            mir::LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"),
-            mir::LocalKind::Temp => {}
+            LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"),
+            LocalKind::Temp => {}
         }
 
         // If multiple different locals are copied to the return place. We can't pick a
@@ -118,20 +118,54 @@ fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> {
     copied_to_return_place
 }
 
-fn find_local_assigned_to_return_place(
-    start: BasicBlock,
-    body: &mut mir::Body<'_>,
-) -> Option<Local> {
-    let mut block = start;
-    let mut seen = HybridBitSet::new_empty(body.basic_blocks.len());
+#[instrument(level = "trace", skip(body), ret)]
+fn find_local_assigned_to_return_place(start: BasicBlock, body: &mut Body<'_>) -> Option<Local> {
+    // The locals that are assigned-to between `return` and `_0 = _rvo_local`.
+    let mut assigned_locals = BitSet::new_empty(body.local_decls.len());
+    // Whether we have seen an indirect write.
+    let mut seen_indirect = false;
+
+    let mut discard_borrowed_locals = |assigned_locals: &mut BitSet<Local>| {
+        // We have an indirect assignment to a local between the assignment to `_0 = _rvo`
+        // and `return`. This means we may be modifying the RVO local after the assignment.
+        // Discard all borrowed locals to be safe.
+        if !seen_indirect {
+            assigned_locals.union(&borrowed_locals(body));
+            // Mark that we have seen an indirect write to avoid recomputing `borrowed_locals`.
+            seen_indirect = true;
+        }
+    };
 
     // Iterate as long as `block` has exactly one predecessor that we have not yet visited.
-    while seen.insert(block) {
+    let mut block = start;
+    let mut seen_blocks = BitSet::new_empty(body.basic_blocks.len());
+    while seen_blocks.insert(block) {
         trace!("Looking for assignments to `_0` in {:?}", block);
+        let bbdata = &body.basic_blocks[block];
+
+        let mut vis = DiscardWrites { assigned_locals: &mut assigned_locals, seen_indirect: false };
+        vis.visit_terminator(&bbdata.terminator(), body.terminator_loc(block));
+        if vis.seen_indirect {
+            discard_borrowed_locals(&mut assigned_locals);
+        }
+
+        for (statement_index, stmt) in bbdata.statements.iter().enumerate().rev() {
+            if let StatementKind::Assign(box (lhs, ref rhs)) = stmt.kind
+                && lhs.as_local() == Some(RETURN_PLACE)
+                && let Rvalue::Use(rhs) = rhs
+                && let Some(rhs) = rhs.place()
+                && let Some(rhs) = rhs.as_local()
+                && !assigned_locals.contains(rhs)
+            {
+                return Some(rhs);
+            }
 
-        let local = body[block].statements.iter().rev().find_map(as_local_assigned_to_return_place);
-        if local.is_some() {
-            return local;
+            let mut vis =
+                DiscardWrites { assigned_locals: &mut assigned_locals, seen_indirect: false };
+            vis.visit_statement(stmt, Location { block, statement_index });
+            if vis.seen_indirect {
+                discard_borrowed_locals(&mut assigned_locals);
+            }
         }
 
         match body.basic_blocks.predecessors()[block].as_slice() {
@@ -145,10 +179,10 @@ fn find_local_assigned_to_return_place(
 
 // If this statement is an assignment of an unprojected local to the return place,
 // return that local.
-fn as_local_assigned_to_return_place(stmt: &mir::Statement<'_>) -> Option<Local> {
-    if let mir::StatementKind::Assign(box (lhs, rhs)) = &stmt.kind {
-        if lhs.as_local() == Some(mir::RETURN_PLACE) {
-            if let mir::Rvalue::Use(mir::Operand::Copy(rhs) | mir::Operand::Move(rhs)) = rhs {
+fn as_local_assigned_to_return_place(stmt: &Statement<'_>) -> Option<Local> {
+    if let StatementKind::Assign(box (lhs, rhs)) = &stmt.kind {
+        if lhs.as_local() == Some(RETURN_PLACE) {
+            if let Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)) = rhs {
                 return rhs.as_local();
             }
         }
@@ -168,51 +202,38 @@ impl<'tcx> MutVisitor<'tcx> for RenameToReturnPlace<'tcx> {
         self.tcx
     }
 
-    fn visit_statement(&mut self, stmt: &mut mir::Statement<'tcx>, loc: Location) {
+    fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
         // Remove assignments of the local being replaced to the return place, since it is now the
         // return place:
         //     _0 = _1
         if as_local_assigned_to_return_place(stmt) == Some(self.to_rename) {
-            stmt.kind = mir::StatementKind::Nop;
+            stmt.kind = StatementKind::Nop;
             return;
         }
 
         // Remove storage annotations for the local being replaced:
         //     StorageLive(_1)
-        if let mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) =
-            stmt.kind
-        {
+        if let StatementKind::StorageLive(local) | StatementKind::StorageDead(local) = stmt.kind {
             if local == self.to_rename {
-                stmt.kind = mir::StatementKind::Nop;
+                stmt.kind = StatementKind::Nop;
                 return;
             }
         }
 
         self.super_statement(stmt, loc)
     }
 
-    fn visit_terminator(&mut self, terminator: &mut mir::Terminator<'tcx>, loc: Location) {
-        // Ignore the implicit "use" of the return place in a `Return` statement.
-        if let mir::TerminatorKind::Return = terminator.kind {
-            return;
-        }
-
-        self.super_terminator(terminator, loc);
-    }
-
-    fn visit_local(&mut self, l: &mut Local, ctxt: PlaceContext, _: Location) {
-        if *l == mir::RETURN_PLACE {
-            assert_eq!(ctxt, PlaceContext::NonUse(NonUseContext::VarDebugInfo));
-        } else if *l == self.to_rename {
-            *l = mir::RETURN_PLACE;
+    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
+        if *l == self.to_rename {
+            *l = RETURN_PLACE;
         }
     }
 }
 
 struct IsReturnPlaceRead(bool);
 
 impl IsReturnPlaceRead {
-    fn run(body: &mir::Body<'_>) -> bool {
+    fn run(body: &Body<'_>) -> bool {
         let mut vis = IsReturnPlaceRead(false);
         vis.visit_body(body);
         vis.0
@@ -221,17 +242,35 @@ impl IsReturnPlaceRead {
 
 impl<'tcx> Visitor<'tcx> for IsReturnPlaceRead {
     fn visit_local(&mut self, l: Local, ctxt: PlaceContext, _: Location) {
-        if l == mir::RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() {
+        if l == RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() {
             self.0 = true;
         }
     }
 
-    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, loc: Location) {
+    fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, loc: Location) {
         // Ignore the implicit "use" of the return place in a `Return` statement.
-        if let mir::TerminatorKind::Return = terminator.kind {
+        if let TerminatorKind::Return = terminator.kind {
             return;
         }
 
         self.super_terminator(terminator, loc);
     }
 }
+
+struct DiscardWrites<'a> {
+    assigned_locals: &'a mut BitSet<Local>,
+    seen_indirect: bool,
+}
+
+impl<'tcx> Visitor<'tcx> for DiscardWrites<'_> {
+    fn visit_place(&mut self, place: &Place<'tcx>, ctxt: PlaceContext, _: Location) {
+        match ctxt {
+            PlaceContext::MutatingUse(_)
+            | PlaceContext::NonMutatingUse(NonMutatingUseContext::Move) => {
+                self.seen_indirect |= place.is_indirect_first_projection();
+                self.assigned_locals.insert(place.local);
+            }
+            PlaceContext::NonMutatingUse(_) | PlaceContext::NonUse(_) => {}
+        }
+    }
+}

tests/mir-opt/nrvo_miscompile_111005.call.RenameReturnPlace.diff

@@ -0,0 +1,22 @@
+- // MIR for `call` before RenameReturnPlace
++ // MIR for `call` after RenameReturnPlace
+  
+  fn call(_1: char) -> char {
+      let mut _0: char;
+      let mut _2: char;
+  
+      bb0: {
+          _0 = wrong(_1) -> [return: bb1, unwind continue];
+      }
+  
+      bb1: {
+          _2 = _1;
+          _0 = _2;
+          _2 = wrong(_1) -> [return: bb2, unwind continue];
+      }
+  
+      bb2: {
+          return;
+      }
+  }
+  

tests/mir-opt/nrvo_miscompile_111005.call_ok.RenameReturnPlace.diff

@@ -0,0 +1,22 @@
+- // MIR for `call_ok` before RenameReturnPlace
++ // MIR for `call_ok` after RenameReturnPlace
+  
+  fn call_ok(_1: char) -> char {
+      let mut _0: char;
+      let mut _2: char;
+  
+      bb0: {
+          _0 = wrong(_1) -> [return: bb1, unwind continue];
+      }
+  
+      bb1: {
+-         _2 = wrong(_1) -> [return: bb2, unwind continue];
++         _0 = wrong(_1) -> [return: bb2, unwind continue];
+      }
+  
+      bb2: {
+-         _0 = _2;
+          return;
+      }
+  }
+  

tests/mir-opt/nrvo_miscompile_111005.indirect.RenameReturnPlace.diff

@@ -0,0 +1,17 @@
+- // MIR for `indirect` before RenameReturnPlace
++ // MIR for `indirect` after RenameReturnPlace
+  
+  fn indirect(_1: char) -> char {
+      let mut _0: char;
+      let mut _2: char;
+      let mut _3: &mut char;
+  
+      bb0: {
+          _2 = _1;
+          _3 = &mut _2;
+          _0 = _2;
+          (*_3) = const 'b';
+          return;
+      }
+  }
+  

tests/mir-opt/nrvo_miscompile_111005.moved.RenameReturnPlace.diff

@@ -0,0 +1,16 @@
+- // MIR for `moved` before RenameReturnPlace
++ // MIR for `moved` after RenameReturnPlace
+  
+  fn moved(_1: char) -> char {
+      let mut _0: char;
+      let mut _2: char;
+      let mut _3: char;
+  
+      bb0: {
+          _2 = _1;
+          _0 = _2;
+          _3 = move _2;
+          return;
+      }
+  }
+  

tests/mir-opt/nrvo_miscompile_111005.multiple.RenameReturnPlace.diff

@@ -0,0 +1,19 @@
+- // MIR for `multiple` before RenameReturnPlace
++ // MIR for `multiple` after RenameReturnPlace
+  
+  fn multiple(_1: char) -> char {
+      let mut _0: char;
+      let mut _2: char;
+      let mut _3: char;
+  
+      bb0: {
+          _2 = _1;
+-         _3 = _1;
+-         _0 = _3;
++         _0 = _1;
+          _0 = _2;
+          _2 = const 'b';
+          return;
+      }
+  }
+  

tests/mir-opt/nrvo_miscompile_111005.multiple_return.RenameReturnPlace.diff

@@ -0,0 +1,25 @@
+- // MIR for `multiple_return` before RenameReturnPlace
++ // MIR for `multiple_return` after RenameReturnPlace
+  
+  fn multiple_return(_1: char, _2: bool) -> char {
+      let mut _0: char;
+      let mut _3: char;
+      let mut _4: char;
+  
+      bb0: {
+          switchInt(_2) -> [0: bb1, otherwise: bb2];
+      }
+  
+      bb1: {
+          _3 = _1;
+          _0 = _3;
+          return;
+      }
+  
+      bb2: {
+          _4 = const 'z';
+          _0 = _4;
+          return;
+      }
+  }
+