Add methods to create StableMIR constant

compiler/rustc_smir/src/rustc_smir/context.rs

@@ -23,7 +23,8 @@ use stable_mir::mir::Body;
 use stable_mir::target::{MachineInfo, MachineSize};
 use stable_mir::ty::{
     AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, FieldDef, FnDef, ForeignDef,
-    ForeignItemKind, GenericArgs, LineInfo, PolyFnSig, RigidTy, Span, Ty, TyKind, VariantDef,
+    ForeignItemKind, GenericArgs, LineInfo, PolyFnSig, RigidTy, Span, Ty, TyKind, UintTy,
+    VariantDef,
 };
 use stable_mir::{Crate, CrateDef, CrateItem, CrateNum, DefId, Error, Filename, ItemKind, Symbol};
 use std::cell::RefCell;
@@ -341,15 +342,56 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
             .ok_or_else(|| Error::new(format!("Const `{cnst:?}` cannot be encoded as u64")))
     }
 
-    fn usize_to_const(&self, val: u64) -> Result<Const, Error> {
+    fn try_new_const_zst(&self, ty: Ty) -> Result<Const, Error> {
         let mut tables = self.0.borrow_mut();
-        let ty = tables.tcx.types.usize;
+        let tcx = tables.tcx;
+        let ty_internal = ty.internal(&mut *tables, tcx);
+        let size = tables
+            .tcx
+            .layout_of(ParamEnv::empty().and(ty_internal))
+            .map_err(|err| {
+                Error::new(format!(
+                    "Cannot create a zero-sized constant for type `{ty_internal}`: {err}"
+                ))
+            })?
+            .size;
+        if size.bytes() != 0 {
+            return Err(Error::new(format!(
+                "Cannot create a zero-sized constant for type `{ty_internal}`: \
+                 Type `{ty_internal}` has {} bytes",
+                size.bytes()
+            )));
+        }
+
+        Ok(ty::Const::zero_sized(tables.tcx, ty_internal).stable(&mut *tables))
+    }
+
+    fn new_const_str(&self, value: &str) -> Const {
+        let mut tables = self.0.borrow_mut();
+        let tcx = tables.tcx;
+        let ty = ty::Ty::new_static_str(tcx);
+        let bytes = value.as_bytes();
+        let val_tree = ty::ValTree::from_raw_bytes(tcx, bytes);
+
+        ty::Const::new_value(tcx, val_tree, ty).stable(&mut *tables)
+    }
+
+    fn new_const_bool(&self, value: bool) -> Const {
+        let mut tables = self.0.borrow_mut();
+        ty::Const::from_bool(tables.tcx, value).stable(&mut *tables)
+    }
+
+    fn try_new_const_uint(&self, value: u128, uint_ty: UintTy) -> Result<Const, Error> {
+        let mut tables = self.0.borrow_mut();
+        let tcx = tables.tcx;
+        let ty = ty::Ty::new_uint(tcx, uint_ty.internal(&mut *tables, tcx));
         let size = tables.tcx.layout_of(ParamEnv::empty().and(ty)).unwrap().size;
 
-        let scalar = ScalarInt::try_from_uint(val, size).ok_or_else(|| {
-            Error::new(format!("Value overflow: cannot convert `{val}` to usize."))
+        // We don't use Const::from_bits since it doesn't have any error checking.
+        let scalar = ScalarInt::try_from_uint(value, size).ok_or_else(|| {
+            Error::new(format!("Value overflow: cannot convert `{value}` to `{ty}`."))
         })?;
-        Ok(rustc_middle::ty::Const::new_value(tables.tcx, ValTree::from_scalar_int(scalar), ty)
+        Ok(ty::Const::new_value(tables.tcx, ValTree::from_scalar_int(scalar), ty)
             .stable(&mut *tables))
     }
 
@@ -556,7 +598,9 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         global_alloc: &GlobalAlloc,
     ) -> Option<stable_mir::mir::alloc::AllocId> {
         let mut tables = self.0.borrow_mut();
-        let GlobalAlloc::VTable(ty, trait_ref) = global_alloc else { return None };
+        let GlobalAlloc::VTable(ty, trait_ref) = global_alloc else {
+            return None;
+        };
         let tcx = tables.tcx;
         let alloc_id = tables.tcx.vtable_allocation((
             ty.internal(&mut *tables, tcx),

compiler/stable_mir/src/compiler_interface.rs

@@ -14,7 +14,7 @@ use crate::ty::{
     AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, FieldDef, FnDef, ForeignDef,
     ForeignItemKind, ForeignModule, ForeignModuleDef, GenericArgs, GenericPredicates, Generics,
     ImplDef, ImplTrait, LineInfo, PolyFnSig, RigidTy, Span, TraitDecl, TraitDef, Ty, TyKind,
-    VariantDef,
+    UintTy, VariantDef,
 };
 use crate::{
     mir, Crate, CrateItem, CrateItems, CrateNum, DefId, Error, Filename, ImplTraitDecls, ItemKind,
@@ -101,8 +101,17 @@ pub trait Context {
     /// Evaluate constant as a target usize.
     fn eval_target_usize(&self, cnst: &Const) -> Result<u64, Error>;
 
-    /// Create a target usize constant for the given value.
-    fn usize_to_const(&self, val: u64) -> Result<Const, Error>;
+    /// Create a new zero-sized constant.
+    fn try_new_const_zst(&self, ty: Ty) -> Result<Const, Error>;
+
+    /// Create a new constant that represents the given string value.
+    fn new_const_str(&self, value: &str) -> Const;
+
+    /// Create a new constant that represents the given boolean value.
+    fn new_const_bool(&self, value: bool) -> Const;
+
+    /// Create a new constant that represents the given value.
+    fn try_new_const_uint(&self, value: u128, uint_ty: UintTy) -> Result<Const, Error>;
 
     /// Create a new type from the given kind.
     fn new_rigid_ty(&self, kind: RigidTy) -> Ty;
@@ -199,7 +208,7 @@ pub trait Context {
 
 // A thread local variable that stores a pointer to the tables mapping between TyCtxt
 // datastructures and stable MIR datastructures
-scoped_thread_local! (static TLV: Cell<*const ()>);
+scoped_thread_local!(static TLV: Cell<*const ()>);
 
 pub fn run<F, T>(context: &dyn Context, f: F) -> Result<T, Error>
 where

compiler/stable_mir/src/ty.rs

@@ -128,13 +128,38 @@ impl Const {
 
     /// Creates an interned usize constant.
     fn try_from_target_usize(val: u64) -> Result<Self, Error> {
-        with(|cx| cx.usize_to_const(val))
+        with(|cx| cx.try_new_const_uint(val.into(), UintTy::Usize))
     }
 
     /// Try to evaluate to a target `usize`.
     pub fn eval_target_usize(&self) -> Result<u64, Error> {
         with(|cx| cx.eval_target_usize(self))
     }
+
+    /// Create a constant that represents a new zero-sized constant of type T.
+    /// Fails if the type is not a ZST or if it doesn't have a known size.
+    pub fn try_new_zero_sized(ty: Ty) -> Result<Const, Error> {
+        with(|cx| cx.try_new_const_zst(ty))
+    }
+
+    /// Build a new constant that represents the given string.
+    ///
+    /// Note that there is no guarantee today about duplication of the same constant.
+    /// I.e.: Calling this function multiple times with the same argument may or may not return
+    /// the same allocation.
+    pub fn from_str(value: &str) -> Const {
+        with(|cx| cx.new_const_str(value))
+    }
+
+    /// Build a new constant that represents the given boolean value.
+    pub fn from_bool(value: bool) -> Const {
+        with(|cx| cx.new_const_bool(value))
+    }
+
+    /// Build a new constant that represents the given unsigned integer.
+    pub fn try_from_uint(value: u128, uint_ty: UintTy) -> Result<Const, Error> {
+        with(|cx| cx.try_new_const_uint(value, uint_ty))
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]

tests/ui-fulldeps/stable-mir/check_transform.rs

@@ -0,0 +1,147 @@
+//@ run-pass
+//! Test a few methods to transform StableMIR.
+
+//@ ignore-stage1
+//@ ignore-cross-compile
+//@ ignore-remote
+//@ ignore-windows-gnu mingw has troubles with linking https://github.com/rust-lang/rust/pull/116837
+
+#![feature(rustc_private)]
+#![feature(assert_matches)]
+#![feature(control_flow_enum)]
+#![feature(ascii_char, ascii_char_variants)]
+
+extern crate rustc_hir;
+#[macro_use]
+extern crate rustc_smir;
+extern crate rustc_driver;
+extern crate rustc_interface;
+extern crate stable_mir;
+
+use rustc_smir::rustc_internal;
+use stable_mir::mir::alloc::GlobalAlloc;
+use stable_mir::mir::mono::Instance;
+use stable_mir::mir::{Body, Constant, Operand, Rvalue, StatementKind, TerminatorKind};
+use stable_mir::ty::{Const, ConstantKind};
+use stable_mir::{CrateDef, CrateItems, ItemKind};
+use std::convert::TryFrom;
+use std::io::Write;
+use std::ops::ControlFlow;
+
+const CRATE_NAME: &str = "input";
+
+/// This function uses the Stable MIR APIs to transform the MIR.
+fn test_transform() -> ControlFlow<()> {
+    // Find items in the local crate.
+    let items = stable_mir::all_local_items();
+
+    // Test fn_abi
+    let target_fn = *get_item(&items, (ItemKind::Fn, "dummy")).unwrap();
+    let instance = Instance::try_from(target_fn).unwrap();
+    let body = instance.body().unwrap();
+    check_msg(&body, "oops");
+
+    let new_msg = "new panic message";
+    let new_body = change_panic_msg(body, new_msg);
+    check_msg(&new_body, new_msg);
+
+    ControlFlow::Continue(())
+}
+
+/// Check that the body panic message matches the given message.
+fn check_msg(body: &Body, expected: &str) {
+    let msg = body
+        .blocks
+        .iter()
+        .find_map(|bb| match &bb.terminator.kind {
+            TerminatorKind::Call { args, .. } => {
+                assert_eq!(args.len(), 1, "Expected panic message, but found {args:?}");
+                let msg_const = match &args[0] {
+                    Operand::Constant(msg_const) => msg_const,
+                    Operand::Copy(place) | Operand::Move(place) => {
+                        assert!(place.projection.is_empty());
+                        bb.statements
+                            .iter()
+                            .find_map(|stmt| match &stmt.kind {
+                                StatementKind::Assign(
+                                    destination,
+                                    Rvalue::Use(Operand::Constant(msg_const)),
+                                ) if destination == place => Some(msg_const),
+                                _ => None,
+                            })
+                            .unwrap()
+                    }
+                };
+                let ConstantKind::Allocated(alloc) = msg_const.literal.kind() else {
+                    unreachable!()
+                };
+                assert_eq!(alloc.provenance.ptrs.len(), 1);
+
+                let alloc_prov_id = alloc.provenance.ptrs[0].1 .0;
+                let GlobalAlloc::Memory(val) = GlobalAlloc::from(alloc_prov_id) else {
+                    unreachable!()
+                };
+                let bytes = val.raw_bytes().unwrap();
+                Some(std::str::from_utf8(&bytes).unwrap().to_string())
+            }
+            _ => None,
+        })
+        .expect("Failed to find panic message");
+    assert_eq!(&msg, expected);
+}
+
+/// Modify body to use a different panic message.
+fn change_panic_msg(mut body: Body, new_msg: &str) -> Body {
+    for bb in &mut body.blocks {
+        match &mut bb.terminator.kind {
+            TerminatorKind::Call { args, .. } => {
+                let new_const = Const::from_str(new_msg);
+                args[0] = Operand::Constant(Constant {
+                    literal: new_const,
+                    span: bb.terminator.span,
+                    user_ty: None,
+                });
+            }
+            _ => {}
+        }
+    }
+    body
+}
+
+fn get_item<'a>(
+    items: &'a CrateItems,
+    item: (ItemKind, &str),
+) -> Option<&'a stable_mir::CrateItem> {
+    items.iter().find(|crate_item| (item.0 == crate_item.kind()) && crate_item.name() == item.1)
+}
+
+/// This test will generate and analyze a dummy crate using the stable mir.
+/// For that, it will first write the dummy crate into a file.
+/// Then it will create a `StableMir` using custom arguments and then
+/// it will run the compiler.
+fn main() {
+    let path = "transform_input.rs";
+    generate_input(&path).unwrap();
+    let args = vec![
+        "rustc".to_string(),
+        "--crate-type=lib".to_string(),
+        "--crate-name".to_string(),
+        CRATE_NAME.to_string(),
+        path.to_string(),
+    ];
+    run!(args, test_transform).unwrap();
+}
+
+fn generate_input(path: &str) -> std::io::Result<()> {
+    let mut file = std::fs::File::create(path)?;
+    write!(
+        file,
+        r#"
+        #![feature(panic_internals)]
+        pub fn dummy() {{
+            core::panicking::panic_str("oops");
+        }}
+        "#
+    )?;
+    Ok(())
+}