From 4616b9fd1d99e3f6d26858f32f835e60ddc3aaa3 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Wed, 6 Dec 2023 22:13:07 -0800
Subject: [PATCH] Add sanity check to `BinOp::ty()`

---
 compiler/stable_mir/src/mir/body.rs | 22 ++++++++++++++--------
 compiler/stable_mir/src/ty.rs       | 13 +++++++++++++
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
index abc9769ea9d..90bd7aa7d18 100644
--- 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
index c7405c7f582..c922264f8a3 100644
--- 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 }) =