Implement a minimal subset of MIR statements and terminators for smir

5 files changed, 209 insertions, 5 deletions

diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index c17dab6dce0..86e30dd0f6a 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -8,7 +8,7 @@
 //! For now, we are developing everything inside `rustc`, thus, we keep this module private.
 
 use crate::{
-    rustc_internal::crate_item,
+    rustc_internal::{crate_item, item_def_id},
     stable_mir::{self},
 };
 use rustc_middle::ty::{tls::with, TyCtxt};
@@ -47,3 +47,112 @@ fn smir_crate(tcx: TyCtxt<'_>, crate_num: CrateNum) -> stable_mir::Crate {
     debug!(?crate_name, ?crate_num, "smir_crate");
     stable_mir::Crate { id: crate_num.into(), name: crate_name, is_local }
 }
+
+pub fn mir_body(item: &stable_mir::CrateItem) -> stable_mir::mir::Body {
+    with(|tcx| {
+        let def_id = item_def_id(item);
+        let mir = tcx.optimized_mir(def_id);
+        stable_mir::mir::Body {
+            blocks: mir
+                .basic_blocks
+                .iter()
+                .map(|block| stable_mir::mir::BasicBlock {
+                    terminator: rustc_terminator_to_terminator(block.terminator()),
+                    statements: block.statements.iter().map(rustc_statement_to_statement).collect(),
+                })
+                .collect(),
+        }
+    })
+}
+
+fn rustc_statement_to_statement(
+    s: &rustc_middle::mir::Statement<'_>,
+) -> stable_mir::mir::Statement {
+    use rustc_middle::mir::StatementKind::*;
+    match &s.kind {
+        Assign(assign) => stable_mir::mir::Statement::Assign(
+            rustc_place_to_place(&assign.0),
+            rustc_rvalue_to_rvalue(&assign.1),
+        ),
+        FakeRead(_) => todo!(),
+        SetDiscriminant { .. } => todo!(),
+        Deinit(_) => todo!(),
+        StorageLive(_) => todo!(),
+        StorageDead(_) => todo!(),
+        Retag(_, _) => todo!(),
+        PlaceMention(_) => todo!(),
+        AscribeUserType(_, _) => todo!(),
+        Coverage(_) => todo!(),
+        Intrinsic(_) => todo!(),
+        ConstEvalCounter => todo!(),
+        Nop => stable_mir::mir::Statement::Nop,
+    }
+}
+
+fn rustc_rvalue_to_rvalue(rvalue: &rustc_middle::mir::Rvalue<'_>) -> stable_mir::mir::Operand {
+    use rustc_middle::mir::Rvalue::*;
+    match rvalue {
+        Use(op) => rustc_op_to_op(op),
+        Repeat(_, _) => todo!(),
+        Ref(_, _, _) => todo!(),
+        ThreadLocalRef(_) => todo!(),
+        AddressOf(_, _) => todo!(),
+        Len(_) => todo!(),
+        Cast(_, _, _) => todo!(),
+        BinaryOp(_, _) => todo!(),
+        CheckedBinaryOp(_, _) => todo!(),
+        NullaryOp(_, _) => todo!(),
+        UnaryOp(_, _) => todo!(),
+        Discriminant(_) => todo!(),
+        Aggregate(_, _) => todo!(),
+        ShallowInitBox(_, _) => todo!(),
+        CopyForDeref(_) => todo!(),
+    }
+}
+
+fn rustc_op_to_op(op: &rustc_middle::mir::Operand<'_>) -> stable_mir::mir::Operand {
+    use rustc_middle::mir::Operand::*;
+    match op {
+        Copy(place) => stable_mir::mir::Operand::Copy(rustc_place_to_place(place)),
+        Move(place) => stable_mir::mir::Operand::Move(rustc_place_to_place(place)),
+        Constant(c) => stable_mir::mir::Operand::Constant(c.to_string()),
+    }
+}
+
+fn rustc_place_to_place(place: &rustc_middle::mir::Place<'_>) -> stable_mir::mir::Place {
+    assert_eq!(&place.projection[..], &[]);
+    stable_mir::mir::Place { local: place.local.as_usize() }
+}
+
+fn rustc_terminator_to_terminator(
+    terminator: &rustc_middle::mir::Terminator<'_>,
+) -> stable_mir::mir::Terminator {
+    use rustc_middle::mir::TerminatorKind::*;
+    use stable_mir::mir::Terminator;
+    match &terminator.kind {
+        Goto { target } => Terminator::Goto { target: target.as_usize() },
+        SwitchInt { discr, targets } => Terminator::SwitchInt {
+            discr: rustc_op_to_op(discr),
+            targets: targets
+                .iter()
+                .map(|(value, target)| stable_mir::mir::SwitchTarget {
+                    value,
+                    target: target.as_usize(),
+                })
+                .collect(),
+            otherwise: targets.otherwise().as_usize(),
+        },
+        Resume => Terminator::Resume,
+        Abort => Terminator::Abort,
+        Return => Terminator::Return,
+        Unreachable => Terminator::Unreachable,
+        Drop { .. } => todo!(),
+        Call { .. } => todo!(),
+        Assert { .. } => todo!(),
+        Yield { .. } => todo!(),
+        GeneratorDrop => todo!(),
+        FalseEdge { .. } => todo!(),
+        FalseUnwind { .. } => todo!(),
+        InlineAsm { .. } => todo!(),
+    }
+}
diff --git a/compiler/rustc_smir/src/stable_mir/mir.rs b/compiler/rustc_smir/src/stable_mir/mir.rs
new file mode 100644
index 00000000000..a9dbc3463f8
--- /dev/null
+++ b/compiler/rustc_smir/src/stable_mir/mir.rs
@@ -0,0 +1,3 @@
+mod body;
+
+pub use body::*;
diff --git a/compiler/rustc_smir/src/stable_mir/mir/body.rs b/compiler/rustc_smir/src/stable_mir/mir/body.rs
new file mode 100644
index 00000000000..613141952a7
--- /dev/null
+++ b/compiler/rustc_smir/src/stable_mir/mir/body.rs
@@ -0,0 +1,62 @@
+pub struct Body {
+    pub blocks: Vec<BasicBlock>,
+}
+
+pub struct BasicBlock {
+    pub statements: Vec<Statement>,
+    pub terminator: Terminator,
+}
+
+pub enum Terminator {
+    Goto {
+        target: usize,
+    },
+    SwitchInt {
+        discr: Operand,
+        targets: Vec<SwitchTarget>,
+        otherwise: usize,
+    },
+    Resume,
+    Abort,
+    Return,
+    Unreachable,
+    Drop {
+        place: Place,
+        target: usize,
+        unwind: Option<usize>,
+    },
+    Call {
+        func: Operand,
+        args: Vec<Operand>,
+        destination: Place,
+        target: Option<usize>,
+        cleanup: Option<usize>,
+    },
+    Assert {
+        cond: Operand,
+        expected: bool,
+        msg: String,
+        target: usize,
+        cleanup: Option<usize>,
+    },
+}
+
+pub enum Statement {
+    Assign(Place, Operand),
+    Nop,
+}
+
+pub enum Operand {
+    Copy(Place),
+    Move(Place),
+    Constant(String),
+}
+
+pub struct Place {
+    pub local: usize,
+}
+
+pub struct SwitchTarget {
+    pub value: u128,
+    pub target: usize,
+}
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index 6b3ab620e2d..ba23186224a 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -11,6 +11,8 @@
 //! There shouldn't be any direct references to internal compiler constructs in this module.
 //! If you need an internal construct, consider using `rustc_internal` or `rustc_smir`.
 
+pub mod mir;
+
 /// Use String for now but we should replace it.
 pub type Symbol = String;
 
@@ -37,6 +39,12 @@ pub struct Crate {
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct CrateItem(pub(crate) DefId);
 
+impl CrateItem {
+    pub fn body(&self) -> mir::Body {
+        crate::rustc_smir::mir_body(self)
+    }
+}
+
 /// Access to the local crate.
 pub fn local_crate() -> Crate {
     crate::rustc_smir::local_crate()
diff --git a/tests/ui-fulldeps/stable-mir/crate-info.rs b/tests/ui-fulldeps/stable-mir/crate-info.rs
index 4458ab0162e..95797ddf073 100644
--- a/tests/ui-fulldeps/stable-mir/crate-info.rs
+++ b/tests/ui-fulldeps/stable-mir/crate-info.rs
@@ -30,16 +30,34 @@ fn test_stable_mir(tcx: TyCtxt<'_>) {
 
     // Find items in the local crate.
     let items = stable_mir::all_local_items();
-    assert!(has_item(tcx, &items, (DefKind::Fn, "foo_bar")));
-    assert!(has_item(tcx, &items, (DefKind::Fn, "foo::bar")));
+    assert!(get_item(tcx, &items, (DefKind::Fn, "foo_bar")).is_some());
+    assert!(get_item(tcx, &items, (DefKind::Fn, "foo::bar")).is_some());
 
     // Find the `std` crate.
     assert!(stable_mir::find_crate("std").is_some());
+
+    let bar = get_item(tcx, &items, (DefKind::Fn, "bar")).unwrap();
+    let body = bar.body();
+    assert_eq!(body.blocks.len(), 1);
+    let block = &body.blocks[0];
+    assert_eq!(block.statements.len(), 1);
+    match &block.statements[0] {
+        stable_mir::mir::Statement::Assign(..) => {}
+        _ => panic!(),
+    }
+    match &block.terminator {
+        stable_mir::mir::Terminator::Return => {}
+        _ => panic!(),
+    }
 }
 
 // Use internal API to find a function in a crate.
-fn has_item(tcx: TyCtxt, items: &stable_mir::CrateItems, item: (DefKind, &str)) -> bool {
-    items.iter().any(|crate_item| {
+fn get_item<'a>(
+    tcx: TyCtxt,
+    items: &'a stable_mir::CrateItems,
+    item: (DefKind, &str),
+) -> Option<&'a stable_mir::CrateItem> {
+    items.iter().find(|crate_item| {
         let def_id = rustc_internal::item_def_id(crate_item);
         tcx.def_kind(def_id) == item.0 && tcx.def_path_str(def_id) == item.1
     })
@@ -94,6 +112,10 @@ fn generate_input(path: &str) -> std::io::Result<()> {
         }}
     }}
 
+    pub fn bar(x: i32) -> i32 {{
+        x
+    }}
+
     pub fn foo_bar(x: i32, y: i32) -> i64 {{
         let x_64 = foo::bar(x);
         let y_64 = foo::bar(y);