Rollup merge of #118688 - celinval:smir-rvalue-ty, r=compiler-errors

matthiaskrgr · web-flow · commit 20ed80523a4f · 2023-12-07T21:38:08.000+01:00
Add method to get type of an Rvalue in StableMIR

Provide a method to StableMIR users to retrieve the type of an Rvalue operation. There were two possible implementation:

1. Create the logic inside stable_mir to process the type according to the Rvalue semantics, which duplicates the logic of `rustc_middle::mir::Rvalue::ty()`.
2. Implement the Rvalue translation from StableMIR back to internal representation, invoke the `rustc_middle::mir::Rvalue::ty()`, and translate the return value to StableMIR.

I chose the first one for now since the duplication was fairly small, and the option 2 would require way more work to translate everything back to rustc internal representation. If we eventually add those translations, we could easily swap to the option 2.

```@compiler-errors``` / ```@ouz-a``` Please let me know if you have any strong opinion here.

r? ```@compiler-errors```
diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs
@@ -250,6 +250,13 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         tables.tcx.mk_ty_from_kind(internal_kind).stable(&mut *tables)
     }
 
+    #[allow(rustc::usage_of_qualified_ty)]
+    fn new_box_ty(&self, ty: stable_mir::ty::Ty) -> stable_mir::ty::Ty {
+        let mut tables = self.0.borrow_mut();
+        let inner = ty.internal(&mut *tables);
+        ty::Ty::new_box(tables.tcx, inner).stable(&mut *tables)
+    }
+
     fn def_ty(&self, item: stable_mir::DefId) -> stable_mir::ty::Ty {
         let mut tables = self.0.borrow_mut();
         tables.tcx.type_of(item.internal(&mut *tables)).instantiate_identity().stable(&mut *tables)
@@ -276,6 +283,13 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         tables.types[ty].kind().stable(&mut *tables)
     }
 
+    fn rigid_ty_discriminant_ty(&self, ty: &RigidTy) -> stable_mir::ty::Ty {
+        let mut tables = self.0.borrow_mut();
+        let internal_kind = ty.internal(&mut *tables);
+        let internal_ty = tables.tcx.mk_ty_from_kind(internal_kind);
+        internal_ty.discriminant_ty(tables.tcx).stable(&mut *tables)
+    }
+
     fn instance_body(&self, def: InstanceDef) -> Option<Body> {
         let mut tables = self.0.borrow_mut();
         let instance = tables.instances[def];
@@ -308,9 +322,9 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         matches!(instance.def, ty::InstanceDef::DropGlue(_, None))
     }
 
-    fn mono_instance(&self, item: stable_mir::CrateItem) -> stable_mir::mir::mono::Instance {
+    fn mono_instance(&self, def_id: stable_mir::DefId) -> stable_mir::mir::mono::Instance {
         let mut tables = self.0.borrow_mut();
-        let def_id = tables[item.0];
+        let def_id = tables[def_id];
         Instance::mono(tables.tcx, def_id).stable(&mut *tables)
     }
 
diff --git a/compiler/stable_mir/src/compiler_interface.rs b/compiler/stable_mir/src/compiler_interface.rs
@@ -87,6 +87,9 @@ pub trait Context {
     /// Create a new type from the given kind.
     fn new_rigid_ty(&self, kind: RigidTy) -> Ty;
 
+    /// Create a new box type, `Box<T>`, for the given inner type `T`.
+    fn new_box_ty(&self, ty: Ty) -> Ty;
+
     /// Returns the type of given crate item.
     fn def_ty(&self, item: DefId) -> Ty;
 
@@ -102,6 +105,9 @@ pub trait Context {
     /// Obtain the representation of a type.
     fn ty_kind(&self, ty: Ty) -> TyKind;
 
+    // Get the discriminant Ty for this Ty if there's one.
+    fn rigid_ty_discriminant_ty(&self, ty: &RigidTy) -> Ty;
+
     /// Get the body of an Instance which is already monomorphized.
     fn instance_body(&self, instance: InstanceDef) -> Option<Body>;
 
@@ -119,7 +125,7 @@ pub trait Context {
 
     /// Convert a non-generic crate item into an instance.
     /// This function will panic if the item is generic.
-    fn mono_instance(&self, item: CrateItem) -> Instance;
+    fn mono_instance(&self, def_id: DefId) -> Instance;
 
     /// Item requires monomorphization.
     fn requires_monomorphization(&self, def_id: DefId) -> bool;
diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
@@ -274,6 +274,38 @@ pub enum BinOp {
     Offset,
 }
 
+impl BinOp {
+    /// Return the type of this operation for the given input Ty.
+    /// This function does not perform type checking, and it currently doesn't handle SIMD.
+    pub fn ty(&self, lhs_ty: Ty, rhs_ty: Ty) -> Ty {
+        assert!(lhs_ty.kind().is_primitive());
+        assert!(rhs_ty.kind().is_primitive());
+        match self {
+            BinOp::Add
+            | BinOp::AddUnchecked
+            | BinOp::Sub
+            | BinOp::SubUnchecked
+            | BinOp::Mul
+            | BinOp::MulUnchecked
+            | BinOp::Div
+            | BinOp::Rem
+            | BinOp::BitXor
+            | BinOp::BitAnd
+            | BinOp::BitOr => {
+                assert_eq!(lhs_ty, rhs_ty);
+                lhs_ty
+            }
+            BinOp::Shl | BinOp::ShlUnchecked | BinOp::Shr | BinOp::ShrUnchecked | BinOp::Offset => {
+                lhs_ty
+            }
+            BinOp::Eq | BinOp::Lt | BinOp::Le | BinOp::Ne | BinOp::Ge | BinOp::Gt => {
+                assert_eq!(lhs_ty, rhs_ty);
+                Ty::bool_ty()
+            }
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum UnOp {
     Not,
@@ -475,6 +507,63 @@ pub enum Rvalue {
     Use(Operand),
 }
 
+impl Rvalue {
+    pub fn ty(&self, locals: &[LocalDecl]) -> Result<Ty, Error> {
+        match self {
+            Rvalue::Use(operand) => operand.ty(locals),
+            Rvalue::Repeat(operand, count) => {
+                Ok(Ty::new_array_with_const_len(operand.ty(locals)?, count.clone()))
+            }
+            Rvalue::ThreadLocalRef(did) => Ok(did.ty()),
+            Rvalue::Ref(reg, bk, place) => {
+                let place_ty = place.ty(locals)?;
+                Ok(Ty::new_ref(reg.clone(), place_ty, bk.to_mutable_lossy()))
+            }
+            Rvalue::AddressOf(mutability, place) => {
+                let place_ty = place.ty(locals)?;
+                Ok(Ty::new_ptr(place_ty, *mutability))
+            }
+            Rvalue::Len(..) => Ok(Ty::usize_ty()),
+            Rvalue::Cast(.., ty) => Ok(*ty),
+            Rvalue::BinaryOp(op, lhs, rhs) => {
+                let lhs_ty = lhs.ty(locals)?;
+                let rhs_ty = rhs.ty(locals)?;
+                Ok(op.ty(lhs_ty, rhs_ty))
+            }
+            Rvalue::CheckedBinaryOp(op, lhs, rhs) => {
+                let lhs_ty = lhs.ty(locals)?;
+                let rhs_ty = rhs.ty(locals)?;
+                let ty = op.ty(lhs_ty, rhs_ty);
+                Ok(Ty::new_tuple(&[ty, Ty::bool_ty()]))
+            }
+            Rvalue::UnaryOp(UnOp::Not | UnOp::Neg, operand) => operand.ty(locals),
+            Rvalue::Discriminant(place) => {
+                let place_ty = place.ty(locals)?;
+                place_ty
+                    .kind()
+                    .discriminant_ty()
+                    .ok_or_else(|| error!("Expected a `RigidTy` but found: {place_ty:?}"))
+            }
+            Rvalue::NullaryOp(NullOp::SizeOf | NullOp::AlignOf | NullOp::OffsetOf(..), _) => {
+                Ok(Ty::usize_ty())
+            }
+            Rvalue::Aggregate(ak, ops) => match *ak {
+                AggregateKind::Array(ty) => Ty::try_new_array(ty, ops.len() as u64),
+                AggregateKind::Tuple => Ok(Ty::new_tuple(
+                    &ops.iter().map(|op| op.ty(locals)).collect::<Result<Vec<_>, _>>()?,
+                )),
+                AggregateKind::Adt(def, _, ref args, _, _) => Ok(def.ty_with_args(args)),
+                AggregateKind::Closure(def, ref args) => Ok(Ty::new_closure(def, args.clone())),
+                AggregateKind::Coroutine(def, ref args, mov) => {
+                    Ok(Ty::new_coroutine(def, args.clone(), mov))
+                }
+            },
+            Rvalue::ShallowInitBox(_, ty) => Ok(Ty::new_box(*ty)),
+            Rvalue::CopyForDeref(place) => place.ty(locals),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum AggregateKind {
     Array(Ty),
@@ -725,6 +814,17 @@ pub enum BorrowKind {
     },
 }
 
+impl BorrowKind {
+    pub fn to_mutable_lossy(self) -> Mutability {
+        match self {
+            BorrowKind::Mut { .. } => Mutability::Mut,
+            BorrowKind::Shared => Mutability::Not,
+            // FIXME: There's no type corresponding to a shallow borrow, so use `&` as an approximation.
+            BorrowKind::Fake => Mutability::Not,
+        }
+    }
+}
+
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
 pub enum MutBorrowKind {
     Default,
diff --git a/compiler/stable_mir/src/mir/mono.rs b/compiler/stable_mir/src/mir/mono.rs
@@ -150,8 +150,9 @@ impl TryFrom<CrateItem> for Instance {
 
     fn try_from(item: CrateItem) -> Result<Self, Self::Error> {
         with(|context| {
-            if !context.requires_monomorphization(item.0) {
-                Ok(context.mono_instance(item))
+            let def_id = item.def_id();
+            if !context.requires_monomorphization(def_id) {
+                Ok(context.mono_instance(def_id))
             } else {
                 Err(Error::new("Item requires monomorphization".to_string()))
             }
@@ -219,6 +220,21 @@ impl TryFrom<CrateItem> for StaticDef {
     }
 }
 
+impl TryFrom<Instance> for StaticDef {
+    type Error = crate::Error;
+
+    fn try_from(value: Instance) -> Result<Self, Self::Error> {
+        StaticDef::try_from(CrateItem::try_from(value)?)
+    }
+}
+
+impl From<StaticDef> for Instance {
+    fn from(value: StaticDef) -> Self {
+        // A static definition should always be convertible to an instance.
+        with(|cx| cx.mono_instance(value.def_id()))
+    }
+}
+
 impl StaticDef {
     /// Return the type of this static definition.
     pub fn ty(&self) -> Ty {
diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
@@ -31,15 +31,50 @@ impl Ty {
         Ok(Ty::from_rigid_kind(RigidTy::Array(elem_ty, Const::try_from_target_usize(size)?)))
     }
 
+    /// Create a new array type from Const length.
+    pub fn new_array_with_const_len(elem_ty: Ty, len: Const) -> Ty {
+        Ty::from_rigid_kind(RigidTy::Array(elem_ty, len))
+    }
+
     /// Create a new pointer type.
     pub fn new_ptr(pointee_ty: Ty, mutability: Mutability) -> Ty {
         Ty::from_rigid_kind(RigidTy::RawPtr(pointee_ty, mutability))
     }
 
+    /// Create a new reference type.
+    pub fn new_ref(reg: Region, pointee_ty: Ty, mutability: Mutability) -> Ty {
+        Ty::from_rigid_kind(RigidTy::Ref(reg, pointee_ty, mutability))
+    }
+
+    /// Create a new pointer type.
+    pub fn new_tuple(tys: &[Ty]) -> Ty {
+        Ty::from_rigid_kind(RigidTy::Tuple(Vec::from(tys)))
+    }
+
+    /// Create a new closure type.
+    pub fn new_closure(def: ClosureDef, args: GenericArgs) -> Ty {
+        Ty::from_rigid_kind(RigidTy::Closure(def, args))
+    }
+
+    /// Create a new coroutine type.
+    pub fn new_coroutine(def: CoroutineDef, args: GenericArgs, mov: Movability) -> Ty {
+        Ty::from_rigid_kind(RigidTy::Coroutine(def, args, mov))
+    }
+
+    /// Create a new box type that represents `Box<T>`, for the given inner type `T`.
+    pub fn new_box(inner_ty: Ty) -> Ty {
+        with(|cx| cx.new_box_ty(inner_ty))
+    }
+
     /// Create a type representing `usize`.
     pub fn usize_ty() -> Ty {
         Ty::from_rigid_kind(RigidTy::Uint(UintTy::Usize))
     }
+
+    /// Create a type representing `bool`.
+    pub fn bool_ty() -> Ty {
+        Ty::from_rigid_kind(RigidTy::Bool)
+    }
 }
 
 impl Ty {
@@ -209,6 +244,19 @@ impl TyKind {
         matches!(self, TyKind::RigidTy(RigidTy::FnPtr(..)))
     }
 
+    pub fn is_primitive(&self) -> bool {
+        matches!(
+            self,
+            TyKind::RigidTy(
+                RigidTy::Bool
+                    | RigidTy::Char
+                    | RigidTy::Int(_)
+                    | RigidTy::Uint(_)
+                    | RigidTy::Float(_)
+            )
+        )
+    }
+
     pub fn trait_principal(&self) -> Option<Binder<ExistentialTraitRef>> {
         if let TyKind::RigidTy(RigidTy::Dynamic(predicates, _, _)) = self {
             if let Some(Binder { value: ExistentialPredicate::Trait(trait_ref), bound_vars }) =
@@ -251,13 +299,19 @@ impl TyKind {
     }
 
     /// Get the function signature for function like types (Fn, FnPtr, Closure, Coroutine)
+    /// FIXME(closure)
     pub fn fn_sig(&self) -> Option<PolyFnSig> {
         match self {
             TyKind::RigidTy(RigidTy::FnDef(def, args)) => Some(with(|cx| cx.fn_sig(*def, args))),
             TyKind::RigidTy(RigidTy::FnPtr(sig)) => Some(sig.clone()),
             _ => None,
         }
     }
+
+    /// Get the discriminant type for this type.
+    pub fn discriminant_ty(&self) -> Option<Ty> {
+        self.rigid().map(|ty| with(|cx| cx.rigid_ty_discriminant_ty(ty)))
+    }
 }
 
 pub struct TypeAndMut {
@@ -289,6 +343,13 @@ pub enum RigidTy {
     CoroutineWitness(CoroutineWitnessDef, GenericArgs),
 }
 
+impl RigidTy {
+    /// Get the discriminant type for this type.
+    pub fn discriminant_ty(&self) -> Ty {
+        with(|cx| cx.rigid_ty_discriminant_ty(self))
+    }
+}
+
 impl From<RigidTy> for TyKind {
     fn from(value: RigidTy) -> Self {
         TyKind::RigidTy(value)
