add successors and their formatter

2 files changed, 159 insertions, 23 deletions

diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
index 0966f9198a2..02a28687676 100644
--- a/compiler/stable_mir/src/mir/body.rs
+++ b/compiler/stable_mir/src/mir/body.rs
@@ -3,8 +3,7 @@ use crate::ty::{
     AdtDef, ClosureDef, Const, CoroutineDef, GenericArgs, Movability, Region, RigidTy, Ty, TyKind,
 };
 use crate::{Error, Opaque, Span, Symbol};
-use std::io;
-
+use std::{io, slice};
 /// The SMIR representation of a single function.
 #[derive(Clone, Debug)]
 pub struct Body {
@@ -83,7 +82,8 @@ impl Body {
                         Ok(())
                     })
                     .collect::<Vec<_>>();
-                writeln!(w, "{}", pretty_terminator(&block.terminator.kind))?;
+                pretty_terminator(&block.terminator.kind, w)?;
+                writeln!(w, "").unwrap();
                 writeln!(w, "    }}").unwrap();
                 Ok(())
             })
@@ -101,7 +101,7 @@ pub struct LocalDecl {
     pub mutability: Mutability,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, PartialEq, Eq, Debug)]
 pub struct BasicBlock {
     pub statements: Vec<Statement>,
     pub terminator: Terminator,
@@ -113,6 +113,14 @@ pub struct Terminator {
     pub span: Span,
 }
 
+impl Terminator {
+    pub fn successors(&self) -> Successors<'_> {
+        self.kind.successors()
+    }
+}
+
+pub type Successors<'a> = impl Iterator<Item = usize> + 'a;
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum TerminatorKind {
     Goto {
@@ -120,7 +128,7 @@ pub enum TerminatorKind {
     },
     SwitchInt {
         discr: Operand,
-        targets: Vec<SwitchTarget>,
+        targets: SwitchTargets,
         otherwise: usize,
     },
     Resume,
@@ -157,6 +165,58 @@ pub enum TerminatorKind {
     },
 }
 
+impl TerminatorKind {
+    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())
+            }
+            Goto { target: t }
+            | Call { target: None, unwind: UnwindAction::Cleanup(t), .. }
+            | Call { target: Some(t), unwind: _, .. }
+            | Drop { target: t, unwind: _, .. }
+            | Assert { target: t, unwind: _, .. }
+            | InlineAsm { destination: None, unwind: UnwindAction::Cleanup(t), .. }
+            | InlineAsm { destination: Some(t), unwind: _, .. } => {
+                Some(t).into_iter().chain((&[]).into_iter().copied())
+            }
+
+            CoroutineDrop
+            | 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())
+            }
+        }
+    }
+
+    pub fn unwind(&self) -> Option<&UnwindAction> {
+        match *self {
+            TerminatorKind::Goto { .. }
+            | TerminatorKind::Return
+            | TerminatorKind::Unreachable
+            | TerminatorKind::CoroutineDrop
+            | TerminatorKind::Resume
+            | TerminatorKind::Abort
+            | TerminatorKind::SwitchInt { .. } => None,
+            TerminatorKind::Call { ref unwind, .. }
+            | TerminatorKind::Assert { ref unwind, .. }
+            | TerminatorKind::Drop { ref unwind, .. }
+            | TerminatorKind::InlineAsm { ref unwind, .. } => Some(unwind),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct InlineAsmOperand {
     pub in_value: Option<Operand>,
@@ -603,9 +663,9 @@ pub struct Constant {
 }
 
 #[derive(Clone, Debug, Eq, PartialEq)]
-pub struct SwitchTarget {
-    pub value: u128,
-    pub target: usize,
+pub struct SwitchTargets {
+    pub value: Vec<u128>,
+    pub targets: Vec<usize>,
 }
 
 #[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 6fa44088742..ee2237394ca 100644
--- a/compiler/stable_mir/src/mir/pretty.rs
+++ b/compiler/stable_mir/src/mir/pretty.rs
@@ -1,7 +1,9 @@
 use crate::crate_def::CrateDef;
-use crate::mir::{Operand, Rvalue, StatementKind};
+use crate::mir::{Operand, Rvalue, StatementKind, UnwindAction};
 use crate::ty::{DynKind, FloatTy, IntTy, RigidTy, TyKind, UintTy};
 use crate::{with, Body, CrateItem, Mutability};
+use std::io::Write;
+use std::{io, iter};
 
 use super::{AssertMessage, BinOp, TerminatorKind};
 
@@ -72,21 +74,68 @@ pub fn pretty_statement(statement: &StatementKind) -> String {
     pretty
 }
 
-pub fn pretty_terminator(terminator: &TerminatorKind) -> String {
+pub fn pretty_terminator<W: io::Write>(terminator: &TerminatorKind, w: &mut W) -> io::Result<()> {
+    write!(w, "{}", pretty_terminator_head(terminator))?;
+    let successor_count = terminator.successors().count();
+    let labels = pretty_successor_labels(terminator);
+
+    let show_unwind = !matches!(terminator.unwind(), None | Some(UnwindAction::Cleanup(_)));
+    let fmt_unwind = |fmt: &mut dyn Write| -> io::Result<()> {
+        write!(fmt, "unwind ")?;
+        match terminator.unwind() {
+            None | Some(UnwindAction::Cleanup(_)) => unreachable!(),
+            Some(UnwindAction::Continue) => write!(fmt, "continue"),
+            Some(UnwindAction::Unreachable) => write!(fmt, "unreachable"),
+            Some(UnwindAction::Terminate) => write!(fmt, "terminate"),
+        }
+    };
+
+    match (successor_count, show_unwind) {
+        (0, false) => Ok(()),
+        (0, true) => {
+            write!(w, " -> ")?;
+            fmt_unwind(w)?;
+            Ok(())
+        }
+        (1, false) => {
+            write!(w, " -> {:?}", terminator.successors().next().unwrap())?;
+            Ok(())
+        }
+        _ => {
+            write!(w, " -> [")?;
+            for (i, target) in terminator.successors().enumerate() {
+                if i > 0 {
+                    write!(w, ", ")?;
+                }
+                write!(w, "{}: {:?}", labels[i], target)?;
+            }
+            if show_unwind {
+                write!(w, ", ")?;
+                fmt_unwind(w)?;
+            }
+            write!(w, "]")
+        }
+    }?;
+
+    Ok(())
+}
+
+pub fn pretty_terminator_head(terminator: &TerminatorKind) -> String {
+    use self::TerminatorKind::*;
     let mut pretty = String::new();
     match terminator {
-        TerminatorKind::Goto { .. } => format!("        goto"),
-        TerminatorKind::SwitchInt { discr, .. } => {
+        Goto { .. } => format!("        goto"),
+        SwitchInt { discr, .. } => {
             format!("        switch({})", pretty_operand(discr))
         }
-        TerminatorKind::Resume => format!("        resume"),
-        TerminatorKind::Abort => format!("        abort"),
-        TerminatorKind::Return => format!("        return"),
-        TerminatorKind::Unreachable => format!("        unreachable"),
-        TerminatorKind::Drop { place, .. } => format!("        drop({:?})", place.local),
-        TerminatorKind::Call { func, args, destination, .. } => {
+        Resume => format!("        resume"),
+        Abort => format!("        abort"),
+        Return => format!("        return"),
+        Unreachable => format!("        unreachable"),
+        Drop { place, .. } => format!("        drop(_{:?})", place.local),
+        Call { func, args, destination, .. } => {
             pretty.push_str("        ");
-            pretty.push_str(format!("{} = ", destination.local).as_str());
+            pretty.push_str(format!("_{} = ", destination.local).as_str());
             pretty.push_str(&pretty_operand(func));
             pretty.push_str("(");
             args.iter().enumerate().for_each(|(i, arg)| {
@@ -98,18 +147,45 @@ pub fn pretty_terminator(terminator: &TerminatorKind) -> String {
             pretty.push_str(")");
             pretty
         }
-        TerminatorKind::Assert { cond, expected, msg, target: _, unwind: _ } => {
+        Assert { cond, expected, msg, target: _, unwind: _ } => {
             pretty.push_str("        assert(");
             if !expected {
                 pretty.push_str("!");
             }
-            pretty.push_str(&pretty_operand(cond));
+            pretty.push_str(format!("{} bool),", &pretty_operand(cond)).as_str());
             pretty.push_str(&pretty_assert_message(msg));
             pretty.push_str(")");
             pretty
         }
-        TerminatorKind::CoroutineDrop => format!("        coroutine_drop"),
-        TerminatorKind::InlineAsm { .. } => todo!(),
+        CoroutineDrop => format!("        coroutine_drop"),
+        InlineAsm { .. } => todo!(),
+    }
+}
+
+pub fn pretty_successor_labels(terminator: &TerminatorKind) -> Vec<String> {
+    use self::TerminatorKind::*;
+    match terminator {
+        Resume | Abort | Return | Unreachable | CoroutineDrop => vec![],
+        Goto { .. } => vec!["".to_string()],
+        SwitchInt { targets, .. } => targets
+            .value
+            .iter()
+            .map(|target| format!("{}", target))
+            .chain(iter::once("otherwise".into()))
+            .collect(),
+        Drop { unwind: UnwindAction::Cleanup(_), .. } => vec!["return".into(), "unwind".into()],
+        Drop { unwind: _, .. } => vec!["return".into()],
+        Call { target: Some(_), unwind: UnwindAction::Cleanup(_), .. } => {
+            vec!["return".into(), "unwind".into()]
+        }
+        Call { target: Some(_), unwind: _, .. } => vec!["return".into()],
+        Call { target: None, unwind: UnwindAction::Cleanup(_), .. } => vec!["unwind".into()],
+        Call { target: None, unwind: _, .. } => vec![],
+        Assert { unwind: UnwindAction::Cleanup(_), .. } => {
+            vec!["success".into(), "unwind".into()]
+        }
+        Assert { unwind: _, .. } => vec!["success".into()],
+        InlineAsm { .. } => todo!(),
     }
 }