[rustc.git] / compiler / rustc_mir_transform / src / separate_const_switch.rs

//! A pass that duplicates switch-terminated blocks
//! into a new copy for each predecessor, provided
//! the predecessor sets the value being switched
//! over to a constant.
//!
//! The purpose of this pass is to help constant
//! propagation passes to simplify the switch terminator
//! of the copied blocks into gotos when some predecessors
//! statically determine the output of switches.
//!
//! ```text
//!     x = 12 ---              ---> something
//!               \            / 12
//!                --> switch x
//!               /            \ otherwise
//!     x = y  ---              ---> something else
//! ```
//! becomes
//! ```text
//!     x = 12 ---> switch x ------> something
//!                          \ / 12
//!                           X
//!                          / \ otherwise
//!     x = y  ---> switch x ------> something else
//! ```
//! so it can hopefully later be turned by another pass into
//! ```text
//!     x = 12 --------------------> something
//!                            / 12
//!                           /
//!                          /   otherwise
//!     x = y  ---- switch x ------> something else
//! ```
//!
//! This optimization is meant to cover simple cases
//! like `?` desugaring. For now, it thus focuses on
//! simplicity rather than completeness (it notably
//! sometimes duplicates abusively).

use crate::MirPass;
use rustc_middle::mir::*;
use rustc_middle::ty::TyCtxt;
use smallvec::SmallVec;

pub struct SeparateConstSwitch;

impl<'tcx> MirPass<'tcx> for SeparateConstSwitch {
    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
        sess.mir_opt_level() >= 4
    }

    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        // If execution did something, applying a simplification layer
        // helps later passes optimize the copy away.
        if separate_const_switch(body) > 0 {
            super::simplify::simplify_cfg(tcx, body);
        }
    }
}

/// Returns the amount of blocks that were duplicated
pub fn separate_const_switch(body: &mut Body<'_>) -> usize {
    let mut new_blocks: SmallVec<[(BasicBlock, BasicBlock); 6]> = SmallVec::new();
    let predecessors = body.predecessors();
    'block_iter: for (block_id, block) in body.basic_blocks().iter_enumerated() {
        if let TerminatorKind::SwitchInt {
            discr: Operand::Copy(switch_place) | Operand::Move(switch_place),
            ..
        } = block.terminator().kind
        {
            // If the block is on an unwind path, do not
            // apply the optimization as unwind paths
            // rely on a unique parent invariant
            if block.is_cleanup {
                continue 'block_iter;
            }

            // If the block has fewer than 2 predecessors, ignore it
            // we could maybe chain blocks that have exactly one
            // predecessor, but for now we ignore that
            if predecessors[block_id].len() < 2 {
                continue 'block_iter;
            }

            // First, let's find a non-const place
            // that determines the result of the switch
            if let Some(switch_place) = find_determining_place(switch_place, block) {
                // We now have an input place for which it would
                // be interesting if predecessors assigned it from a const

                let mut predecessors_left = predecessors[block_id].len();
                'predec_iter: for predecessor_id in predecessors[block_id].iter().copied() {
                    let predecessor = &body.basic_blocks()[predecessor_id];

                    // First we make sure the predecessor jumps
                    // in a reasonable way
                    match &predecessor.terminator().kind {
                        // The following terminators are
                        // unconditionally valid
                        TerminatorKind::Goto { .. } | TerminatorKind::SwitchInt { .. } => {}

                        TerminatorKind::FalseEdge { real_target, .. } => {
                            if *real_target != block_id {
                                continue 'predec_iter;
                            }
                        }

                        // The following terminators are not allowed
                        TerminatorKind::Resume
                        | TerminatorKind::Drop { .. }
                        | TerminatorKind::DropAndReplace { .. }
                        | TerminatorKind::Call { .. }
                        | TerminatorKind::Assert { .. }
                        | TerminatorKind::FalseUnwind { .. }
                        | TerminatorKind::Yield { .. }
                        | TerminatorKind::Abort
                        | TerminatorKind::Return
                        | TerminatorKind::Unreachable
                        | TerminatorKind::InlineAsm { .. }
                        | TerminatorKind::GeneratorDrop => {
                            continue 'predec_iter;
                        }
                    }

                    if is_likely_const(switch_place, predecessor) {
                        new_blocks.push((predecessor_id, block_id));
                        predecessors_left -= 1;
                        if predecessors_left < 2 {
                            // If the original block only has one predecessor left,
                            // we have nothing left to do
                            break 'predec_iter;
                        }
                    }
                }
            }
        }
    }

    // Once the analysis is done, perform the duplication
    let body_span = body.span;
    let copied_blocks = new_blocks.len();
    let blocks = body.basic_blocks_mut();
    for (pred_id, target_id) in new_blocks {
        let new_block = blocks[target_id].clone();
        let new_block_id = blocks.push(new_block);
        let terminator = blocks[pred_id].terminator_mut();

        match terminator.kind {
            TerminatorKind::Goto { ref mut target } => {
                *target = new_block_id;
            }

            TerminatorKind::FalseEdge { ref mut real_target, .. } => {
                if *real_target == target_id {
                    *real_target = new_block_id;
                }
            }

            TerminatorKind::SwitchInt { ref mut targets, .. } => {
                targets.all_targets_mut().iter_mut().for_each(|x| {
                    if *x == target_id {
                        *x = new_block_id;
                    }
                });
            }

            TerminatorKind::Resume
            | TerminatorKind::Abort
            | TerminatorKind::Return
            | TerminatorKind::Unreachable
            | TerminatorKind::GeneratorDrop
            | TerminatorKind::Assert { .. }
            | TerminatorKind::DropAndReplace { .. }
            | TerminatorKind::FalseUnwind { .. }
            | TerminatorKind::Drop { .. }
            | TerminatorKind::Call { .. }
            | TerminatorKind::InlineAsm { .. }
            | TerminatorKind::Yield { .. } => {
                span_bug!(
                    body_span,
                    "basic block terminator had unexpected kind {:?}",
                    &terminator.kind
                )
            }
        }
    }

    copied_blocks
}

/// This function describes a rough heuristic guessing
/// whether a place is last set with a const within the block.
/// Notably, it will be overly pessimistic in cases that are already
/// not handled by `separate_const_switch`.
fn is_likely_const<'tcx>(mut tracked_place: Place<'tcx>, block: &BasicBlockData<'tcx>) -> bool {
    for statement in block.statements.iter().rev() {
        match &statement.kind {
            StatementKind::Assign(assign) => {
                if assign.0 == tracked_place {
                    match assign.1 {
                        // These rvalues are definitely constant
                        Rvalue::Use(Operand::Constant(_))
                        | Rvalue::Ref(_, _, _)
                        | Rvalue::AddressOf(_, _)
                        | Rvalue::Cast(_, Operand::Constant(_), _)
                        | Rvalue::NullaryOp(_, _)
                        | Rvalue::ShallowInitBox(_, _)
                        | Rvalue::UnaryOp(_, Operand::Constant(_)) => return true,

                        // These rvalues make things ambiguous
                        Rvalue::Repeat(_, _)
                        | Rvalue::ThreadLocalRef(_)
                        | Rvalue::Len(_)
                        | Rvalue::BinaryOp(_, _)
                        | Rvalue::CheckedBinaryOp(_, _)
                        | Rvalue::Aggregate(_, _) => return false,

                        // These rvalues move the place to track
                        Rvalue::Cast(_, Operand::Copy(place) | Operand::Move(place), _)
                        | Rvalue::Use(Operand::Copy(place) | Operand::Move(place))
                        | Rvalue::UnaryOp(_, Operand::Copy(place) | Operand::Move(place))
                        | Rvalue::Discriminant(place) => tracked_place = place,
                    }
                }
            }

            // If the discriminant is set, it is always set
            // as a constant, so the job is done.
            // As we are **ignoring projections**, if the place
            // we are tracking sees its discriminant be set,
            // that means we had to be tracking the discriminant
            // specifically (as it is impossible to switch over
            // an enum directly, and if we were switching over
            // its content, we would have had to at least cast it to
            // some variant first)
            StatementKind::SetDiscriminant { place, .. } => {
                if **place == tracked_place {
                    return true;
                }
            }

            // These statements have no influence on the place
            // we are interested in
            StatementKind::FakeRead(_)
            | StatementKind::StorageLive(_)
            | StatementKind::Retag(_, _)
            | StatementKind::AscribeUserType(_, _)
            | StatementKind::Coverage(_)
            | StatementKind::StorageDead(_)
            | StatementKind::CopyNonOverlapping(_)
            | StatementKind::Nop => {}
        }
    }

    // If no good reason for the place to be const is found,
    // give up. We could maybe go up predecessors, but in
    // most cases giving up now should be sufficient.
    false
}

/// Finds a unique place that entirely determines the value
/// of `switch_place`, if it exists. This is only a heuristic.
/// Ideally we would like to track multiple determining places
/// for some edge cases, but one is enough for a lot of situations.
fn find_determining_place<'tcx>(
    mut switch_place: Place<'tcx>,
    block: &BasicBlockData<'tcx>,
) -> Option<Place<'tcx>> {
    for statement in block.statements.iter().rev() {
        match &statement.kind {
            StatementKind::Assign(op) => {
                if op.0 != switch_place {
                    continue;
                }

                match op.1 {
                    // The following rvalues move the place
                    // that may be const in the predecessor
                    Rvalue::Use(Operand::Move(new) | Operand::Copy(new))
                    | Rvalue::UnaryOp(_, Operand::Copy(new) | Operand::Move(new))
                    | Rvalue::Cast(_, Operand::Move(new) | Operand::Copy(new), _)
                    | Rvalue::Repeat(Operand::Move(new) | Operand::Copy(new), _)
                    | Rvalue::Discriminant(new)
                    => switch_place = new,

                    // The following rvalues might still make the block
                    // be valid but for now we reject them
                    Rvalue::Len(_)
                    | Rvalue::Ref(_, _, _)
                    | Rvalue::BinaryOp(_, _)
                    | Rvalue::CheckedBinaryOp(_, _)
                    | Rvalue::Aggregate(_, _)

                    // The following rvalues definitely mean we cannot
                    // or should not apply this optimization
                    | Rvalue::Use(Operand::Constant(_))
                    | Rvalue::Repeat(Operand::Constant(_), _)
                    | Rvalue::ThreadLocalRef(_)
                    | Rvalue::AddressOf(_, _)
                    | Rvalue::NullaryOp(_, _)
                    | Rvalue::ShallowInitBox(_, _)
                    | Rvalue::UnaryOp(_, Operand::Constant(_))
                    | Rvalue::Cast(_, Operand::Constant(_), _)
                    => return None,
                }
            }

            // These statements have no influence on the place
            // we are interested in
            StatementKind::FakeRead(_)
            | StatementKind::StorageLive(_)
            | StatementKind::StorageDead(_)
            | StatementKind::Retag(_, _)
            | StatementKind::AscribeUserType(_, _)
            | StatementKind::Coverage(_)
            | StatementKind::CopyNonOverlapping(_)
            | StatementKind::Nop => {}

            // If the discriminant is set, it is always set
            // as a constant, so the job is already done.
            // As we are **ignoring projections**, if the place
            // we are tracking sees its discriminant be set,
            // that means we had to be tracking the discriminant
            // specifically (as it is impossible to switch over
            // an enum directly, and if we were switching over
            // its content, we would have had to at least cast it to
            // some variant first)
            StatementKind::SetDiscriminant { place, .. } => {
                if **place == switch_place {
                    return None;
                }
            }
        }
    }

    Some(switch_place)
}
Commit	Line	Data
94222f64 XL	1	//! A pass that duplicates switch-terminated blocks
	2	//! into a new copy for each predecessor, provided
	3	//! the predecessor sets the value being switched
	4	//! over to a constant.
	5	//!
	6	//! The purpose of this pass is to help constant
	7	//! propagation passes to simplify the switch terminator
	8	//! of the copied blocks into gotos when some predecessors
	9	//! statically determine the output of switches.
	10	//!
	11	//! ```text
	12	//! x = 12 --- ---> something
	13	//! \ / 12
	14	//! --> switch x
	15	//! / \ otherwise
	16	//! x = y --- ---> something else
	17	//! ```
	18	//! becomes
	19	//! ```text
	20	//! x = 12 ---> switch x ------> something
	21	//! \ / 12
	22	//! X
	23	//! / \ otherwise
	24	//! x = y ---> switch x ------> something else
	25	//! ```
	26	//! so it can hopefully later be turned by another pass into
	27	//! ```text
	28	//! x = 12 --------------------> something
	29	//! / 12
	30	//! /
	31	//! / otherwise
	32	//! x = y ---- switch x ------> something else
	33	//! ```
	34	//!
	35	//! This optimization is meant to cover simple cases
	36	//! like `?` desugaring. For now, it thus focuses on
	37	//! simplicity rather than completeness (it notably
	38	//! sometimes duplicates abusively).
	39
c295e0f8	40	use crate::MirPass;
94222f64 XL	41	use rustc_middle::mir::*;
	42	use rustc_middle::ty::TyCtxt;
	43	use smallvec::SmallVec;
	44
	45	pub struct SeparateConstSwitch;
	46
	47	impl<'tcx> MirPass<'tcx> for SeparateConstSwitch {
a2a8927a XL	48	fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
	49	sess.mir_opt_level() >= 4
	50	}
94222f64	51
a2a8927a	52	fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
94222f64 XL	53	// If execution did something, applying a simplification layer
	54	// helps later passes optimize the copy away.
	55	if separate_const_switch(body) > 0 {
	56	super::simplify::simplify_cfg(tcx, body);
	57	}
	58	}
	59	}
	60
	61	/// Returns the amount of blocks that were duplicated
a2a8927a	62	pub fn separate_const_switch(body: &mut Body<'_>) -> usize {
94222f64 XL	63	let mut new_blocks: SmallVec<[(BasicBlock, BasicBlock); 6]> = SmallVec::new();
	64	let predecessors = body.predecessors();
	65	'block_iter: for (block_id, block) in body.basic_blocks().iter_enumerated() {
	66	if let TerminatorKind::SwitchInt {
	67	discr: Operand::Copy(switch_place) \| Operand::Move(switch_place),
	68	..
	69	} = block.terminator().kind
	70	{
	71	// If the block is on an unwind path, do not
	72	// apply the optimization as unwind paths
	73	// rely on a unique parent invariant
	74	if block.is_cleanup {
	75	continue 'block_iter;
	76	}
	77
	78	// If the block has fewer than 2 predecessors, ignore it
	79	// we could maybe chain blocks that have exactly one
	80	// predecessor, but for now we ignore that
	81	if predecessors[block_id].len() < 2 {
	82	continue 'block_iter;
	83	}
	84
	85	// First, let's find a non-const place
	86	// that determines the result of the switch
	87	if let Some(switch_place) = find_determining_place(switch_place, block) {
	88	// We now have an input place for which it would
	89	// be interesting if predecessors assigned it from a const
	90
	91	let mut predecessors_left = predecessors[block_id].len();
	92	'predec_iter: for predecessor_id in predecessors[block_id].iter().copied() {
	93	let predecessor = &body.basic_blocks()[predecessor_id];
	94
	95	// First we make sure the predecessor jumps
	96	// in a reasonable way
	97	match &predecessor.terminator().kind {
	98	// The following terminators are
	99	// unconditionally valid
	100	TerminatorKind::Goto { .. } \| TerminatorKind::SwitchInt { .. } => {}
	101
	102	TerminatorKind::FalseEdge { real_target, .. } => {
	103	if *real_target != block_id {
	104	continue 'predec_iter;
	105	}
	106	}
	107
	108	// The following terminators are not allowed
	109	TerminatorKind::Resume
	110	\| TerminatorKind::Drop { .. }
	111	\| TerminatorKind::DropAndReplace { .. }
	112	\| TerminatorKind::Call { .. }
	113	\| TerminatorKind::Assert { .. }
	114	\| TerminatorKind::FalseUnwind { .. }
	115	\| TerminatorKind::Yield { .. }
	116	\| TerminatorKind::Abort
	117	\| TerminatorKind::Return
	118	\| TerminatorKind::Unreachable
	119	\| TerminatorKind::InlineAsm { .. }
	120	\| TerminatorKind::GeneratorDrop => {
	121	continue 'predec_iter;
	122	}
	123	}
	124
	125	if is_likely_const(switch_place, predecessor) {
	126	new_blocks.push((predecessor_id, block_id));
127	predecessors_left -= 1;
128	if predecessors_left < 2 {
129	// If the original block only has one predecessor left,
130	// we have nothing left to do
131	break 'predec_iter;
132	}
133	}
134	}
135	}
136	}
137	}
138
139	// Once the analysis is done, perform the duplication
140	let body_span = body.span;
141	let copied_blocks = new_blocks.len();
142	let blocks = body.basic_blocks_mut();
143	for (pred_id, target_id) in new_blocks {
144	let new_block = blocks[target_id].clone();
145	let new_block_id = blocks.push(new_block);
146	let terminator = blocks[pred_id].terminator_mut();
147
148	match terminator.kind {
149	TerminatorKind::Goto { ref mut target } => {
150	*target = new_block_id;
151	}
152
153	TerminatorKind::FalseEdge { ref mut real_target, .. } => {
154	if *real_target == target_id {
155	*real_target = new_block_id;
156	}
157	}
158
159	TerminatorKind::SwitchInt { ref mut targets, .. } => {
160	targets.all_targets_mut().iter_mut().for_each(\|x\| {
161	if *x == target_id {
162	*x = new_block_id;
163	}
164	});
165	}
166
167	TerminatorKind::Resume
168	\| TerminatorKind::Abort
169	\| TerminatorKind::Return
170	\| TerminatorKind::Unreachable
171	\| TerminatorKind::GeneratorDrop
172	\| TerminatorKind::Assert { .. }
173	\| TerminatorKind::DropAndReplace { .. }
174	\| TerminatorKind::FalseUnwind { .. }
175	\| TerminatorKind::Drop { .. }
176	\| TerminatorKind::Call { .. }
177	\| TerminatorKind::InlineAsm { .. }
178	\| TerminatorKind::Yield { .. } => {
179	span_bug!(
180	body_span,
181	"basic block terminator had unexpected kind {:?}",
182	&terminator.kind
183	)
184	}
185	}
186	}
187
188	copied_blocks
189	}
190
191	/// This function describes a rough heuristic guessing
192	/// whether a place is last set with a const within the block.
193	/// Notably, it will be overly pessimistic in cases that are already
194	/// not handled by `separate_const_switch`.
195	fn is_likely_const<'tcx>(mut tracked_place: Place<'tcx>, block: &BasicBlockData<'tcx>) -> bool {
196	for statement in block.statements.iter().rev() {
197	match &statement.kind {
198	StatementKind::Assign(assign) => {
199	if assign.0 == tracked_place {
200	match assign.1 {
201	// These rvalues are definitely constant
202	Rvalue::Use(Operand::Constant(_))
203	\| Rvalue::Ref(_, _, _)
204	\| Rvalue::AddressOf(_, _)
205	\| Rvalue::Cast(_, Operand::Constant(_), _)
206	\| Rvalue::NullaryOp(_, _)
c295e0f8	207	\| Rvalue::ShallowInitBox(_, _)
94222f64 XL	208	\| Rvalue::UnaryOp(_, Operand::Constant(_)) => return true,
	209
	210	// These rvalues make things ambiguous
	211	Rvalue::Repeat(_, _)
	212	\| Rvalue::ThreadLocalRef(_)
	213	\| Rvalue::Len(_)
	214	\| Rvalue::BinaryOp(_, _)
	215	\| Rvalue::CheckedBinaryOp(_, _)
	216	\| Rvalue::Aggregate(_, _) => return false,
	217
	218	// These rvalues move the place to track
	219	Rvalue::Cast(_, Operand::Copy(place) \| Operand::Move(place), _)
	220	\| Rvalue::Use(Operand::Copy(place) \| Operand::Move(place))
	221	\| Rvalue::UnaryOp(_, Operand::Copy(place) \| Operand::Move(place))
	222	\| Rvalue::Discriminant(place) => tracked_place = place,
	223	}
	224	}
	225	}
	226
	227	// If the discriminant is set, it is always set
	228	// as a constant, so the job is done.
	229	// As we are ignoring projections, if the place
	230	// we are tracking sees its discriminant be set,
	231	// that means we had to be tracking the discriminant
	232	// specifically (as it is impossible to switch over
	233	// an enum directly, and if we were switching over
	234	// its content, we would have had to at least cast it to
	235	// some variant first)
	236	StatementKind::SetDiscriminant { place, .. } => {
	237	if **place == tracked_place {
	238	return true;
	239	}
	240	}
	241
94222f64 XL	242	// These statements have no influence on the place
	243	// we are interested in
	244	StatementKind::FakeRead(_)
	245	\| StatementKind::StorageLive(_)
	246	\| StatementKind::Retag(_, _)
	247	\| StatementKind::AscribeUserType(_, _)
	248	\| StatementKind::Coverage(_)
	249	\| StatementKind::StorageDead(_)
	250	\| StatementKind::CopyNonOverlapping(_)
	251	\| StatementKind::Nop => {}
	252	}
	253	}
	254
	255	// If no good reason for the place to be const is found,
	256	// give up. We could maybe go up predecessors, but in
	257	// most cases giving up now should be sufficient.
	258	false
	259	}
	260
	261	/// Finds a unique place that entirely determines the value
	262	/// of `switch_place`, if it exists. This is only a heuristic.
	263	/// Ideally we would like to track multiple determining places
	264	/// for some edge cases, but one is enough for a lot of situations.
	265	fn find_determining_place<'tcx>(
	266	mut switch_place: Place<'tcx>,
	267	block: &BasicBlockData<'tcx>,
	268	) -> Option<Place<'tcx>> {
	269	for statement in block.statements.iter().rev() {
	270	match &statement.kind {
	271	StatementKind::Assign(op) => {
	272	if op.0 != switch_place {
	273	continue;
	274	}
	275
	276	match op.1 {
	277	// The following rvalues move the place
	278	// that may be const in the predecessor
	279	Rvalue::Use(Operand::Move(new) \| Operand::Copy(new))
	280	\| Rvalue::UnaryOp(_, Operand::Copy(new) \| Operand::Move(new))
	281	\| Rvalue::Cast(_, Operand::Move(new) \| Operand::Copy(new), _)
	282	\| Rvalue::Repeat(Operand::Move(new) \| Operand::Copy(new), _)
	283	\| Rvalue::Discriminant(new)
	284	=> switch_place = new,
	285
	286	// The following rvalues might still make the block
	287	// be valid but for now we reject them
	288	Rvalue::Len(_)
	289	\| Rvalue::Ref(_, _, _)
	290	\| Rvalue::BinaryOp(_, _)
	291	\| Rvalue::CheckedBinaryOp(_, _)
	292	\| Rvalue::Aggregate(_, _)
	293
	294	// The following rvalues definitely mean we cannot
	295	// or should not apply this optimization
	296	\| Rvalue::Use(Operand::Constant(_))
	297	\| Rvalue::Repeat(Operand::Constant(_), _)
	298	\| Rvalue::ThreadLocalRef(_)
	299	\| Rvalue::AddressOf(_, _)
	300	\| Rvalue::NullaryOp(_, _)
c295e0f8	301	\| Rvalue::ShallowInitBox(_, _)
94222f64 XL	302	\| Rvalue::UnaryOp(_, Operand::Constant(_))
	303	\| Rvalue::Cast(_, Operand::Constant(_), _)
	304	=> return None,
	305	}
	306	}
	307
	308	// These statements have no influence on the place
	309	// we are interested in
	310	StatementKind::FakeRead(_)
	311	\| StatementKind::StorageLive(_)
	312	\| StatementKind::StorageDead(_)
	313	\| StatementKind::Retag(_, _)
	314	\| StatementKind::AscribeUserType(_, _)
	315	\| StatementKind::Coverage(_)
	316	\| StatementKind::CopyNonOverlapping(_)
	317	\| StatementKind::Nop => {}
	318
94222f64 XL	319	// If the discriminant is set, it is always set
	320	// as a constant, so the job is already done.
	321	// As we are ignoring projections, if the place
	322	// we are tracking sees its discriminant be set,
	323	// that means we had to be tracking the discriminant
	324	// specifically (as it is impossible to switch over
	325	// an enum directly, and if we were switching over
	326	// its content, we would have had to at least cast it to
	327	// some variant first)
	328	StatementKind::SetDiscriminant { place, .. } => {
	329	if **place == switch_place {
	330	return None;
	331	}
	332	}
	333	}
	334	}
	335
	336	Some(switch_place)
	337	}