Fix using incorrect type for variants in DefWithBody::body_type

Author: Veykril

crates/hir-def/src/adt.rs

@@ -219,6 +219,13 @@ impl EnumData {
         let (id, _) = self.variants.iter().find(|(_id, data)| &data.name == name)?;
         Some(id)
     }
+
+    pub fn variant_body_type(&self) -> Either<BuiltinInt, BuiltinUint> {
+        match self.repr {
+            Some(ReprData { kind: ReprKind::BuiltinInt { builtin, .. }, .. }) => builtin,
+            _ => Either::Right(BuiltinUint::U32),
+        }
+    }
 }
 
 impl HasChildSource<LocalEnumVariantId> for EnumId {

crates/hir-ty/src/db.rs

@@ -7,8 +7,7 @@ use arrayvec::ArrayVec;
 use base_db::{impl_intern_key, salsa, CrateId, Upcast};
 use hir_def::{
     db::DefDatabase, expr::ExprId, BlockId, ConstId, ConstParamId, DefWithBodyId, EnumVariantId,
-    FunctionId, GenericDefId, ImplId, LifetimeParamId, LocalFieldId, Lookup, TypeOrConstParamId,
-    VariantId,
+    FunctionId, GenericDefId, ImplId, LifetimeParamId, LocalFieldId, TypeOrConstParamId, VariantId,
 };
 use la_arena::ArenaMap;
 
@@ -194,11 +193,7 @@ fn infer_wait(db: &dyn HirDatabase, def: DefWithBodyId) -> Arc<InferenceResult>
             db.const_data(it).name.clone().unwrap_or_else(Name::missing).to_string()
         }
         DefWithBodyId::VariantId(it) => {
-            let up_db: &dyn DefDatabase = db.upcast();
-            let loc = it.parent.lookup(up_db);
-            let item_tree = loc.id.item_tree(up_db);
-            let konst = &item_tree[loc.id.value];
-            konst.name.to_string()
+            db.enum_data(it.parent).variants[it.local_id].name.to_string()
         }
     });
     db.infer_query(def)

crates/hir-ty/src/infer.rs

@@ -18,7 +18,6 @@ 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},
@@ -70,15 +69,10 @@ pub(crate) fn infer_query(db: &dyn HirDatabase, def: DefWithBodyId) -> Arc<Infer
         DefWithBodyId::FunctionId(f) => ctx.collect_fn(f),
         DefWithBodyId::StaticId(s) => ctx.collect_static(&db.static_data(s)),
         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)),
-            };
+            ctx.return_ty = TyBuilder::builtin(match db.enum_data(v.parent).variant_body_type() {
+                Either::Left(builtin) => BuiltinType::Int(builtin),
+                Either::Right(builtin) => BuiltinType::Uint(builtin),
+            });
         }
     }
 

crates/hir/src/lib.rs

@@ -953,6 +953,17 @@ impl Enum {
         Type::from_def(db, self.id)
     }
 
+    /// The type of the enum variant bodies.
+    pub fn variant_body_ty(self, db: &dyn HirDatabase) -> Type {
+        Type::new_for_crate(
+            self.id.lookup(db.upcast()).container.krate(),
+            TyBuilder::builtin(match db.enum_data(self.id).variant_body_type() {
+                Either::Left(builtin) => hir_def::builtin_type::BuiltinType::Int(builtin),
+                Either::Right(builtin) => hir_def::builtin_type::BuiltinType::Uint(builtin),
+            }),
+        )
+    }
+
     pub fn is_data_carrying(self, db: &dyn HirDatabase) -> bool {
         self.variants(db).iter().any(|v| !matches!(v.kind(db), StructKind::Unit))
     }
@@ -1176,7 +1187,7 @@ impl DefWithBody {
             DefWithBody::Function(it) => it.ret_type(db),
             DefWithBody::Static(it) => it.ty(db),
             DefWithBody::Const(it) => it.ty(db),
-            DefWithBody::Variant(it) => it.parent.ty(db),
+            DefWithBody::Variant(it) => it.parent.variant_body_ty(db),
         }
     }