Auto merge of #66282 - Centril:simplify-try, r=<try>

[WIP] [mir-opt] asking `?`s in a more optimized fashion

This PR works towards #66234 by providing two optimization passes meant to run in sequence:

- `SimplifyArmIdentity` which transforms something like:
  ```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
  ```

- `SimplifyBranchSame` which transforms `SwitchInt`s to identical basic blocks into a `goto` to the first reachable target.

Together, these are meant to simplify the following into just `res`:
```rust
match res {
    Ok(x) => Ok(x),
    Err(x) => Err(x),
}
```

It should be noted however that the desugaring of `?` includes a function call and so the first pass in this PR relies on inlining to substitute that function call for identity on `x`. Inlining requires `mir-opt-level=2` so this might not have any effect in perf-bot but let's find out.

r? @oli-obk -- This is WIP, but I'd appreciate feedback. :)

Author: bors

Cargo.lock

@@ -3668,6 +3668,7 @@ dependencies = [
  "arena",
  "either",
  "graphviz",
+ "itertools 0.8.0",
  "log",
  "log_settings",
  "polonius-engine",

src/librustc/hir/mod.rs

@@ -2101,7 +2101,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,
@@ -2111,7 +2111,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 InlineAsm {
     pub asm: Symbol,
     pub asm_str_style: StrStyle,

src/librustc/mir/interpret/error.rs

@@ -246,7 +246,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,

src/librustc/mir/mod.rs

@@ -1040,7 +1040,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 },
@@ -1568,7 +1568,7 @@ impl Statement<'_> {
     }
 }
 
-#[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable)]
+#[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable, PartialEq)]
 pub enum StatementKind<'tcx> {
     /// Write the RHS Rvalue to the LHS Place.
     Assign(Box<(Place<'tcx>, Rvalue<'tcx>)>),
@@ -1634,7 +1634,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.
@@ -1667,7 +1667,7 @@ pub enum FakeReadCause {
     ForLet,
 }
 
-#[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable)]
+#[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable, PartialEq)]
 pub struct InlineAsm<'tcx> {
     pub asm: HirInlineAsm,
     pub outputs: Box<[Place<'tcx>]>,
@@ -2105,7 +2105,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>),
@@ -2487,7 +2487,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>,

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"

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)]
 #![feature(never_type)]
 #![feature(specialization)]
 #![feature(try_trait)]

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;
@@ -275,6 +276,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-noo-landing-pads"),
+        &simplify_try::SimplifyArmIdentity,
+        &simplify_try::SimplifyBranchSame,
         &simplify::SimplifyCfg::new("final"),
         &simplify::SimplifyLocals,
 

src/librustc_mir/transform/simplify_try.rs

@@ -0,0 +1,276 @@
+//! 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 bbs = body.basic_blocks_mut();
+        for bb_idx in bbs.indices() {
+            let targets = match &bbs[bb_idx].terminator().kind {
+                TerminatorKind::SwitchInt { targets, .. } => targets,
+                _ => continue,
+            };
+
+            // Reaching `unreachable` is UB so assume it doesn't happen.
+            let mut iter_bbs_reachable = targets
+                .iter()
+                .map(|idx| (*idx, &bbs[*idx]))
+                .filter(|(_, bb)| bb.terminator().kind != TerminatorKind::Unreachable)
+                .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 };
+            }
+        }
+
+        // We may have dead blocks now, so remvoe those.
+        simplify::remove_dead_blocks(body);
+    }
+}
+
+/*
+
+KEEPSAKE: REMOVE IF NOT NECESSARY!
+
+fn statement_semantic_eq(sa: &StatementKind<'_>, sb: &StatementKind<'_>) -> bool {
+    use StatementKind::*;
+    match (sb, sa) {
+        (AscribeUserType(pa, va), AscribeUserType(pb, vb)) => pa == pb && va == vb,
+        (Assign(a), Assign(b)) => a == b,
+        (FakeRead(fa, pa), FakeRead(fb, pb)) => fa == fb && pa == pb,
+        (InlineAsm(a), InlineAsm(b)) => a == b,
+        (Nop, StatementKind::Nop) => true,
+        (Retag(ra, pa), Retag(rb, pb)) => ra == rb && pa == pb,
+        (
+            SetDiscriminant { place: pa, variant_index: va },
+            SetDiscriminant { place: pb, variant_index: vb },
+        ) => pa == pb && va == vb,
+        (StorageDead(a), StorageDead(b)) => a == b,
+        (StorageLive(a), StorageLive(b)) => a == b,
+        (AscribeUserType(..), _) | (_, AscribeUserType(..))
+        | (StorageDead(..), _) | (_, StorageDead(..))
+        | (Assign(..), _) | (_, Assign(..))
+        | (FakeRead(..), _) | (_, FakeRead(..))
+        | (InlineAsm(..), _) | (_, InlineAsm(..))
+        | (Nop, _) | (_, Nop)
+        | (Retag(..), _) | (_, Retag(..))
+        | (SetDiscriminant { .. }, _) | (_, SetDiscriminant { .. }) => true,
+    }
+}
+
+fn terminator_semantic_eq(ta: &TerminatorKind<'_>, tb: &TerminatorKind<'_>) -> bool {
+    use TerminatorKind::*;
+    match (ta, tb) {
+        (Goto { target: a }, Goto { target: b }) => a == b,
+        (Resume, Resume)
+        | (Abort, Abort)
+        | (Return, Return)
+        | (Unreachable, Unreachable)
+        | (GeneratorDrop, GeneratorDrop) => true,
+        (
+            SwitchInt { discr: da, switch_ty: sa, targets: ta, values: _ },
+            SwitchInt { discr: db, switch_ty: sb, targets: tb, values: _ },
+        ) => da == db && sa == sb && ta == tb,
+        (
+            Drop { location: la, target: ta, unwind: ua },
+            Drop { location: lb, target: tb, unwind: ub },
+        ) => la == lb && ta == tb && ua == ub,
+        (
+            DropAndReplace { location: la, target: ta, unwind: ua, value: va },
+            DropAndReplace { location: lb, target: tb, unwind: ub, value: vb },
+        ) => la == lb && ta == tb && ua == ub && va == vb,
+        (
+            Call { func: fa, args: aa, destination: da, cleanup: ca, from_hir_call: _ },
+            Call { func: fb, args: ab, destination: db, cleanup: cb, from_hir_call: _ },
+        ) => fa == fb && aa == ab && da == db && ca == cb,
+        (
+            Assert { cond: coa, expected: ea, msg: ma, target: ta, cleanup: cla },
+            Assert { cond: cob, expected: eb, msg: mb, target: tb, cleanup: clb },
+        ) => coa == cob && ea == eb && ma == mb && ta == tb && cla == clb,
+        (
+            Yield { value: va, resume: ra, drop: da },
+            Yield { value: vb, resume: rb, drop: db },
+        ) => va == vb && ra == rb && da == db,
+        (
+            FalseEdges { real_target: ra, imaginary_target: ia },
+            FalseEdges { real_target: rb, imaginary_target: ib },
+        ) => ra == rb && ia == ib,
+        (
+            FalseUnwind { real_target: ra, unwind: ua },
+            FalseUnwind { real_target: rb, unwind: ub },
+        ) => ra == rb && ua == ub,
+        (Goto { .. }, _) | (_, Goto { .. })
+        | (Resume, _) | (_, Resume)
+        | (Abort, _) | (_, Abort)
+        | (Return, _) | (_, Return)
+        | (Unreachable, _) | (_, Unreachable)
+        | (GeneratorDrop, _) | (_, GeneratorDrop)
+        | (SwitchInt { .. }, _) | (_, SwitchInt { .. })
+        | (Drop { .. }, _) | (_, Drop { .. })
+        | (DropAndReplace { .. }, _) | (_, DropAndReplace { .. })
+        | (Call { .. }, _) | (_, Call { .. })
+        | (Assert { .. }, _) | (_, Assert { .. })
+        | (Yield { .. }, _) | (_, Yield { .. })
+        | (FalseEdges { .. }, _) | (_, FalseEdges { .. }) => false,
+    }
+}
+*/

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