From c7ef03aeb7c33825e3d020c9e1791171e71cd445 Mon Sep 17 00:00:00 2001 From: Amogh Shivaram Date: Tue, 8 Jul 2025 18:15:05 -0500 Subject: [PATCH] Apply effects to otherwise edge in dataflow analysis --- .../src/drop_flag_effects.rs | 28 +++-- .../src/framework/direction.rs | 11 +- .../rustc_mir_dataflow/src/framework/mod.rs | 1 + .../src/impls/initialized.rs | 88 ++++++++++---- .../src/elaborate_drops.rs | 2 + .../src/remove_uninit_drops.rs | 1 + ...erwise_drops.result_ok.ElaborateDrops.diff | 108 ++++++++++++++++++ tests/mir-opt/otherwise_drops.rs | 13 +++ ...mes.foo.ScalarReplacementOfAggregates.diff | 11 +- 9 files changed, 219 insertions(+), 44 deletions(-) create mode 100644 tests/mir-opt/otherwise_drops.result_ok.ElaborateDrops.diff create mode 100644 tests/mir-opt/otherwise_drops.rs diff --git a/compiler/rustc_mir_dataflow/src/drop_flag_effects.rs b/compiler/rustc_mir_dataflow/src/drop_flag_effects.rs index 496a342cf4c55..c9c7fddae5a87 100644 --- a/compiler/rustc_mir_dataflow/src/drop_flag_effects.rs +++ b/compiler/rustc_mir_dataflow/src/drop_flag_effects.rs @@ -1,5 +1,6 @@ use rustc_abi::VariantIdx; use rustc_middle::mir::{self, Body, Location, Terminator, TerminatorKind}; +use smallvec::SmallVec; use tracing::debug; use super::move_paths::{InitKind, LookupResult, MoveData, MovePathIndex}; @@ -155,15 +156,28 @@ where } } -/// Calls `handle_inactive_variant` for each descendant move path of `enum_place` that contains a -/// `Downcast` to a variant besides the `active_variant`. -/// -/// NOTE: If there are no move paths corresponding to an inactive variant, -/// `handle_inactive_variant` will not be called for that variant. +/// Indicates which variants are inactive at a `SwitchInt` edge by listing their `VariantIdx`s or +/// specifying the single active variant's `VariantIdx`. +pub(crate) enum InactiveVariants { + Inactives(SmallVec<[VariantIdx; 4]>), + Active(VariantIdx), +} + +impl InactiveVariants { + fn contains(&self, variant_idx: VariantIdx) -> bool { + match self { + InactiveVariants::Inactives(inactives) => inactives.contains(&variant_idx), + InactiveVariants::Active(active) => variant_idx != *active, + } + } +} + +/// Calls `handle_inactive_variant` for each child move path of `enum_place` corresponding to an +/// inactive variant at a particular `SwitchInt` edge. pub(crate) fn on_all_inactive_variants<'tcx>( move_data: &MoveData<'tcx>, enum_place: mir::Place<'tcx>, - active_variant: VariantIdx, + inactive_variants: &InactiveVariants, mut handle_inactive_variant: impl FnMut(MovePathIndex), ) { let LookupResult::Exact(enum_mpi) = move_data.rev_lookup.find(enum_place.as_ref()) else { @@ -182,7 +196,7 @@ pub(crate) fn on_all_inactive_variants<'tcx>( unreachable!(); }; - if variant_idx != active_variant { + if inactive_variants.contains(variant_idx) { on_all_children_bits(move_data, variant_mpi, |mpi| handle_inactive_variant(mpi)); } } diff --git a/compiler/rustc_mir_dataflow/src/framework/direction.rs b/compiler/rustc_mir_dataflow/src/framework/direction.rs index e955e38ad10fa..cb647476db8f3 100644 --- a/compiler/rustc_mir_dataflow/src/framework/direction.rs +++ b/compiler/rustc_mir_dataflow/src/framework/direction.rs @@ -112,12 +112,13 @@ impl Direction for Backward { propagate(pred, &tmp); } - mir::TerminatorKind::SwitchInt { targets: _, ref discr } => { + mir::TerminatorKind::SwitchInt { ref targets, ref discr } => { if let Some(mut data) = analysis.get_switch_int_data(block, discr) { let mut tmp = analysis.bottom_value(body); for &value in &body.basic_blocks.switch_sources()[&(block, pred)] { tmp.clone_from(exit_state); - analysis.apply_switch_int_edge_effect(&mut data, &mut tmp, value); + analysis + .apply_switch_int_edge_effect(&mut data, &mut tmp, value, targets); propagate(pred, &tmp); } } else { @@ -290,20 +291,20 @@ impl Direction for Forward { for (value, target) in targets.iter() { tmp.clone_from(exit_state); let value = SwitchTargetValue::Normal(value); - analysis.apply_switch_int_edge_effect(&mut data, &mut tmp, value); + analysis.apply_switch_int_edge_effect(&mut data, &mut tmp, value, targets); propagate(target, &tmp); } // Once we get to the final, "otherwise" branch, there is no need to preserve // `exit_state`, so pass it directly to `apply_switch_int_edge_effect` to save // a clone of the dataflow state. - let otherwise = targets.otherwise(); analysis.apply_switch_int_edge_effect( &mut data, exit_state, SwitchTargetValue::Otherwise, + targets, ); - propagate(otherwise, exit_state); + propagate(targets.otherwise(), exit_state); } else { for target in targets.all_targets() { propagate(*target, exit_state); diff --git a/compiler/rustc_mir_dataflow/src/framework/mod.rs b/compiler/rustc_mir_dataflow/src/framework/mod.rs index 9cadec100b534..b6a5603601959 100644 --- a/compiler/rustc_mir_dataflow/src/framework/mod.rs +++ b/compiler/rustc_mir_dataflow/src/framework/mod.rs @@ -224,6 +224,7 @@ pub trait Analysis<'tcx> { _data: &mut Self::SwitchIntData, _state: &mut Self::Domain, _value: SwitchTargetValue, + _targets: &mir::SwitchTargets, ) { unreachable!(); } diff --git a/compiler/rustc_mir_dataflow/src/impls/initialized.rs b/compiler/rustc_mir_dataflow/src/impls/initialized.rs index 18165b0b9bd08..085757f0fb6ee 100644 --- a/compiler/rustc_mir_dataflow/src/impls/initialized.rs +++ b/compiler/rustc_mir_dataflow/src/impls/initialized.rs @@ -9,9 +9,10 @@ use rustc_middle::mir::{ }; use rustc_middle::ty::util::Discr; use rustc_middle::ty::{self, TyCtxt}; +use smallvec::SmallVec; use tracing::{debug, instrument}; -use crate::drop_flag_effects::DropFlagState; +use crate::drop_flag_effects::{DropFlagState, InactiveVariants}; use crate::move_paths::{HasMoveData, InitIndex, InitKind, LookupResult, MoveData, MovePathIndex}; use crate::{ Analysis, GenKill, MaybeReachable, drop_flag_effects, drop_flag_effects_for_function_entry, @@ -26,6 +27,12 @@ pub struct MaybePlacesSwitchIntData<'tcx> { } impl<'tcx> MaybePlacesSwitchIntData<'tcx> { + /// Creates a `SmallVec` mapping each target in `targets` to its `VariantIdx`. + fn variants(&mut self, targets: &mir::SwitchTargets) -> SmallVec<[VariantIdx; 4]> { + self.index = 0; + targets.all_values().iter().map(|value| self.next_discr(value.get())).collect() + } + // The discriminant order in the `SwitchInt` targets should match the order yielded by // `AdtDef::discriminants`. We rely on this to match each discriminant in the targets to its // corresponding variant in linear time. @@ -131,12 +138,26 @@ pub struct MaybeInitializedPlaces<'a, 'tcx> { tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>, move_data: &'a MoveData<'tcx>, + exclude_inactive_in_otherwise: bool, skip_unreachable_unwind: bool, } impl<'a, 'tcx> MaybeInitializedPlaces<'a, 'tcx> { pub fn new(tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>, move_data: &'a MoveData<'tcx>) -> Self { - MaybeInitializedPlaces { tcx, body, move_data, skip_unreachable_unwind: false } + MaybeInitializedPlaces { + tcx, + body, + move_data, + exclude_inactive_in_otherwise: false, + skip_unreachable_unwind: false, + } + } + + /// Ensures definitely inactive variants are excluded from the set of initialized places for + /// blocks reached through an `otherwise` edge. + pub fn exclude_inactive_in_otherwise(mut self) -> Self { + self.exclude_inactive_in_otherwise = true; + self } pub fn skipping_unreachable_unwind(mut self) -> Self { @@ -208,6 +229,7 @@ pub struct MaybeUninitializedPlaces<'a, 'tcx> { move_data: &'a MoveData<'tcx>, mark_inactive_variants_as_uninit: bool, + include_inactive_in_otherwise: bool, skip_unreachable_unwind: DenseBitSet, } @@ -218,6 +240,7 @@ impl<'a, 'tcx> MaybeUninitializedPlaces<'a, 'tcx> { body, move_data, mark_inactive_variants_as_uninit: false, + include_inactive_in_otherwise: false, skip_unreachable_unwind: DenseBitSet::new_empty(body.basic_blocks.len()), } } @@ -232,6 +255,13 @@ impl<'a, 'tcx> MaybeUninitializedPlaces<'a, 'tcx> { self } + /// Ensures definitely inactive variants are included in the set of uninitialized places for + /// blocks reached through an `otherwise` edge. + pub fn include_inactive_in_otherwise(mut self) -> Self { + self.include_inactive_in_otherwise = true; + self + } + pub fn skipping_unreachable_unwind( mut self, unreachable_unwind: DenseBitSet, @@ -431,17 +461,24 @@ impl<'tcx> Analysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> { data: &mut Self::SwitchIntData, state: &mut Self::Domain, value: SwitchTargetValue, + targets: &mir::SwitchTargets, ) { - if let SwitchTargetValue::Normal(value) = value { - // Kill all move paths that correspond to variants we know to be inactive along this - // particular outgoing edge of a `SwitchInt`. - drop_flag_effects::on_all_inactive_variants( - self.move_data, - data.enum_place, - data.next_discr(value), - |mpi| state.kill(mpi), - ); - } + let inactive_variants = match value { + SwitchTargetValue::Normal(value) => InactiveVariants::Active(data.next_discr(value)), + SwitchTargetValue::Otherwise if self.exclude_inactive_in_otherwise => { + InactiveVariants::Inactives(data.variants(targets)) + } + _ => return, + }; + + // Kill all move paths that correspond to variants we know to be inactive along this + // particular outgoing edge of a `SwitchInt`. + drop_flag_effects::on_all_inactive_variants( + self.move_data, + data.enum_place, + &inactive_variants, + |mpi| state.kill(mpi), + ); } } @@ -544,17 +581,24 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> { data: &mut Self::SwitchIntData, state: &mut Self::Domain, value: SwitchTargetValue, + targets: &mir::SwitchTargets, ) { - if let SwitchTargetValue::Normal(value) = value { - // Mark all move paths that correspond to variants other than this one as maybe - // uninitialized (in reality, they are *definitely* uninitialized). - drop_flag_effects::on_all_inactive_variants( - self.move_data, - data.enum_place, - data.next_discr(value), - |mpi| state.gen_(mpi), - ); - } + let inactive_variants = match value { + SwitchTargetValue::Normal(value) => InactiveVariants::Active(data.next_discr(value)), + SwitchTargetValue::Otherwise if self.include_inactive_in_otherwise => { + InactiveVariants::Inactives(data.variants(targets)) + } + _ => return, + }; + + // Mark all move paths that correspond to variants other than this one as maybe + // uninitialized (in reality, they are *definitely* uninitialized). + drop_flag_effects::on_all_inactive_variants( + self.move_data, + data.enum_place, + &inactive_variants, + |mpi| state.gen_(mpi), + ); } } diff --git a/compiler/rustc_mir_transform/src/elaborate_drops.rs b/compiler/rustc_mir_transform/src/elaborate_drops.rs index 42c8cb0b906e8..b4fa2be1d00e7 100644 --- a/compiler/rustc_mir_transform/src/elaborate_drops.rs +++ b/compiler/rustc_mir_transform/src/elaborate_drops.rs @@ -62,12 +62,14 @@ impl<'tcx> crate::MirPass<'tcx> for ElaborateDrops { let env = MoveDataTypingEnv { move_data, typing_env }; let mut inits = MaybeInitializedPlaces::new(tcx, body, &env.move_data) + .exclude_inactive_in_otherwise() .skipping_unreachable_unwind() .iterate_to_fixpoint(tcx, body, Some("elaborate_drops")) .into_results_cursor(body); let dead_unwinds = compute_dead_unwinds(body, &mut inits); let uninits = MaybeUninitializedPlaces::new(tcx, body, &env.move_data) + .include_inactive_in_otherwise() .mark_inactive_variants_as_uninit() .skipping_unreachable_unwind(dead_unwinds) .iterate_to_fixpoint(tcx, body, Some("elaborate_drops")) diff --git a/compiler/rustc_mir_transform/src/remove_uninit_drops.rs b/compiler/rustc_mir_transform/src/remove_uninit_drops.rs index 9044a88295cc4..25d6fa1bab9fd 100644 --- a/compiler/rustc_mir_transform/src/remove_uninit_drops.rs +++ b/compiler/rustc_mir_transform/src/remove_uninit_drops.rs @@ -22,6 +22,7 @@ impl<'tcx> crate::MirPass<'tcx> for RemoveUninitDrops { let move_data = MoveData::gather_moves(body, tcx, |ty| ty.needs_drop(tcx, typing_env)); let mut maybe_inits = MaybeInitializedPlaces::new(tcx, body, &move_data) + .exclude_inactive_in_otherwise() .iterate_to_fixpoint(tcx, body, Some("remove_uninit_drops")) .into_results_cursor(body); diff --git a/tests/mir-opt/otherwise_drops.result_ok.ElaborateDrops.diff b/tests/mir-opt/otherwise_drops.result_ok.ElaborateDrops.diff new file mode 100644 index 0000000000000..9bd4db723d401 --- /dev/null +++ b/tests/mir-opt/otherwise_drops.result_ok.ElaborateDrops.diff @@ -0,0 +1,108 @@ +- // MIR for `result_ok` before ElaborateDrops ++ // MIR for `result_ok` after ElaborateDrops + + fn result_ok(_1: Result) -> Option { + debug result => _1; + let mut _0: std::option::Option; + let mut _2: isize; + let _3: std::string::String; + let mut _4: std::string::String; ++ let mut _5: bool; ++ let mut _6: isize; ++ let mut _7: isize; + scope 1 { + debug s => _3; + } + + bb0: { ++ _5 = const false; ++ _5 = const true; + PlaceMention(_1); + _2 = discriminant(_1); + switchInt(move _2) -> [0: bb2, otherwise: bb1]; + } + + bb1: { + _0 = Option::::None; + goto -> bb5; + } + + bb2: { + StorageLive(_3); + _3 = move ((_1 as Ok).0: std::string::String); + StorageLive(_4); + _4 = move _3; + _0 = Option::::Some(move _4); +- drop(_4) -> [return: bb3, unwind: bb7]; ++ goto -> bb3; + } + + bb3: { + StorageDead(_4); +- drop(_3) -> [return: bb4, unwind: bb8]; ++ goto -> bb4; + } + + bb4: { + StorageDead(_3); + goto -> bb5; + } + + bb5: { +- drop(_1) -> [return: bb6, unwind: bb9]; ++ goto -> bb16; + } + + bb6: { + return; + } + + bb7 (cleanup): { +- drop(_3) -> [return: bb8, unwind terminate(cleanup)]; ++ goto -> bb8; + } + + bb8 (cleanup): { +- drop(_1) -> [return: bb9, unwind terminate(cleanup)]; ++ goto -> bb9; + } + + bb9 (cleanup): { + resume; ++ } ++ ++ bb10: { ++ goto -> bb6; ++ } ++ ++ bb11 (cleanup): { ++ goto -> bb9; ++ } ++ ++ bb12 (cleanup): { ++ goto -> bb9; ++ } ++ ++ bb13: { ++ goto -> bb10; ++ } ++ ++ bb14: { ++ goto -> bb10; ++ } ++ ++ bb15 (cleanup): { ++ goto -> bb9; ++ } ++ ++ bb16: { ++ _6 = discriminant(_1); ++ switchInt(move _6) -> [0: bb13, otherwise: bb14]; ++ } ++ ++ bb17 (cleanup): { ++ _7 = discriminant(_1); ++ switchInt(move _7) -> [0: bb11, otherwise: bb15]; + } + } + diff --git a/tests/mir-opt/otherwise_drops.rs b/tests/mir-opt/otherwise_drops.rs new file mode 100644 index 0000000000000..c4bc05b593c8b --- /dev/null +++ b/tests/mir-opt/otherwise_drops.rs @@ -0,0 +1,13 @@ +//@ compile-flags: -C panic=abort +//@ test-mir-pass: ElaborateDrops + +// Ensures there are no drops for the wildcard match arm. + +// EMIT_MIR otherwise_drops.result_ok.ElaborateDrops.diff +fn result_ok(result: Result) -> Option { + // CHECK-NOT: drop + match result { + Ok(s) => Some(s), + _ => None, + } +} diff --git a/tests/mir-opt/sroa/lifetimes.foo.ScalarReplacementOfAggregates.diff b/tests/mir-opt/sroa/lifetimes.foo.ScalarReplacementOfAggregates.diff index e2d3c6c41b897..0d5fcf9ef1432 100644 --- a/tests/mir-opt/sroa/lifetimes.foo.ScalarReplacementOfAggregates.diff +++ b/tests/mir-opt/sroa/lifetimes.foo.ScalarReplacementOfAggregates.diff @@ -96,7 +96,6 @@ bb2: { StorageLive(_8); - _28 = const false; _8 = move ((_5 as Ok).0: std::boxed::Box); StorageLive(_9); StorageLive(_10); @@ -186,7 +185,7 @@ bb9: { StorageDead(_6); _29 = discriminant(_5); - switchInt(move _29) -> [0: bb11, otherwise: bb13]; + switchInt(move _29) -> [0: bb10, otherwise: bb11]; } bb10: { @@ -200,14 +199,6 @@ } bb11: { - switchInt(copy _28) -> [0: bb10, otherwise: bb12]; - } - - bb12: { - drop(((_5 as Ok).0: std::boxed::Box)) -> [return: bb10, unwind unreachable]; - } - - bb13: { drop(_5) -> [return: bb10, unwind unreachable]; } }