|
1 | | -use std::iter; |
| 1 | +use std::{iter, usize}; |
2 | 2 |
|
| 3 | +use rustc_const_eval::const_eval::mk_eval_cx_for_const_val; |
| 4 | +use rustc_index::bit_set::BitSet; |
3 | 5 | use rustc_index::IndexSlice; |
4 | 6 | use rustc_middle::mir::patch::MirPatch; |
5 | 7 | use rustc_middle::mir::*; |
| 8 | +use rustc_middle::ty; |
6 | 9 | use rustc_middle::ty::layout::{IntegerExt, TyAndLayout}; |
| 10 | +use rustc_middle::ty::util::Discr; |
7 | 11 | use rustc_middle::ty::{ParamEnv, ScalarInt, Ty, TyCtxt}; |
| 12 | +use rustc_mir_dataflow::impls::{borrowed_locals, MaybeTransitiveLiveLocals}; |
| 13 | +use rustc_mir_dataflow::Analysis; |
8 | 14 | use rustc_target::abi::Integer; |
9 | 15 | use rustc_type_ir::TyKind::*; |
10 | 16 |
|
@@ -48,6 +54,10 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification { |
48 | 54 | should_cleanup = true; |
49 | 55 | continue; |
50 | 56 | } |
| 57 | + if simplify_to_copy(tcx, body, bb_idx, param_env).is_some() { |
| 58 | + should_cleanup = true; |
| 59 | + continue; |
| 60 | + } |
51 | 61 | } |
52 | 62 |
|
53 | 63 | if should_cleanup { |
@@ -515,3 +525,212 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp { |
515 | 525 | } |
516 | 526 | } |
517 | 527 | } |
| 528 | + |
| 529 | +/// This is primarily used to merge these copy statements that simplified the canonical enum clone method by GVN. |
| 530 | +/// The GVN simplified |
| 531 | +/// ```ignore (syntax-highlighting-only) |
| 532 | +/// match a { |
| 533 | +/// Foo::A(x) => Foo::A(*x), |
| 534 | +/// Foo::B => Foo::B |
| 535 | +/// } |
| 536 | +/// ``` |
| 537 | +/// to |
| 538 | +/// ```ignore (syntax-highlighting-only) |
| 539 | +/// match a { |
| 540 | +/// Foo::A(_x) => a, // copy a |
| 541 | +/// Foo::B => Foo::B |
| 542 | +/// } |
| 543 | +/// ``` |
| 544 | +/// This function will simplify into a copy statement. |
| 545 | +fn simplify_to_copy<'tcx>( |
| 546 | + tcx: TyCtxt<'tcx>, |
| 547 | + body: &mut Body<'tcx>, |
| 548 | + switch_bb_idx: BasicBlock, |
| 549 | + param_env: ParamEnv<'tcx>, |
| 550 | +) -> Option<()> { |
| 551 | + // To save compile time, only consider the first BB has a switch terminator. |
| 552 | + if switch_bb_idx != START_BLOCK { |
| 553 | + return None; |
| 554 | + } |
| 555 | + let bbs = &body.basic_blocks; |
| 556 | + // Check if the copy source matches the following pattern. |
| 557 | + // _2 = discriminant(*_1); // "*_1" is the expected the copy source. |
| 558 | + // switchInt(move _2) -> [0: bb3, 1: bb2, otherwise: bb1]; |
| 559 | + let &Statement { |
| 560 | + kind: StatementKind::Assign(box (discr_place, Rvalue::Discriminant(expected_src_place))), |
| 561 | + .. |
| 562 | + } = bbs[switch_bb_idx].statements.last()? |
| 563 | + else { |
| 564 | + return None; |
| 565 | + }; |
| 566 | + let expected_src_ty = expected_src_place.ty(body.local_decls(), tcx); |
| 567 | + if !expected_src_ty.ty.is_enum() || expected_src_ty.variant_index.is_some() { |
| 568 | + return None; |
| 569 | + } |
| 570 | + // To save compile time, only consider the copy source is assigned to the return place. |
| 571 | + let expected_dest_place = Place::return_place(); |
| 572 | + let expected_dest_ty = expected_dest_place.ty(body.local_decls(), tcx); |
| 573 | + if expected_dest_ty.ty != expected_src_ty.ty || expected_dest_ty.variant_index.is_some() { |
| 574 | + return None; |
| 575 | + } |
| 576 | + let targets = match bbs[switch_bb_idx].terminator().kind { |
| 577 | + TerminatorKind::SwitchInt { ref discr, ref targets, .. } |
| 578 | + if discr.place() == Some(discr_place) => |
| 579 | + { |
| 580 | + targets |
| 581 | + } |
| 582 | + _ => return None, |
| 583 | + }; |
| 584 | + // We require that the possible target blocks all be distinct. |
| 585 | + if !targets.is_distinct() { |
| 586 | + return None; |
| 587 | + } |
| 588 | + if !bbs[targets.otherwise()].is_empty_unreachable() { |
| 589 | + return None; |
| 590 | + } |
| 591 | + // Check that destinations are identical, and if not, then don't optimize this block. |
| 592 | + let mut target_iter = targets.iter(); |
| 593 | + let first_terminator_kind = &bbs[target_iter.next().unwrap().1].terminator().kind; |
| 594 | + if !target_iter |
| 595 | + .all(|(_, other_target)| first_terminator_kind == &bbs[other_target].terminator().kind) |
| 596 | + { |
| 597 | + return None; |
| 598 | + } |
| 599 | + |
| 600 | + let borrowed_locals = borrowed_locals(body); |
| 601 | + let mut live = None; |
| 602 | + |
| 603 | + for (index, target_bb) in targets.iter() { |
| 604 | + let stmts = &bbs[target_bb].statements; |
| 605 | + if stmts.is_empty() { |
| 606 | + return None; |
| 607 | + } |
| 608 | + if let [Statement { kind: StatementKind::Assign(box (place, rvalue)), .. }] = |
| 609 | + bbs[target_bb].statements.as_slice() |
| 610 | + { |
| 611 | + let dest_ty = place.ty(body.local_decls(), tcx); |
| 612 | + if dest_ty.ty != expected_src_ty.ty || dest_ty.variant_index.is_some() { |
| 613 | + return None; |
| 614 | + } |
| 615 | + let ty::Adt(def, _) = dest_ty.ty.kind() else { |
| 616 | + return None; |
| 617 | + }; |
| 618 | + if expected_dest_place != *place { |
| 619 | + return None; |
| 620 | + } |
| 621 | + match rvalue { |
| 622 | + // Check if `_3 = const Foo::B` can be transformed to `_3 = copy *_1`. |
| 623 | + Rvalue::Use(Operand::Constant(box constant)) |
| 624 | + if let Const::Val(const_, ty) = constant.const_ => |
| 625 | + { |
| 626 | + let (ecx, op) = |
| 627 | + mk_eval_cx_for_const_val(tcx.at(constant.span), param_env, const_, ty)?; |
| 628 | + let variant = ecx.read_discriminant(&op).ok()?; |
| 629 | + if !def.variants()[variant].fields.is_empty() { |
| 630 | + return None; |
| 631 | + } |
| 632 | + let Discr { val, .. } = ty.discriminant_for_variant(tcx, variant)?; |
| 633 | + if val != index { |
| 634 | + return None; |
| 635 | + } |
| 636 | + } |
| 637 | + Rvalue::Use(Operand::Copy(src_place)) if *src_place == expected_src_place => {} |
| 638 | + // Check if `_3 = Foo::B` can be transformed to `_3 = copy *_1`. |
| 639 | + Rvalue::Aggregate(box AggregateKind::Adt(_, variant_index, _, _, None), fields) |
| 640 | + if fields.is_empty() |
| 641 | + && let Some(Discr { val, .. }) = |
| 642 | + expected_src_ty.ty.discriminant_for_variant(tcx, *variant_index) |
| 643 | + && val == index => {} |
| 644 | + _ => return None, |
| 645 | + } |
| 646 | + } else { |
| 647 | + // If the BB contains more than one statement, we have to check if these statements can be ignored. |
| 648 | + let mut lived_stmts: BitSet<usize> = |
| 649 | + BitSet::new_filled(bbs[target_bb].statements.len()); |
| 650 | + let mut expected_copy_stmt = None; |
| 651 | + for (statement_index, statement) in bbs[target_bb].statements.iter().enumerate().rev() { |
| 652 | + let loc = Location { block: target_bb, statement_index }; |
| 653 | + if let StatementKind::Assign(assign) = &statement.kind { |
| 654 | + if !assign.1.is_safe_to_remove() { |
| 655 | + return None; |
| 656 | + } |
| 657 | + } |
| 658 | + match &statement.kind { |
| 659 | + StatementKind::Assign(box (place, _)) |
| 660 | + | StatementKind::SetDiscriminant { place: box place, .. } |
| 661 | + | StatementKind::Deinit(box place) => { |
| 662 | + if place.is_indirect() || borrowed_locals.contains(place.local) { |
| 663 | + return None; |
| 664 | + } |
| 665 | + let live = live.get_or_insert_with(|| { |
| 666 | + MaybeTransitiveLiveLocals::new(&borrowed_locals) |
| 667 | + .into_engine(tcx, body) |
| 668 | + .iterate_to_fixpoint() |
| 669 | + .into_results_cursor(body) |
| 670 | + }); |
| 671 | + live.seek_before_primary_effect(loc); |
| 672 | + if !live.get().contains(place.local) { |
| 673 | + lived_stmts.remove(statement_index); |
| 674 | + } else if let StatementKind::Assign(box ( |
| 675 | + _, |
| 676 | + Rvalue::Use(Operand::Copy(src_place)), |
| 677 | + )) = statement.kind |
| 678 | + && expected_copy_stmt.is_none() |
| 679 | + && expected_src_place == src_place |
| 680 | + && expected_dest_place == *place |
| 681 | + { |
| 682 | + // There is only one statement that cannot be ignored that can be used as an expected copy statement. |
| 683 | + expected_copy_stmt = Some(statement_index); |
| 684 | + } else { |
| 685 | + return None; |
| 686 | + } |
| 687 | + } |
| 688 | + StatementKind::StorageLive(_) |
| 689 | + | StatementKind::StorageDead(_) |
| 690 | + | StatementKind::Nop => (), |
| 691 | + |
| 692 | + StatementKind::Retag(_, _) |
| 693 | + | StatementKind::Coverage(_) |
| 694 | + | StatementKind::Intrinsic(_) |
| 695 | + | StatementKind::ConstEvalCounter |
| 696 | + | StatementKind::PlaceMention(_) |
| 697 | + | StatementKind::FakeRead(_) |
| 698 | + | StatementKind::AscribeUserType(_, _) => { |
| 699 | + return None; |
| 700 | + } |
| 701 | + } |
| 702 | + } |
| 703 | + let expected_copy_stmt = expected_copy_stmt?; |
| 704 | + // We can ignore the paired StorageLive and StorageDead. |
| 705 | + let mut storage_live_locals: BitSet<Local> = BitSet::new_empty(body.local_decls.len()); |
| 706 | + for stmt_index in lived_stmts.iter() { |
| 707 | + let statement = &bbs[target_bb].statements[stmt_index]; |
| 708 | + match &statement.kind { |
| 709 | + StatementKind::Assign(_) if expected_copy_stmt == stmt_index => {} |
| 710 | + StatementKind::StorageLive(local) |
| 711 | + if *local != expected_dest_place.local |
| 712 | + && storage_live_locals.insert(*local) => {} |
| 713 | + StatementKind::StorageDead(local) |
| 714 | + if *local != expected_dest_place.local |
| 715 | + && storage_live_locals.remove(*local) => {} |
| 716 | + StatementKind::Nop => {} |
| 717 | + _ => return None, |
| 718 | + } |
| 719 | + } |
| 720 | + if !storage_live_locals.is_empty() { |
| 721 | + return None; |
| 722 | + } |
| 723 | + } |
| 724 | + } |
| 725 | + let statement_index = bbs[switch_bb_idx].statements.len(); |
| 726 | + let parent_end = Location { block: switch_bb_idx, statement_index }; |
| 727 | + let mut patch = MirPatch::new(body); |
| 728 | + patch.add_assign( |
| 729 | + parent_end, |
| 730 | + expected_dest_place, |
| 731 | + Rvalue::Use(Operand::Copy(expected_src_place)), |
| 732 | + ); |
| 733 | + patch.patch_terminator(switch_bb_idx, first_terminator_kind.clone()); |
| 734 | + patch.apply(body); |
| 735 | + Some(()) |
| 736 | +} |
0 commit comments