[rustc.git] / compiler / rustc_mir / src / transform / 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::transform::MirPass;
use rustc_middle::mir::*;
use rustc_middle::ty::TyCtxt;
use smallvec::SmallVec;

pub struct SeparateConstSwitch;

impl<'tcx> MirPass<'tcx> for SeparateConstSwitch {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        if tcx.sess.mir_opt_level() < 4 {
            return;
        }

        // 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<'tcx>(body: &mut Body<'tcx>) -> 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::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;
                }
            }

            // If inline assembly is found, we probably should
            // not try to analyze the code
            StatementKind::LlvmInlineAsm(_) => return false,

            // 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::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 inline assembly is found, we probably should
            // not try to analyze the code
            StatementKind::LlvmInlineAsm(_) => return None,

            // 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
	40	use crate::transform::MirPass;
	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 {
	48	fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
	49	if tcx.sess.mir_opt_level() < 4 {
	50	return;
	51	}
	52
	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
	62	pub fn separate_const_switch<'tcx>(body: &mut Body<'tcx>) -> usize {
	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(_, _)
207	\| Rvalue::UnaryOp(_, Operand::Constant(_)) => return true,
208
209	// These rvalues make things ambiguous
210	Rvalue::Repeat(_, _)
211	\| Rvalue::ThreadLocalRef(_)
212	\| Rvalue::Len(_)
213	\| Rvalue::BinaryOp(_, _)
214	\| Rvalue::CheckedBinaryOp(_, _)
215	\| Rvalue::Aggregate(_, _) => return false,
216
217	// These rvalues move the place to track
218	Rvalue::Cast(_, Operand::Copy(place) \| Operand::Move(place), _)
219	\| Rvalue::Use(Operand::Copy(place) \| Operand::Move(place))
220	\| Rvalue::UnaryOp(_, Operand::Copy(place) \| Operand::Move(place))
221	\| Rvalue::Discriminant(place) => tracked_place = place,
222	}
223	}
224	}
225
226	// If the discriminant is set, it is always set
227	// as a constant, so the job is done.
228	// As we are ignoring projections, if the place
229	// we are tracking sees its discriminant be set,
230	// that means we had to be tracking the discriminant
231	// specifically (as it is impossible to switch over
232	// an enum directly, and if we were switching over
233	// its content, we would have had to at least cast it to
234	// some variant first)
235	StatementKind::SetDiscriminant { place, .. } => {
236	if **place == tracked_place {
237	return true;
238	}
239	}
240
241	// If inline assembly is found, we probably should
242	// not try to analyze the code
243	StatementKind::LlvmInlineAsm(_) => return false,
244
245	// These statements have no influence on the place
246	// we are interested in
247	StatementKind::FakeRead(_)
248	\| StatementKind::StorageLive(_)
249	\| StatementKind::Retag(_, _)
250	\| StatementKind::AscribeUserType(_, _)
251	\| StatementKind::Coverage(_)
252	\| StatementKind::StorageDead(_)
253	\| StatementKind::CopyNonOverlapping(_)
254	\| StatementKind::Nop => {}
255	}
256	}
257
258	// If no good reason for the place to be const is found,
259	// give up. We could maybe go up predecessors, but in
260	// most cases giving up now should be sufficient.
261	false
262	}
263
264	/// Finds a unique place that entirely determines the value
265	/// of `switch_place`, if it exists. This is only a heuristic.
266	/// Ideally we would like to track multiple determining places
267	/// for some edge cases, but one is enough for a lot of situations.
268	fn find_determining_place<'tcx>(
269	mut switch_place: Place<'tcx>,
270	block: &BasicBlockData<'tcx>,
271	) -> Option<Place<'tcx>> {
272	for statement in block.statements.iter().rev() {
273	match &statement.kind {
274	StatementKind::Assign(op) => {
275	if op.0 != switch_place {
276	continue;
277	}
278
279	match op.1 {
280	// The following rvalues move the place
281	// that may be const in the predecessor
282	Rvalue::Use(Operand::Move(new) \| Operand::Copy(new))
283	\| Rvalue::UnaryOp(_, Operand::Copy(new) \| Operand::Move(new))
284	\| Rvalue::Cast(_, Operand::Move(new) \| Operand::Copy(new), _)
285	\| Rvalue::Repeat(Operand::Move(new) \| Operand::Copy(new), _)
286	\| Rvalue::Discriminant(new)
287	=> switch_place = new,
288
289	// The following rvalues might still make the block
290	// be valid but for now we reject them
291	Rvalue::Len(_)
292	\| Rvalue::Ref(_, _, _)
293	\| Rvalue::BinaryOp(_, _)
294	\| Rvalue::CheckedBinaryOp(_, _)
295	\| Rvalue::Aggregate(_, _)
296
297	// The following rvalues definitely mean we cannot
298	// or should not apply this optimization
299	\| Rvalue::Use(Operand::Constant(_))
300	\| Rvalue::Repeat(Operand::Constant(_), _)
301	\| Rvalue::ThreadLocalRef(_)
302	\| Rvalue::AddressOf(_, _)
303	\| Rvalue::NullaryOp(_, _)
304	\| Rvalue::UnaryOp(_, Operand::Constant(_))
305	\| Rvalue::Cast(_, Operand::Constant(_), _)
306	=> return None,
307	}
308	}
309
310	// These statements have no influence on the place
311	// we are interested in
312	StatementKind::FakeRead(_)
313	\| StatementKind::StorageLive(_)
314	\| StatementKind::StorageDead(_)
315	\| StatementKind::Retag(_, _)
316	\| StatementKind::AscribeUserType(_, _)
317	\| StatementKind::Coverage(_)
318	\| StatementKind::CopyNonOverlapping(_)
319	\| StatementKind::Nop => {}
320
321	// If inline assembly is found, we probably should
322	// not try to analyze the code
323	StatementKind::LlvmInlineAsm(_) => return None,
324
325	// If the discriminant is set, it is always set
326	// as a constant, so the job is already done.
327	// As we are ignoring projections, if the place
328	// we are tracking sees its discriminant be set,
329	// that means we had to be tracking the discriminant
330	// specifically (as it is impossible to switch over
331	// an enum directly, and if we were switching over
332	// its content, we would have had to at least cast it to
333	// some variant first)
334	StatementKind::SetDiscriminant { place, .. } => {
335	if **place == switch_place {
336	return None;
337	}
338	}
339	}
340	}
341
342	Some(switch_place)
343	}