Introduce MIR optimizations for simplifying `x?` on `Result`s.

This optimization depends on inlining for the identity
conversions introduced by the lowering of the `?`.
To take advantage of `SimplifyArmIdentity`, `-Z mir-opt-level=2`
is required because that triggers the inlining MIR optimization.

11 files changed, 432 insertions, 12 deletions

diff --git a/Cargo.lock b/Cargo.lock
index d5d66f89987..4f2df2ee945 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3709,6 +3709,7 @@ dependencies = [
  "arena",
  "either",
  "graphviz",
+ "itertools 0.8.0",
  "log",
  "log_settings",
  "polonius-engine",
diff --git a/src/librustc/hir/mod.rs b/src/librustc/hir/mod.rs
index 17b13dae37f..9d81aa15d33 100644
--- a/src/librustc/hir/mod.rs
+++ b/src/librustc/hir/mod.rs
@@ -2052,7 +2052,7 @@ pub enum TyKind {
     Err,
 }
 
-#[derive(Copy, Clone, RustcEncodable, RustcDecodable, Debug, HashStable)]
+#[derive(Copy, Clone, RustcEncodable, RustcDecodable, Debug, HashStable, PartialEq)]
 pub struct InlineAsmOutput {
     pub constraint: Symbol,
     pub is_rw: bool,
@@ -2062,7 +2062,7 @@ pub struct InlineAsmOutput {
 
 // NOTE(eddyb) This is used within MIR as well, so unlike the rest of the HIR,
 // it needs to be `Clone` and use plain `Vec<T>` instead of `HirVec<T>`.
-#[derive(Clone, RustcEncodable, RustcDecodable, Debug, HashStable)]
+#[derive(Clone, RustcEncodable, RustcDecodable, Debug, HashStable, PartialEq)]
 pub struct InlineAsmInner {
     pub asm: Symbol,
     pub asm_str_style: StrStyle,
diff --git a/src/librustc/mir/interpret/error.rs b/src/librustc/mir/interpret/error.rs
index 5f605756438..4fe82f03b03 100644
--- a/src/librustc/mir/interpret/error.rs
+++ b/src/librustc/mir/interpret/error.rs
@@ -248,7 +248,7 @@ impl<'tcx> From<InterpError<'tcx>> for InterpErrorInfo<'tcx> {
     }
 }
 
-#[derive(Clone, RustcEncodable, RustcDecodable, HashStable)]
+#[derive(Clone, RustcEncodable, RustcDecodable, HashStable, PartialEq)]
 pub enum PanicInfo<O> {
     Panic {
         msg: Symbol,
diff --git a/src/librustc/mir/mod.rs b/src/librustc/mir/mod.rs
index b7d0f538db5..c66075e81b8 100644
--- a/src/librustc/mir/mod.rs
+++ b/src/librustc/mir/mod.rs
@@ -999,7 +999,7 @@ pub struct Terminator<'tcx> {
     pub kind: TerminatorKind<'tcx>,
 }
 
-#[derive(Clone, RustcEncodable, RustcDecodable, HashStable)]
+#[derive(Clone, RustcEncodable, RustcDecodable, HashStable, PartialEq)]
 pub enum TerminatorKind<'tcx> {
     /// Block should have one successor in the graph; we jump there.
     Goto { target: BasicBlock },
@@ -1528,7 +1528,7 @@ impl Statement<'_> {
     }
 }
 
-#[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable, TypeFoldable)]
+#[derive(Clone, Debug, PartialEq, RustcEncodable, RustcDecodable, HashStable, TypeFoldable)]
 pub enum StatementKind<'tcx> {
     /// Write the RHS Rvalue to the LHS Place.
     Assign(Box<(Place<'tcx>, Rvalue<'tcx>)>),
@@ -1594,7 +1594,7 @@ pub enum RetagKind {
 }
 
 /// The `FakeReadCause` describes the type of pattern why a FakeRead statement exists.
-#[derive(Copy, Clone, RustcEncodable, RustcDecodable, Debug, HashStable)]
+#[derive(Copy, Clone, RustcEncodable, RustcDecodable, Debug, HashStable, PartialEq)]
 pub enum FakeReadCause {
     /// Inject a fake read of the borrowed input at the end of each guards
     /// code.
@@ -1636,7 +1636,7 @@ pub enum FakeReadCause {
     ForIndex,
 }
 
-#[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable, TypeFoldable)]
+#[derive(Clone, Debug, PartialEq, RustcEncodable, RustcDecodable, HashStable, TypeFoldable)]
 pub struct InlineAsm<'tcx> {
     pub asm: hir::InlineAsmInner,
     pub outputs: Box<[Place<'tcx>]>,
@@ -2068,7 +2068,7 @@ impl<'tcx> Operand<'tcx> {
 ///////////////////////////////////////////////////////////////////////////
 /// Rvalues
 
-#[derive(Clone, RustcEncodable, RustcDecodable, HashStable)]
+#[derive(Clone, RustcEncodable, RustcDecodable, HashStable, PartialEq)]
 pub enum Rvalue<'tcx> {
     /// x (either a move or copy, depending on type of x)
     Use(Operand<'tcx>),
@@ -2444,7 +2444,7 @@ impl<'tcx> UserTypeProjections {
 /// * `let (x, _): T = ...` -- here, the `projs` vector would contain
 ///   `field[0]` (aka `.0`), indicating that the type of `s` is
 ///   determined by finding the type of the `.0` field from `T`.
-#[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable)]
+#[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable, PartialEq)]
 pub struct UserTypeProjection {
     pub base: UserTypeAnnotationIndex,
     pub projs: Vec<ProjectionKind>,
diff --git a/src/librustc_mir/Cargo.toml b/src/librustc_mir/Cargo.toml
index 9bc6d32b7cb..4afbb4d85d0 100644
--- a/src/librustc_mir/Cargo.toml
+++ b/src/librustc_mir/Cargo.toml
@@ -13,6 +13,7 @@ doctest = false
 arena = { path = "../libarena" }
 either = "1.5.0"
 dot = { path = "../libgraphviz", package = "graphviz" }
+itertools = "0.8"
 log = "0.4"
 log_settings = "0.1.1"
 polonius-engine  = "0.10.0"
diff --git a/src/librustc_mir/lib.rs b/src/librustc_mir/lib.rs
index cbb6408126a..f2707969517 100644
--- a/src/librustc_mir/lib.rs
+++ b/src/librustc_mir/lib.rs
@@ -16,6 +16,7 @@ Rust MIR: a lowered representation of Rust. Also: an experiment!
 #![feature(decl_macro)]
 #![feature(drain_filter)]
 #![feature(exhaustive_patterns)]
+#![feature(iter_order_by)]
 #![cfg_attr(bootstrap, feature(never_type))]
 #![feature(specialization)]
 #![feature(try_trait)]
diff --git a/src/librustc_mir/transform/mod.rs b/src/librustc_mir/transform/mod.rs
index 897e37858a6..2b2b52971ef 100644
--- a/src/librustc_mir/transform/mod.rs
+++ b/src/librustc_mir/transform/mod.rs
@@ -18,6 +18,7 @@ pub mod cleanup_post_borrowck;
 pub mod check_consts;
 pub mod check_unsafety;
 pub mod simplify_branches;
+pub mod simplify_try;
 pub mod simplify;
 pub mod erase_regions;
 pub mod no_landing_pads;
@@ -305,6 +306,9 @@ fn run_optimization_passes<'tcx>(
         &copy_prop::CopyPropagation,
         &simplify_branches::SimplifyBranches::new("after-copy-prop"),
         &remove_noop_landing_pads::RemoveNoopLandingPads,
+        &simplify::SimplifyCfg::new("after-remove-noop-landing-pads"),
+        &simplify_try::SimplifyArmIdentity,
+        &simplify_try::SimplifyBranchSame,
         &simplify::SimplifyCfg::new("final"),
         &simplify::SimplifyLocals,
 
diff --git a/src/librustc_mir/transform/simplify_try.rs b/src/librustc_mir/transform/simplify_try.rs
new file mode 100644
index 00000000000..de5c2ebb571
--- /dev/null
+++ b/src/librustc_mir/transform/simplify_try.rs
@@ -0,0 +1,201 @@
+//! The general point of the optimizations provided here is to simplify something like:
+//!
+//! ```rust
+//! match x {
+//!     Ok(x) => Ok(x),
+//!     Err(x) => Err(x)
+//! }
+//! ```
+//!
+//! into just `x`.
+
+use crate::transform::{MirPass, MirSource, simplify};
+use rustc::ty::{TyCtxt, Ty};
+use rustc::mir::*;
+use rustc_target::abi::VariantIdx;
+use itertools::Itertools as _;
+
+/// Simplifies arms of form `Variant(x) => Variant(x)` to just a move.
+///
+/// This is done by transforming basic blocks where the statements match:
+///
+/// ```rust
+/// _LOCAL_TMP = ((_LOCAL_1 as Variant ).FIELD: TY );
+/// ((_LOCAL_0 as Variant).FIELD: TY) = move _LOCAL_TMP;
+/// discriminant(_LOCAL_0) = VAR_IDX;
+/// ```
+///
+/// into:
+///
+/// ```rust
+/// _LOCAL_0 = move _LOCAL_1
+/// ```
+pub struct SimplifyArmIdentity;
+
+impl<'tcx> MirPass<'tcx> for SimplifyArmIdentity {
+    fn run_pass(&self, _: TyCtxt<'tcx>, _: MirSource<'tcx>, body: &mut Body<'tcx>) {
+        for bb in body.basic_blocks_mut() {
+            // Need 3 statements:
+            let (s0, s1, s2) = match &mut *bb.statements {
+                [s0, s1, s2] => (s0, s1, s2),
+                _ => continue,
+            };
+
+            // Pattern match on the form we want:
+            let (local_tmp_s0, local_1, vf_s0) = match match_get_variant_field(s0) {
+                None => continue,
+                Some(x) => x,
+            };
+            let (local_tmp_s1, local_0, vf_s1) = match match_set_variant_field(s1) {
+                None => continue,
+                Some(x) => x,
+            };
+            if local_tmp_s0 != local_tmp_s1
+                || vf_s0 != vf_s1
+                || Some((local_0, vf_s0.var_idx)) != match_set_discr(s2)
+            {
+                continue;
+            }
+
+            // Right shape; transform!
+            match &mut s0.kind {
+                StatementKind::Assign(box (place, rvalue)) => {
+                    *place = local_0.into();
+                    *rvalue = Rvalue::Use(Operand::Move(local_1.into()));
+                }
+                _ => unreachable!(),
+            }
+            s1.make_nop();
+            s2.make_nop();
+        }
+    }
+}
+
+/// 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>)> {
+    match &stmt.kind {
+        StatementKind::Assign(box (place_into, rvalue_from)) => match rvalue_from {
+            Rvalue::Use(Operand::Copy(pf)) | Rvalue::Use(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))
+            }
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
+/// Match on:
+/// ```rust
+/// ((_LOCAL_FROM as Variant).FIELD: TY) = move _LOCAL_INTO;
+/// ```
+fn match_set_variant_field<'tcx>(stmt: &Statement<'tcx>) -> Option<(Local, Local, VarField<'tcx>)> {
+    match &stmt.kind {
+        StatementKind::Assign(box (place_from, rvalue_into)) => match rvalue_into {
+            Rvalue::Use(Operand::Move(place_into)) => {
+                let local_into = place_into.as_local()?;
+                let (local_from, vf) = match_variant_field_place(&place_from)?;
+                Some((local_into, local_from, vf))
+            }
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
+/// Match on:
+/// ```rust
+/// discriminant(_LOCAL_TO_SET) = VAR_IDX;
+/// ```
+fn match_set_discr<'tcx>(stmt: &Statement<'tcx>) -> Option<(Local, VariantIdx)> {
+    match &stmt.kind {
+        StatementKind::SetDiscriminant { place, variant_index } => Some((
+            place.as_local()?,
+            *variant_index
+        )),
+        _ => None,
+    }
+}
+
+#[derive(PartialEq)]
+struct VarField<'tcx> {
+    field: Field,
+    field_ty: Ty<'tcx>,
+    var_idx: VariantIdx,
+}
+
+/// Match on `((_LOCAL as Variant).FIELD: TY)`.
+fn match_variant_field_place<'tcx>(place: &Place<'tcx>) -> Option<(Local, VarField<'tcx>)> {
+    match place.as_ref() {
+        PlaceRef {
+            base: &PlaceBase::Local(local),
+            projection: &[ProjectionElem::Downcast(_, var_idx), ProjectionElem::Field(field, ty)],
+        } => Some((local, VarField { field, field_ty: ty, var_idx })),
+        _ => None,
+    }
+}
+
+/// Simplifies `SwitchInt(_) -> [targets]`,
+/// where all the `targets` have the same form,
+/// into `goto -> target_first`.
+pub struct SimplifyBranchSame;
+
+impl<'tcx> MirPass<'tcx> for SimplifyBranchSame {
+    fn run_pass(&self, _: TyCtxt<'tcx>, _: MirSource<'tcx>, body: &mut Body<'tcx>) {
+        let mut did_remove_blocks = false;
+        let bbs = body.basic_blocks_mut();
+        for bb_idx in bbs.indices() {
+            let targets = match &bbs[bb_idx].terminator().kind {
+                TerminatorKind::SwitchInt { targets, .. } => targets,
+                _ => continue,
+            };
+
+            let mut iter_bbs_reachable = targets
+                .iter()
+                .map(|idx| (*idx, &bbs[*idx]))
+                .filter(|(_, bb)| {
+                    // Reaching `unreachable` is UB so assume it doesn't happen.
+                    bb.terminator().kind != TerminatorKind::Unreachable
+                    // But `asm!(...)` could abort the program,
+                    // so we cannot assume that the `unreachable` terminator itself is reachable.
+                    // FIXME(Centril): use a normalization pass instead of a check.
+                    || bb.statements.iter().any(|stmt| match stmt.kind {
+                        StatementKind::InlineAsm(..) => true,
+                        _ => false,
+                    })
+                })
+                .peekable();
+
+            // We want to `goto -> bb_first`.
+            let bb_first = iter_bbs_reachable
+                .peek()
+                .map(|(idx, _)| *idx)
+                .unwrap_or(targets[0]);
+
+            // All successor basic blocks should have the exact same form.
+            let all_successors_equivalent = iter_bbs_reachable
+                .map(|(_, bb)| bb)
+                .tuple_windows()
+                .all(|(bb_l, bb_r)| {
+                    bb_l.is_cleanup == bb_r.is_cleanup
+                    && bb_l.terminator().kind == bb_r.terminator().kind
+                    && bb_l.statements.iter().eq_by(&bb_r.statements, |x, y| x.kind == y.kind)
+                });
+
+            if all_successors_equivalent {
+                // Replace `SwitchInt(..) -> [bb_first, ..];` with a `goto -> bb_first;`.
+                bbs[bb_idx].terminator_mut().kind = TerminatorKind::Goto { target: bb_first };
+                did_remove_blocks = true;
+            }
+        }
+
+        if did_remove_blocks {
+            // We have dead blocks now, so remove those.
+            simplify::remove_dead_blocks(body);
+        }
+    }
+}
diff --git a/src/test/codegen/match.rs b/src/test/codegen/match.rs
index 145d4ba6b4c..b5ee5ba2394 100644
--- a/src/test/codegen/match.rs
+++ b/src/test/codegen/match.rs
@@ -9,19 +9,21 @@ pub enum E {
 
 // CHECK-LABEL: @exhaustive_match
 #[no_mangle]
-pub fn exhaustive_match(e: E, unit: ()) {
+pub fn exhaustive_match(e: E) -> u8 {
 // CHECK: switch{{.*}}, label %[[OTHERWISE:[a-zA-Z0-9_]+]] [
 // CHECK-NEXT: i[[TY:[0-9]+]] [[DISCR:[0-9]+]], label %[[A:[a-zA-Z0-9_]+]]
 // CHECK-NEXT: i[[TY:[0-9]+]] [[DISCR:[0-9]+]], label %[[B:[a-zA-Z0-9_]+]]
 // CHECK-NEXT: ]
 // CHECK: [[B]]:
+// CHECK-NEXT: store i8 1, i8* %1, align 1
 // CHECK-NEXT: br label %[[EXIT:[a-zA-Z0-9_]+]]
 // CHECK: [[OTHERWISE]]:
 // CHECK-NEXT: unreachable
 // CHECK: [[A]]:
+// CHECK-NEXT: store i8 0, i8* %1, align 1
 // CHECK-NEXT: br label %[[EXIT:[a-zA-Z0-9_]+]]
     match e {
-        E::A => unit,
-        E::B => unit,
+        E::A => 0,
+        E::B => 1,
     }
 }
diff --git a/src/test/codegen/try_identity.rs b/src/test/codegen/try_identity.rs
new file mode 100644
index 00000000000..30e7adfddf7
--- /dev/null
+++ b/src/test/codegen/try_identity.rs
@@ -0,0 +1,17 @@
+// compile-flags: -C no-prepopulate-passes -Z mir-opt-level=2
+
+// Ensure that `x?` has no overhead on `Result<T, E>` due to identity `match`es in lowering.
+// This requires inlining to trigger the MIR optimizations in `SimplifyArmIdentity`.
+
+#![crate_type = "lib"]
+
+type R = Result<u64, i32>;
+
+#[no_mangle]
+fn try_identity(x: R) -> R {
+// CHECK: start:
+// CHECK-NOT: br {{.*}}
+// CHECK ret void
+    let y = x?;
+    Ok(y)
+}
diff --git a/src/test/mir-opt/simplify_try.rs b/src/test/mir-opt/simplify_try.rs
new file mode 100644
index 00000000000..7911fbd0a98
--- /dev/null
+++ b/src/test/mir-opt/simplify_try.rs
@@ -0,0 +1,193 @@
+fn try_identity(x: Result<u32, i32>) -> Result<u32, i32> {
+    let y = x?;
+    Ok(y)
+}
+
+fn main() {
+    let _ = try_identity(Ok(0));
+}
+
+// END RUST SOURCE
+// START rustc.try_identity.SimplifyArmIdentity.before.mir
+// fn try_identity(_1: std::result::Result<u32, i32>) -> std::result::Result<u32, i32> {
+//     let mut _0: std::result::Result<u32, i32>;
+//     let _2: u32;
+//     let mut _3: std::result::Result<u32, i32>;
+//     let mut _4: std::result::Result<u32, i32>;
+//     let mut _5: isize;
+//     let _6: i32;
+//     let mut _7: !;
+//     let mut _8: i32;
+//     let mut _9: i32;
+//     let _10: u32;
+//     let mut _11: u32;
+//     scope 1 {
+//     }
+//     scope 2 {
+//         scope 3 {
+//             scope 7 {
+//             }
+//             scope 8 {
+//                 let mut _12: i32;
+//             }
+//         }
+//     }
+//     scope 4 {
+//         scope 5 {
+//         }
+//     }
+//     scope 6 {
+//     }
+//     bb0: {
+//         _5 = discriminant(_1);
+//         switchInt(move _5) -> [0isize: bb4, 1isize: bb2, otherwise: bb1];
+//     }
+//     bb1: {
+//         unreachable;
+//     }
+//     bb2: {
+//         _6 = ((_1 as Err).0: i32);
+//         ((_0 as Err).0: i32) = move _6;
+//         discriminant(_0) = 1;
+//         goto -> bb3;
+//     }
+//     bb3: {
+//         return;
+//     }
+//     bb4: {
+//         _10 = ((_1 as Ok).0: u32);
+//         ((_0 as Ok).0: u32) = move _10;
+//         discriminant(_0) = 0;
+//         goto -> bb3;
+//     }
+// }
+// END rustc.try_identity.SimplifyArmIdentity.before.mir
+
+// START rustc.try_identity.SimplifyArmIdentity.after.mir
+// fn try_identity(_1: std::result::Result<u32, i32>) -> std::result::Result<u32, i32> {
+//     let mut _0: std::result::Result<u32, i32>;
+//     let _2: u32;
+//     let mut _3: std::result::Result<u32, i32>;
+//     let mut _4: std::result::Result<u32, i32>;
+//     let mut _5: isize;
+//     let _6: i32;
+//     let mut _7: !;
+//     let mut _8: i32;
+//     let mut _9: i32;
+//     let _10: u32;
+//     let mut _11: u32;
+//     scope 1 {
+//     }
+//     scope 2 {
+//         scope 3 {
+//             scope 7 {
+//             }
+//             scope 8 {
+//                 let mut _12: i32;
+//             }
+//         }
+//     }
+//     scope 4 {
+//         scope 5 {
+//         }
+//     }
+//     scope 6 {
+//     }
+//     bb0: {
+//         _5 = discriminant(_1);
+//         switchInt(move _5) -> [0isize: bb4, 1isize: bb2, otherwise: bb1];
+//     }
+//     bb1: {
+//         unreachable;
+//     }
+//     bb2: {
+//         _0 = move _1;
+//         nop;
+//         nop;
+//         goto -> bb3;
+//     }
+//     bb3: {
+//         return;
+//     }
+//     bb4: {
+//         _0 = move _1;
+//         nop;
+//         nop;
+//         goto -> bb3;
+//     }
+// }
+// END rustc.try_identity.SimplifyArmIdentity.after.mir
+
+// START rustc.try_identity.SimplifyBranchSame.after.mir
+// fn try_identity(_1: std::result::Result<u32, i32>) -> std::result::Result<u32, i32> {
+//     let mut _0: std::result::Result<u32, i32>;
+//     let _2: u32;
+//     let mut _3: std::result::Result<u32, i32>;
+//     let mut _4: std::result::Result<u32, i32>;
+//     let mut _5: isize;
+//     let _6: i32;
+//     let mut _7: !;
+//     let mut _8: i32;
+//     let mut _9: i32;
+//     let _10: u32;
+//     let mut _11: u32;
+//     scope 1 {
+//     }
+//     scope 2 {
+//         scope 3 {
+//             scope 7 {
+//             }
+//             scope 8 {
+//                 let mut _12: i32;
+//             }
+//         }
+//     }
+//     scope 4 {
+//         scope 5 {
+//         }
+//     }
+//     scope 6 {
+//     }
+//     bb0: {
+//         _5 = discriminant(_1);
+//         goto -> bb2;
+//     }
+//     bb1: {
+//         return;
+//     }
+//     bb2: {
+//         _0 = move _1;
+//         nop;
+//         nop;
+//         goto -> bb1;
+//     }
+// }
+// END rustc.try_identity.SimplifyBranchSame.after.mir
+
+// START rustc.try_identity.SimplifyLocals.after.mir
+// fn try_identity(_1: std::result::Result<u32, i32>) -> std::result::Result<u32, i32> {
+//     let mut _0: std::result::Result<u32, i32>;
+//     let mut _2: isize;
+//     scope 1 {
+//     }
+//     scope 2 {
+//         scope 3 {
+//             scope 7 {
+//             }
+//             scope 8 {
+//             }
+//         }
+//     }
+//     scope 4 {
+//         scope 5 {
+//         }
+//     }
+//     scope 6 {
+//     }
+//     bb0: {
+//         _2 = discriminant(_1);
+//         _0 = move _1;
+//         return;
+//     }
+// }
+// END rustc.try_identity.SimplifyLocals.after.mir