Properly set the enum variant body expected type

1 files changed, 12 insertions, 6 deletions

diff --git a/crates/hir-ty/src/infer.rs b/crates/hir-ty/src/infer.rs
index 85309d32335..2c6d503def1 100644
--- a/crates/hir-ty/src/infer.rs
+++ b/crates/hir-ty/src/infer.rs
@@ -18,7 +18,9 @@ use std::sync::Arc;
 
 use chalk_ir::{cast::Cast, ConstValue, DebruijnIndex, Mutability, Safety, Scalar, TypeFlags};
 use hir_def::{
+    adt::{ReprData, ReprKind},
     body::Body,
+    builtin_type::BuiltinType,
     data::{ConstData, StaticData},
     expr::{BindingAnnotation, ExprId, PatId},
     lang_item::LangItemTarget,
@@ -67,12 +69,16 @@ pub(crate) fn infer_query(db: &dyn HirDatabase, def: DefWithBodyId) -> Arc<Infer
         DefWithBodyId::ConstId(c) => ctx.collect_const(&db.const_data(c)),
         DefWithBodyId::FunctionId(f) => ctx.collect_fn(f),
         DefWithBodyId::StaticId(s) => ctx.collect_static(&db.static_data(s)),
-        DefWithBodyId::VariantId(_v) => {
-            // db.enum_data(v.parent)
-            // FIXME: This should return the `repr(...)` type of the enum
-            ctx.return_ty = TyBuilder::builtin(hir_def::builtin_type::BuiltinType::Uint(
-                hir_def::builtin_type::BuiltinUint::U32,
-            ));
+        DefWithBodyId::VariantId(v) => {
+            ctx.return_ty = match db.enum_data(v.parent).repr {
+                Some(ReprData { kind: ReprKind::BuiltinInt { builtin, .. }, .. }) => {
+                    TyBuilder::builtin(match builtin {
+                        Either::Left(builtin) => BuiltinType::Int(builtin),
+                        Either::Right(builtin) => BuiltinType::Uint(builtin),
+                    })
+                }
+                _ => TyBuilder::builtin(BuiltinType::Uint(hir_def::builtin_type::BuiltinUint::U32)),
+            };
         }
     }