Add sanity check to BinOp::ty()

rust-lang · Dec 7, 2023 · 4616b9f · 4616b9f
1 parent 77d7e44
commit 4616b9f
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 8 deletions.
diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
@@ -277,7 +277,9 @@ pub enum BinOp {
 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 {
+    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
@@ -289,13 +291,17 @@ impl BinOp {
             | BinOp::Rem
             | BinOp::BitXor
             | BinOp::BitAnd
-            | BinOp::BitOr
-            | BinOp::Shl
-            | BinOp::ShlUnchecked
-            | BinOp::Shr
-            | BinOp::ShrUnchecked
-            | BinOp::Offset => lhs_ty,
-            BinOp::Eq | BinOp::Lt | BinOp::Le | BinOp::Ne | BinOp::Ge | BinOp::Gt => Ty::bool_ty(),
+            | 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()
+            }
         }
     }
 }

diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
@@ -244,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 }) =