[mir-opt] Fix mis-optimization and other issues with the SimplifyArmIdentity pass

src/librustc_middle/mir/mod.rs

@@ -256,6 +256,18 @@ impl<'tcx> Body<'tcx> {
         (&mut self.basic_blocks, &mut self.local_decls)
     }
 
+    #[inline]
+    pub fn basic_blocks_local_decls_mut_and_var_debug_info(
+        &mut self,
+    ) -> (
+        &mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+        &mut LocalDecls<'tcx>,
+        &mut Vec<VarDebugInfo<'tcx>>,
+    ) {
+        self.predecessor_cache.invalidate();
+        (&mut self.basic_blocks, &mut self.local_decls, &mut self.var_debug_info)
+    }
+
     /// Returns `true` if a cycle exists in the control-flow graph that is reachable from the
     /// `START_BLOCK`.
     pub fn is_cfg_cyclic(&self) -> bool {

src/librustc_mir/transform/simplify_try.rs

@@ -11,9 +11,10 @@
 
 use crate::transform::{simplify, MirPass, MirSource};
 use itertools::Itertools as _;
-use rustc_index::vec::IndexVec;
+use rustc_index::{bit_set::BitSet, vec::IndexVec};
+use rustc_middle::mir::visit::{NonUseContext, PlaceContext, Visitor};
 use rustc_middle::mir::*;
-use rustc_middle::ty::{Ty, TyCtxt};
+use rustc_middle::ty::{List, Ty, TyCtxt};
 use rustc_target::abi::VariantIdx;
 use std::iter::{Enumerate, Peekable};
 use std::slice::Iter;
@@ -73,9 +74,20 @@ struct ArmIdentityInfo<'tcx> {
 
     /// The statements that should be removed (turned into nops)
     stmts_to_remove: Vec<usize>,
+
+    /// Indices of debug variables that need to be adjusted to point to
+    // `{local_0}.{dbg_projection}`.
+    dbg_info_to_adjust: Vec<usize>,
+
+    /// The projection used to rewrite debug info.
+    dbg_projection: &'tcx List<PlaceElem<'tcx>>,
 }
 
-fn get_arm_identity_info<'a, 'tcx>(stmts: &'a [Statement<'tcx>]) -> Option<ArmIdentityInfo<'tcx>> {
+fn get_arm_identity_info<'a, 'tcx>(
+    stmts: &'a [Statement<'tcx>],
+    locals_count: usize,
+    debug_info: &'a [VarDebugInfo<'tcx>],
+) -> Option<ArmIdentityInfo<'tcx>> {
     // This can't possibly match unless there are at least 3 statements in the block
     // so fail fast on tiny blocks.
     if stmts.len() < 3 {
@@ -187,7 +199,7 @@ fn get_arm_identity_info<'a, 'tcx>(stmts: &'a [Statement<'tcx>]) -> Option<ArmId
     try_eat_storage_stmts(&mut stmt_iter, &mut storage_live_stmts, &mut storage_dead_stmts);
 
     let (get_variant_field_stmt, stmt) = stmt_iter.next()?;
-    let (local_tmp_s0, local_1, vf_s0) = match_get_variant_field(stmt)?;
+    let (local_tmp_s0, local_1, vf_s0, dbg_projection) = match_get_variant_field(stmt)?;
 
     try_eat_storage_stmts(&mut stmt_iter, &mut storage_live_stmts, &mut storage_dead_stmts);
 
@@ -228,6 +240,19 @@ fn get_arm_identity_info<'a, 'tcx>(stmts: &'a [Statement<'tcx>]) -> Option<ArmId
     let stmt_to_overwrite =
         nop_stmts.iter().find(|stmt_idx| live_idx < **stmt_idx && **stmt_idx < dead_idx);
 
+    let mut tmp_assigned_vars = BitSet::new_empty(locals_count);
+    for (l, r) in &tmp_assigns {
+        tmp_assigned_vars.insert(*l);
+        tmp_assigned_vars.insert(*r);
+    }
+
+    let mut dbg_info_to_adjust = Vec::new();
+    for (i, var_info) in debug_info.iter().enumerate() {
+        if tmp_assigned_vars.contains(var_info.place.local) {
+            dbg_info_to_adjust.push(i);
+        }
+    }
+
     Some(ArmIdentityInfo {
         local_temp_0: local_tmp_s0,
         local_1,
@@ -243,12 +268,16 @@ fn get_arm_identity_info<'a, 'tcx>(stmts: &'a [Statement<'tcx>]) -> Option<ArmId
         source_info: discr_stmt_source_info,
         storage_stmts,
         stmts_to_remove: nop_stmts,
+        dbg_info_to_adjust,
+        dbg_projection,
     })
 }
 
 fn optimization_applies<'tcx>(
     opt_info: &ArmIdentityInfo<'tcx>,
     local_decls: &IndexVec<Local, LocalDecl<'tcx>>,
+    local_uses: &IndexVec<Local, usize>,
+    var_debug_info: &[VarDebugInfo<'tcx>],
 ) -> bool {
     trace!("testing if optimization applies...");
 
@@ -273,6 +302,7 @@ fn optimization_applies<'tcx>(
     // Verify the assigment chain consists of the form b = a; c = b; d = c; etc...
     if opt_info.field_tmp_assignments.is_empty() {
         trace!("NO: no assignments found");
+        return false;
     }
     let mut last_assigned_to = opt_info.field_tmp_assignments[0].1;
     let source_local = last_assigned_to;
@@ -285,6 +315,35 @@ fn optimization_applies<'tcx>(
         last_assigned_to = *l;
     }
 
+    // Check that the first and last used locals are only used twice
+    // since they are of the form:
+    //
+    // ```
+    // _first = ((_x as Variant).n: ty);
+    // _n = _first;
+    // ...
+    // ((_y as Variant).n: ty) = _n;
+    // discriminant(_y) = z;
+    // ```
+    for (l, r) in &opt_info.field_tmp_assignments {
+        if local_uses[*l] != 2 {
+            warn!("NO: FAILED assignment chain local {:?} was used more than twice", l);
+            return false;
+        } else if local_uses[*r] != 2 {
+            warn!("NO: FAILED assignment chain local {:?} was used more than twice", r);
+            return false;
+        }
+    }
+
+    // Check that debug info only points to full Locals and not projections.
+    for dbg_idx in &opt_info.dbg_info_to_adjust {
+        let dbg_info = &var_debug_info[*dbg_idx];
+        if !dbg_info.place.projection.is_empty() {
+            trace!("NO: debug info for {:?} had a projection {:?}", dbg_info.name, dbg_info.place);
+            return false;
+        }
+    }
+
     if source_local != opt_info.local_temp_0 {
         trace!(
             "NO: start of assignment chain does not match enum variant temp: {:?} != {:?}",
@@ -312,11 +371,15 @@ impl<'tcx> MirPass<'tcx> for SimplifyArmIdentity {
         }
 
         trace!("running SimplifyArmIdentity on {:?}", source);
-        let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
+        let local_uses = LocalUseCounter::get_local_uses(body);
+        let (basic_blocks, local_decls, debug_info) =
+            body.basic_blocks_local_decls_mut_and_var_debug_info();
         for bb in basic_blocks {
-            if let Some(opt_info) = get_arm_identity_info(&bb.statements) {
+            if let Some(opt_info) =
+                get_arm_identity_info(&bb.statements, local_decls.len(), debug_info)
+            {
                 trace!("got opt_info = {:#?}", opt_info);
-                if !optimization_applies(&opt_info, local_decls) {
+                if !optimization_applies(&opt_info, local_decls, &local_uses, &debug_info) {
                     debug!("optimization skipped for {:?}", source);
                     continue;
                 }
@@ -352,23 +415,57 @@ impl<'tcx> MirPass<'tcx> for SimplifyArmIdentity {
 
                 bb.statements.retain(|stmt| stmt.kind != StatementKind::Nop);
 
+                // Fix the debug info to point to the right local
+                for dbg_index in opt_info.dbg_info_to_adjust {
+                    let dbg_info = &mut debug_info[dbg_index];
+                    assert!(dbg_info.place.projection.is_empty());
+                    dbg_info.place.local = opt_info.local_0;
+                    dbg_info.place.projection = opt_info.dbg_projection;
+                }
+
                 trace!("block is now {:?}", bb.statements);
             }
         }
     }
 }
 
+struct LocalUseCounter {
+    local_uses: IndexVec<Local, usize>,
+}
+
+impl LocalUseCounter {
+    fn get_local_uses<'tcx>(body: &Body<'tcx>) -> IndexVec<Local, usize> {
+        let mut counter = LocalUseCounter { local_uses: IndexVec::from_elem(0, &body.local_decls) };
+        counter.visit_body(body);
+        counter.local_uses
+    }
+}
+
+impl<'tcx> Visitor<'tcx> for LocalUseCounter {
+    fn visit_local(&mut self, local: &Local, context: PlaceContext, _location: Location) {
+        if context.is_storage_marker()
+            || context == PlaceContext::NonUse(NonUseContext::VarDebugInfo)
+        {
+            return;
+        }
+
+        self.local_uses[*local] += 1;
+    }
+}
+
 /// Match on:
 /// ```rust
 /// _LOCAL_INTO = ((_LOCAL_FROM as Variant).FIELD: TY);
 /// ```
-fn match_get_variant_field<'tcx>(stmt: &Statement<'tcx>) -> Option<(Local, Local, VarField<'tcx>)> {
+fn match_get_variant_field<'tcx>(
+    stmt: &Statement<'tcx>,
+) -> Option<(Local, Local, VarField<'tcx>, &'tcx List<PlaceElem<'tcx>>)> {
     match &stmt.kind {
         StatementKind::Assign(box (place_into, rvalue_from)) => match rvalue_from {
             Rvalue::Use(Operand::Copy(pf) | Operand::Move(pf)) => {
                 let local_into = place_into.as_local()?;
                 let (local_from, vf) = match_variant_field_place(*pf)?;
-                Some((local_into, local_from, vf))
+                Some((local_into, local_from, vf, pf.projection))
             }
             _ => None,
         },

src/test/mir-opt/issue-73223.rs

@@ -0,0 +1,13 @@
+fn main() {
+    let split = match Some(1) {
+        Some(v) => v,
+        None => return,
+    };
+
+    let _prev = Some(split);
+    assert_eq!(split, 1);
+}
+
+// EMIT_MIR_FOR_EACH_BIT_WIDTH
+// EMIT_MIR rustc.main.SimplifyArmIdentity.diff
+// EMIT_MIR rustc.main.PreCodegen.diff

src/test/mir-opt/simplify-locals-removes-unused-discriminant-reads/32bit/rustc.map.SimplifyLocals.diff

@@ -5,12 +5,12 @@
       debug x => _1;                       // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:1:8: 1:9
       let mut _0: std::option::Option<std::boxed::Box<()>>; // return place in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:1:31: 1:46
       let mut _2: isize;                   // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:3:9: 3:13
-      let _3: std::boxed::Box<()>;         // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:4:14: 4:15
+-     let _3: std::boxed::Box<()>;         // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:4:14: 4:15
 -     let mut _4: std::boxed::Box<()>;     // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:4:25: 4:26
 -     let mut _5: isize;                   // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:6:1: 6:2
 -     let mut _6: isize;                   // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:6:1: 6:2
       scope 1 {
-          debug x => _3;                   // in scope 1 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:4:14: 4:15
+          debug x => ((_0 as Some).0: std::boxed::Box<()>); // in scope 1 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:4:14: 4:15
       }
 
       bb0: {

src/test/mir-opt/simplify-locals-removes-unused-discriminant-reads/64bit/rustc.map.SimplifyLocals.diff

@@ -5,12 +5,12 @@
       debug x => _1;                       // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:1:8: 1:9
       let mut _0: std::option::Option<std::boxed::Box<()>>; // return place in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:1:31: 1:46
       let mut _2: isize;                   // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:3:9: 3:13
-      let _3: std::boxed::Box<()>;         // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:4:14: 4:15
+-     let _3: std::boxed::Box<()>;         // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:4:14: 4:15
 -     let mut _4: std::boxed::Box<()>;     // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:4:25: 4:26
 -     let mut _5: isize;                   // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:6:1: 6:2
 -     let mut _6: isize;                   // in scope 0 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:6:1: 6:2
       scope 1 {
-          debug x => _3;                   // in scope 1 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:4:14: 4:15
+          debug x => ((_0 as Some).0: std::boxed::Box<()>); // in scope 1 at $DIR/simplify-locals-removes-unused-discriminant-reads.rs:4:14: 4:15
       }
 
       bb0: {

src/test/mir-opt/simplify_try/rustc.try_identity.SimplifyArmIdentity.diff

@@ -15,22 +15,27 @@
       let _10: u32;                        // in scope 0 at $DIR/simplify_try.rs:6:13: 6:15
       let mut _11: u32;                    // in scope 0 at $DIR/simplify_try.rs:7:8: 7:9
       scope 1 {
-          debug y => _2;                   // in scope 1 at $DIR/simplify_try.rs:6:9: 6:10
+-         debug y => _2;                   // in scope 1 at $DIR/simplify_try.rs:6:9: 6:10
++         debug y => ((_0 as Ok).0: u32);  // in scope 1 at $DIR/simplify_try.rs:6:9: 6:10
       }
       scope 2 {
-          debug err => _6;                 // in scope 2 at $DIR/simplify_try.rs:6:14: 6:15
+-         debug err => _6;                 // in scope 2 at $DIR/simplify_try.rs:6:14: 6:15
++         debug err => ((_0 as Err).0: i32); // in scope 2 at $DIR/simplify_try.rs:6:14: 6:15
           scope 3 {
               scope 7 {
-                  debug t => _9;           // in scope 7 at $SRC_DIR/libcore/convert/mod.rs:LL:COL
+-                 debug t => _9;           // in scope 7 at $SRC_DIR/libcore/convert/mod.rs:LL:COL
++                 debug t => ((_0 as Err).0: i32); // in scope 7 at $SRC_DIR/libcore/convert/mod.rs:LL:COL
               }
               scope 8 {
-                  debug v => _8;           // in scope 8 at $SRC_DIR/libcore/result.rs:LL:COL
+-                 debug v => _8;           // in scope 8 at $SRC_DIR/libcore/result.rs:LL:COL
++                 debug v => ((_0 as Err).0: i32); // in scope 8 at $SRC_DIR/libcore/result.rs:LL:COL
                   let mut _12: i32;        // in scope 8 at $DIR/simplify_try.rs:6:14: 6:15
               }
           }
       }
       scope 4 {
-          debug val => _10;                // in scope 4 at $DIR/simplify_try.rs:6:13: 6:15
+-         debug val => _10;                // in scope 4 at $DIR/simplify_try.rs:6:13: 6:15
++         debug val => ((_0 as Ok).0: u32); // in scope 4 at $DIR/simplify_try.rs:6:13: 6:15
           scope 5 {
           }
       }

src/test/mir-opt/simplify_try/rustc.try_identity.SimplifyBranchSame.after.mir

@@ -14,22 +14,22 @@ fn try_identity(_1: std::result::Result<u32, i32>) -> std::result::Result<u32, i
     let _10: u32;                        // in scope 0 at $DIR/simplify_try.rs:6:13: 6:15
     let mut _11: u32;                    // in scope 0 at $DIR/simplify_try.rs:7:8: 7:9
     scope 1 {
-        debug y => _2;                   // in scope 1 at $DIR/simplify_try.rs:6:9: 6:10
+        debug y => ((_0 as Ok).0: u32);  // in scope 1 at $DIR/simplify_try.rs:6:9: 6:10
     }
     scope 2 {
-        debug err => _6;                 // in scope 2 at $DIR/simplify_try.rs:6:14: 6:15
+        debug err => ((_0 as Err).0: i32); // in scope 2 at $DIR/simplify_try.rs:6:14: 6:15
         scope 3 {
             scope 7 {
-                debug t => _9;           // in scope 7 at $SRC_DIR/libcore/convert/mod.rs:LL:COL
+                debug t => ((_0 as Err).0: i32); // in scope 7 at $SRC_DIR/libcore/convert/mod.rs:LL:COL
             }
             scope 8 {
-                debug v => _8;           // in scope 8 at $SRC_DIR/libcore/result.rs:LL:COL
+                debug v => ((_0 as Err).0: i32); // in scope 8 at $SRC_DIR/libcore/result.rs:LL:COL
                 let mut _12: i32;        // in scope 8 at $DIR/simplify_try.rs:6:14: 6:15
             }
         }
     }
     scope 4 {
-        debug val => _10;                // in scope 4 at $DIR/simplify_try.rs:6:13: 6:15
+        debug val => ((_0 as Ok).0: u32); // in scope 4 at $DIR/simplify_try.rs:6:13: 6:15
         scope 5 {
         }
     }

src/test/mir-opt/simplify_try/rustc.try_identity.SimplifyLocals.after.mir

@@ -3,27 +3,22 @@
 fn try_identity(_1: std::result::Result<u32, i32>) -> std::result::Result<u32, i32> {
     debug x => _1;                       // in scope 0 at $DIR/simplify_try.rs:5:17: 5:18
     let mut _0: std::result::Result<u32, i32>; // return place in scope 0 at $DIR/simplify_try.rs:5:41: 5:57
-    let _2: u32;                         // in scope 0 at $DIR/simplify_try.rs:6:9: 6:10
-    let _3: i32;                         // in scope 0 at $DIR/simplify_try.rs:6:14: 6:15
-    let mut _4: i32;                     // in scope 0 at $DIR/simplify_try.rs:6:14: 6:15
-    let mut _5: i32;                     // in scope 0 at $DIR/simplify_try.rs:6:14: 6:15
-    let _6: u32;                         // in scope 0 at $DIR/simplify_try.rs:6:13: 6:15
     scope 1 {
-        debug y => _2;                   // in scope 1 at $DIR/simplify_try.rs:6:9: 6:10
+        debug y => ((_0 as Ok).0: u32);  // in scope 1 at $DIR/simplify_try.rs:6:9: 6:10
     }
     scope 2 {
-        debug err => _3;                 // in scope 2 at $DIR/simplify_try.rs:6:14: 6:15
+        debug err => ((_0 as Err).0: i32); // in scope 2 at $DIR/simplify_try.rs:6:14: 6:15
         scope 3 {
             scope 7 {
-                debug t => _5;           // in scope 7 at $SRC_DIR/libcore/convert/mod.rs:LL:COL
+                debug t => ((_0 as Err).0: i32); // in scope 7 at $SRC_DIR/libcore/convert/mod.rs:LL:COL
             }
             scope 8 {
-                debug v => _4;           // in scope 8 at $SRC_DIR/libcore/result.rs:LL:COL
+                debug v => ((_0 as Err).0: i32); // in scope 8 at $SRC_DIR/libcore/result.rs:LL:COL
             }
         }
     }
     scope 4 {
-        debug val => _6;                 // in scope 4 at $DIR/simplify_try.rs:6:13: 6:15
+        debug val => ((_0 as Ok).0: u32); // in scope 4 at $DIR/simplify_try.rs:6:13: 6:15
         scope 5 {
         }
     }
@@ -32,9 +27,7 @@ fn try_identity(_1: std::result::Result<u32, i32>) -> std::result::Result<u32, i
     }
 
     bb0: {
-        StorageLive(_2);                 // scope 0 at $DIR/simplify_try.rs:6:9: 6:10
         _0 = move _1;                    // scope 1 at $DIR/simplify_try.rs:7:5: 7:10
-        StorageDead(_2);                 // scope 0 at $DIR/simplify_try.rs:8:1: 8:2
         return;                          // scope 0 at $DIR/simplify_try.rs:8:2: 8:2
     }
 }