From 3e0b2fac5dd02378d8a5fe2c8b0c6affed404db0 Mon Sep 17 00:00:00 2001 From: "Celina G. Val" Date: Wed, 29 Nov 2023 12:49:32 -0800 Subject: [PATCH 1/2] Change SwitchTarget representation The new structure encodes its invariant, which reduces the likelihood of having an inconsistent representation. It is also more intuitive and user friendly. I encapsulated the structure for now in case we decide to change it back. --- .../rustc_smir/src/rustc_smir/convert/mir.rs | 11 ++- compiler/stable_mir/src/mir/body.rs | 85 +++++++++++++------ compiler/stable_mir/src/mir/pretty.rs | 15 ++-- compiler/stable_mir/src/mir/visit.rs | 5 +- 4 files changed, 72 insertions(+), 44 deletions(-) diff --git a/compiler/rustc_smir/src/rustc_smir/convert/mir.rs b/compiler/rustc_smir/src/rustc_smir/convert/mir.rs index e9b835f90db24..165b8c50e70f9 100644 --- a/compiler/rustc_smir/src/rustc_smir/convert/mir.rs +++ b/compiler/rustc_smir/src/rustc_smir/convert/mir.rs @@ -579,13 +579,12 @@ impl<'tcx> Stable<'tcx> for mir::TerminatorKind<'tcx> { mir::TerminatorKind::SwitchInt { discr, targets } => TerminatorKind::SwitchInt { discr: discr.stable(tables), targets: { - let (value_vec, mut target_vec): (Vec<_>, Vec<_>) = - targets.iter().map(|(value, target)| (value, target.as_usize())).unzip(); - // We need to push otherwise as last element to ensure it's same as in MIR. - target_vec.push(targets.otherwise().as_usize()); - stable_mir::mir::SwitchTargets { value: value_vec, targets: target_vec } + let branches = targets.iter().map(|(val, target)| (val, target.as_usize())); + stable_mir::mir::SwitchTargets::new( + branches.collect(), + targets.otherwise().as_usize(), + ) }, - otherwise: targets.otherwise().as_usize(), }, mir::TerminatorKind::UnwindResume => TerminatorKind::Resume, mir::TerminatorKind::UnwindTerminate(_) => TerminatorKind::Abort, diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs index 02a2868767659..8ccc6732a8df0 100644 --- a/compiler/stable_mir/src/mir/body.rs +++ b/compiler/stable_mir/src/mir/body.rs @@ -3,7 +3,7 @@ use crate::ty::{ AdtDef, ClosureDef, Const, CoroutineDef, GenericArgs, Movability, Region, RigidTy, Ty, TyKind, }; use crate::{Error, Opaque, Span, Symbol}; -use std::{io, slice}; +use std::io; /// The SMIR representation of a single function. #[derive(Clone, Debug)] pub struct Body { @@ -23,6 +23,8 @@ pub struct Body { pub(super) var_debug_info: Vec, } +pub type BasicBlockIdx = usize; + impl Body { /// Constructs a `Body`. /// @@ -114,22 +116,21 @@ pub struct Terminator { } impl Terminator { - pub fn successors(&self) -> Successors<'_> { + pub fn successors(&self) -> Successors { self.kind.successors() } } -pub type Successors<'a> = impl Iterator + 'a; +pub type Successors = Vec; #[derive(Clone, Debug, Eq, PartialEq)] pub enum TerminatorKind { Goto { - target: usize, + target: BasicBlockIdx, }, SwitchInt { discr: Operand, targets: SwitchTargets, - otherwise: usize, }, Resume, Abort, @@ -137,43 +138,42 @@ pub enum TerminatorKind { Unreachable, Drop { place: Place, - target: usize, + target: BasicBlockIdx, unwind: UnwindAction, }, Call { func: Operand, args: Vec, destination: Place, - target: Option, + target: Option, unwind: UnwindAction, }, Assert { cond: Operand, expected: bool, msg: AssertMessage, - target: usize, + target: BasicBlockIdx, unwind: UnwindAction, }, - CoroutineDrop, InlineAsm { template: String, operands: Vec, options: String, line_spans: String, - destination: Option, + destination: Option, unwind: UnwindAction, }, } impl TerminatorKind { - pub fn successors(&self) -> Successors<'_> { + pub fn successors(&self) -> Successors { use self::TerminatorKind::*; match *self { - Call { target: Some(t), unwind: UnwindAction::Cleanup(ref u), .. } - | Drop { target: t, unwind: UnwindAction::Cleanup(ref u), .. } - | Assert { target: t, unwind: UnwindAction::Cleanup(ref u), .. } - | InlineAsm { destination: Some(t), unwind: UnwindAction::Cleanup(ref u), .. } => { - Some(t).into_iter().chain(slice::from_ref(u).into_iter().copied()) + Call { target: Some(t), unwind: UnwindAction::Cleanup(u), .. } + | Drop { target: t, unwind: UnwindAction::Cleanup(u), .. } + | Assert { target: t, unwind: UnwindAction::Cleanup(u), .. } + | InlineAsm { destination: Some(t), unwind: UnwindAction::Cleanup(u), .. } => { + vec![t, u] } Goto { target: t } | Call { target: None, unwind: UnwindAction::Cleanup(t), .. } @@ -182,21 +182,18 @@ impl TerminatorKind { | Assert { target: t, unwind: _, .. } | InlineAsm { destination: None, unwind: UnwindAction::Cleanup(t), .. } | InlineAsm { destination: Some(t), unwind: _, .. } => { - Some(t).into_iter().chain((&[]).into_iter().copied()) + vec![t] } - CoroutineDrop - | Return + Return | Resume | Abort | Unreachable | Call { target: None, unwind: _, .. } | InlineAsm { destination: None, unwind: _, .. } => { - None.into_iter().chain((&[]).into_iter().copied()) - } - SwitchInt { ref targets, .. } => { - None.into_iter().chain(targets.targets.iter().copied()) + vec![] } + SwitchInt { ref targets, .. } => targets.all_targets(), } } @@ -205,7 +202,6 @@ impl TerminatorKind { TerminatorKind::Goto { .. } | TerminatorKind::Return | TerminatorKind::Unreachable - | TerminatorKind::CoroutineDrop | TerminatorKind::Resume | TerminatorKind::Abort | TerminatorKind::SwitchInt { .. } => None, @@ -231,7 +227,7 @@ pub enum UnwindAction { Continue, Unreachable, Terminate, - Cleanup(usize), + Cleanup(BasicBlockIdx), } #[derive(Clone, Debug, Eq, PartialEq)] @@ -662,10 +658,45 @@ pub struct Constant { pub literal: Const, } +/// The possible branch sites of a [TerminatorKind::SwitchInt]. #[derive(Clone, Debug, Eq, PartialEq)] pub struct SwitchTargets { - pub value: Vec, - pub targets: Vec, + /// The conditional branches where the first element represents the value that guards this + /// branch, and the second element is the branch target. + branches: Vec<(u128, BasicBlockIdx)>, + /// The `otherwise` branch which will be taken in case none of the conditional branches are + /// satisfied. + otherwise: BasicBlockIdx, +} + +impl SwitchTargets { + /// All possible targets including the `otherwise` target. + pub fn all_targets(&self) -> Successors { + Some(self.otherwise) + .into_iter() + .chain(self.branches.iter().map(|(_, target)| *target)) + .collect() + } + + /// The `otherwise` branch target. + pub fn otherwise(&self) -> BasicBlockIdx { + self.otherwise + } + + /// The conditional targets which are only taken if the pattern matches the given value. + pub fn branches(&self) -> impl Iterator + '_ { + self.branches.iter().copied() + } + + /// The number of targets including `otherwise`. + pub fn len(&self) -> usize { + self.branches.len() + 1 + } + + /// Create a new SwitchTargets from the given branches and `otherwise` target. + pub fn new(branches: Vec<(u128, BasicBlockIdx)>, otherwise: BasicBlockIdx) -> SwitchTargets { + SwitchTargets { branches, otherwise } + } } #[derive(Copy, Clone, Debug, Eq, PartialEq)] diff --git a/compiler/stable_mir/src/mir/pretty.rs b/compiler/stable_mir/src/mir/pretty.rs index 759c3b148dfef..3a0eed521dc6e 100644 --- a/compiler/stable_mir/src/mir/pretty.rs +++ b/compiler/stable_mir/src/mir/pretty.rs @@ -76,7 +76,8 @@ pub fn pretty_statement(statement: &StatementKind) -> String { pub fn pretty_terminator(terminator: &TerminatorKind, w: &mut W) -> io::Result<()> { write!(w, "{}", pretty_terminator_head(terminator))?; - let successor_count = terminator.successors().count(); + let successors = terminator.successors(); + let successor_count = successors.len(); let labels = pretty_successor_labels(terminator); let show_unwind = !matches!(terminator.unwind(), None | Some(UnwindAction::Cleanup(_))); @@ -98,12 +99,12 @@ pub fn pretty_terminator(terminator: &TerminatorKind, w: &mut W) - Ok(()) } (1, false) => { - write!(w, " -> {:?}", terminator.successors().next().unwrap())?; + write!(w, " -> {:?}", successors[0])?; Ok(()) } _ => { write!(w, " -> [")?; - for (i, target) in terminator.successors().enumerate() { + for (i, target) in successors.iter().enumerate() { if i > 0 { write!(w, ", ")?; } @@ -157,7 +158,6 @@ pub fn pretty_terminator_head(terminator: &TerminatorKind) -> String { pretty.push_str(")"); pretty } - CoroutineDrop => format!(" coroutine_drop"), InlineAsm { .. } => todo!(), } } @@ -165,12 +165,11 @@ pub fn pretty_terminator_head(terminator: &TerminatorKind) -> String { pub fn pretty_successor_labels(terminator: &TerminatorKind) -> Vec { use self::TerminatorKind::*; match terminator { - Resume | Abort | Return | Unreachable | CoroutineDrop => vec![], + Resume | Abort | Return | Unreachable => vec![], Goto { .. } => vec!["".to_string()], SwitchInt { targets, .. } => targets - .value - .iter() - .map(|target| format!("{}", target)) + .branches() + .map(|(val, _target)| format!("{val}")) .chain(iter::once("otherwise".into())) .collect(), Drop { unwind: UnwindAction::Cleanup(_), .. } => vec!["return".into(), "unwind".into()], diff --git a/compiler/stable_mir/src/mir/visit.rs b/compiler/stable_mir/src/mir/visit.rs index 69bf6ca72b0cf..0c44781c46394 100644 --- a/compiler/stable_mir/src/mir/visit.rs +++ b/compiler/stable_mir/src/mir/visit.rs @@ -237,8 +237,7 @@ pub trait MirVisitor { TerminatorKind::Goto { .. } | TerminatorKind::Resume | TerminatorKind::Abort - | TerminatorKind::Unreachable - | TerminatorKind::CoroutineDrop => {} + | TerminatorKind::Unreachable => {} TerminatorKind::Assert { cond, expected: _, msg, target: _, unwind: _ } => { self.visit_operand(cond, location); self.visit_assert_msg(msg, location); @@ -268,7 +267,7 @@ pub trait MirVisitor { let local = RETURN_LOCAL; self.visit_local(&local, PlaceContext::NON_MUTATING, location); } - TerminatorKind::SwitchInt { discr, targets: _, otherwise: _ } => { + TerminatorKind::SwitchInt { discr, targets: _ } => { self.visit_operand(discr, location); } } From 9d2c92377d1bf4d8b408b57b681bb6a04eb153a6 Mon Sep 17 00:00:00 2001 From: "Celina G. Val" Date: Thu, 30 Nov 2023 11:40:55 -0800 Subject: [PATCH 2/2] Fix SwitchTarget pretty print We currently rely on the order of successors to be conditional branches first, followed by the otherwise target. --- compiler/stable_mir/src/mir/body.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs index 8ccc6732a8df0..3a4f42835622b 100644 --- a/compiler/stable_mir/src/mir/body.rs +++ b/compiler/stable_mir/src/mir/body.rs @@ -672,10 +672,7 @@ pub struct SwitchTargets { impl SwitchTargets { /// All possible targets including the `otherwise` target. pub fn all_targets(&self) -> Successors { - Some(self.otherwise) - .into_iter() - .chain(self.branches.iter().map(|(_, target)| *target)) - .collect() + self.branches.iter().map(|(_, target)| *target).chain(Some(self.otherwise)).collect() } /// The `otherwise` branch target.