Repository: rust-lang/rust-analyzer
Pull request #13269 (2022-09-20 15:16:06)

Properly set the enum variant body type from the repr attribute
Follow up for https://github.com/rust-lang/rust-analyzer/pull/12966, fixes some type inference problems which would cause type mismatches otherwise


Patch:
diff --git a/crates/hir-def/src/adt.rs b/crates/hir-def/src/adt.rs
index 277135d6dc..14f8629056 100644
--- a/crates/hir-def/src/adt.rs
+++ b/crates/hir-def/src/adt.rs
@@ -1,6 +1,6 @@
 //! Defines hir-level representation of structs, enums and unions
 
-use std::sync::Arc;
+use std::{num::NonZeroU32, sync::Arc};
 
 use base_db::CrateId;
 use either::Either;
@@ -14,6 +14,7 @@ use tt::{Delimiter, DelimiterKind, Leaf, Subtree, TokenTree};
 
 use crate::{
     body::{CfgExpander, LowerCtx},
+    builtin_type::{BuiltinInt, BuiltinUint},
     db::DefDatabase,
     intern::Interned,
     item_tree::{AttrOwner, Field, Fields, ItemTree, ModItem, RawVisibilityId},
@@ -31,7 +32,7 @@ use cfg::CfgOptions;
 pub struct StructData {
     pub name: Name,
     pub variant_data: Arc<VariantData>,
-    pub repr: Option<ReprKind>,
+    pub repr: Option<ReprData>,
     pub visibility: RawVisibility,
 }
 
@@ -39,6 +40,7 @@ pub struct StructData {
 pub struct EnumData {
     pub name: Name,
     pub variants: Arena<EnumVariantData>,
+    pub repr: Option<ReprData>,
     pub visibility: RawVisibility,
 }
 
@@ -63,10 +65,19 @@ pub struct FieldData {
     pub visibility: RawVisibility,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Copy, Debug, Clone, PartialEq, Eq)]
 pub enum ReprKind {
-    Packed,
-    Other,
+    C,
+    BuiltinInt { builtin: Either<BuiltinInt, BuiltinUint>, is_c: bool },
+    Transparent,
+    Default,
+}
+
+#[derive(Copy, Debug, Clone, PartialEq, Eq)]
+pub struct ReprData {
+    pub kind: ReprKind,
+    pub packed: bool,
+    pub align: Option<NonZeroU32>,
 }
 
 fn repr_from_value(
@@ -74,21 +85,60 @@ fn repr_from_value(
     krate: CrateId,
     item_tree: &ItemTree,
     of: AttrOwner,
-) -> Option<ReprKind> {
+) -> Option<ReprData> {
     item_tree.attrs(db, krate, of).by_key("repr").tt_values().find_map(parse_repr_tt)
 }
 
-fn parse_repr_tt(tt: &Subtree) -> Option<ReprKind> {
+fn parse_repr_tt(tt: &Subtree) -> Option<ReprData> {
     match tt.delimiter {
         Some(Delimiter { kind: DelimiterKind::Parenthesis, .. }) => {}
         _ => return None,
     }
 
-    let mut it = tt.token_trees.iter();
-    match it.next()? {
-        TokenTree::Leaf(Leaf::Ident(ident)) if ident.text == "packed" => Some(ReprKind::Packed),
-        _ => Some(ReprKind::Other),
+    let mut data = ReprData { kind: ReprKind::Default, packed: false, align: None };
+
+    let mut tts = tt.token_trees.iter().peekable();
+    while let Some(tt) = tts.next() {
+        if let TokenTree::Leaf(Leaf::Ident(ident)) = tt {
+            match &*ident.text {
+                "packed" => {
+                    data.packed = true;
+                    if let Some(TokenTree::Subtree(_)) = tts.peek() {
+                        tts.next();
+                    }
+                }
+                "align" => {
+                    if let Some(TokenTree::Subtree(tt)) = tts.peek() {
+                        tts.next();
+                        if let Some(TokenTree::Leaf(Leaf::Literal(lit))) = tt.token_trees.first() {
+                            if let Ok(align) = lit.text.parse() {
+                                data.align = Some(align);
+                            }
+                        }
+                    }
+                }
+                "C" => {
+                    if let ReprKind::BuiltinInt { is_c, .. } = &mut data.kind {
+                        *is_c = true;
+                    } else {
+                        data.kind = ReprKind::C;
+                    }
+                }
+                "transparent" => data.kind = ReprKind::Transparent,
+                repr => {
+                    let is_c = matches!(data.kind, ReprKind::C);
+                    if let Some(builtin) = BuiltinInt::from_suffix(repr)
+                        .map(Either::Left)
+                        .or_else(|| BuiltinUint::from_suffix(repr).map(Either::Right))
+                    {
+                        data.kind = ReprKind::BuiltinInt { builtin, is_c };
+                    }
+                }
+            }
+        }
     }
+
+    Some(data)
 }
 
 impl StructData {
@@ -108,6 +158,7 @@ impl StructData {
             visibility: item_tree[strukt.visibility].clone(),
         })
     }
+
     pub(crate) fn union_data_query(db: &dyn DefDatabase, id: UnionId) -> Arc<StructData> {
         let loc = id.lookup(db);
         let krate = loc.container.krate;
@@ -133,6 +184,7 @@ impl EnumData {
         let krate = loc.container.krate;
         let item_tree = loc.id.item_tree(db);
         let cfg_options = db.crate_graph()[krate].cfg_options.clone();
+        let repr = repr_from_value(db, krate, &item_tree, ModItem::from(loc.id.value).into());
 
         let enum_ = &item_tree[loc.id.value];
         let mut variants = Arena::new();
@@ -158,6 +210,7 @@ impl EnumData {
         Arc::new(EnumData {
             name: enum_.name.clone(),
             variants,
+            repr,
             visibility: item_tree[enum_.visibility].clone(),
         })
     }
@@ -166,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::Left(BuiltinInt::Isize),
+        }
+    }
 }
 
 impl HasChildSource<LocalEnumVariantId> for EnumId {
diff --git a/crates/hir-def/src/lib.rs b/crates/hir-def/src/lib.rs
index 4c44840e86..5c7aa72349 100644
--- a/crates/hir-def/src/lib.rs
+++ b/crates/hir-def/src/lib.rs
@@ -479,7 +479,6 @@ pub enum DefWithBodyId {
 
 impl_from!(FunctionId, ConstId, StaticId for DefWithBodyId);
 
-// FIXME: Rename EnumVariantId to VariantId so that the macro above can be used
 impl From<EnumVariantId> for DefWithBodyId {
     fn from(id: EnumVariantId) -> Self {
         DefWithBodyId::VariantId(id)
diff --git a/crates/hir-ty/src/consteval.rs b/crates/hir-ty/src/consteval.rs
index 965bf3f2c5..2c0c6e0b83 100644
--- a/crates/hir-ty/src/consteval.rs
+++ b/crates/hir-ty/src/consteval.rs
@@ -153,7 +153,7 @@ fn is_valid(scalar: &Scalar, value: i128) -> bool {
     }
 }
 
-fn get_name(variant: EnumVariantId, ctx: &mut ConstEvalCtx<'_>) -> String {
+fn get_name(ctx: &mut ConstEvalCtx<'_>, variant: EnumVariantId) -> String {
     let loc = variant.parent.lookup(ctx.db.upcast());
     let children = variant.parent.child_source(ctx.db.upcast());
     let item_tree = loc.id.item_tree(ctx.db.upcast());
@@ -167,20 +167,24 @@ pub fn eval_const(
     expr_id: ExprId,
     ctx: &mut ConstEvalCtx<'_>,
 ) -> Result<ComputedExpr, ConstEvalError> {
+    let u128_to_i128 = |it: u128| -> Result<i128, ConstEvalError> {
+        it.try_into().map_err(|_| ConstEvalError::NotSupported("u128 is too big"))
+    };
+
     let expr = &ctx.exprs[expr_id];
     match expr {
         Expr::Missing => match ctx.owner {
+            // evaluate the implicit variant index of an enum variant without expression
+            // FIXME: This should return the type of the enum representation
             DefWithBodyId::VariantId(variant) => {
                 let prev_idx: u32 = variant.local_id.into_raw().into();
-                let prev_idx = prev_idx.checked_sub(1).map(|idx| Idx::from_raw(RawIdx::from(idx)));
+                let prev_idx = prev_idx.checked_sub(1).map(RawIdx::from).map(Idx::from_raw);
                 let value = match prev_idx {
-                    Some(prev) => {
-                        let prev_variant = EnumVariantId { local_id: prev, ..variant };
+                    Some(local_id) => {
+                        let prev_variant = EnumVariantId { local_id, parent: variant.parent };
                         1 + match ctx.db.const_eval_variant(prev_variant)? {
                             ComputedExpr::Literal(Literal::Int(v, _)) => v,
-                            ComputedExpr::Literal(Literal::Uint(v, _)) => v
-                                .try_into()
-                                .map_err(|_| ConstEvalError::NotSupported("too big u128"))?,
+                            ComputedExpr::Literal(Literal::Uint(v, _)) => u128_to_i128(v)?,
                             _ => {
                                 return Err(ConstEvalError::NotSupported(
                                     "Enum can't contain this kind of value",
@@ -206,9 +210,7 @@ pub fn eval_const(
                             return Ok(ComputedExpr::Literal(Literal::Bool(!b)))
                         }
                         ComputedExpr::Literal(Literal::Int(v, _)) => v,
-                        ComputedExpr::Literal(Literal::Uint(v, _)) => v
-                            .try_into()
-                            .map_err(|_| ConstEvalError::NotSupported("too big u128"))?,
+                        ComputedExpr::Literal(Literal::Uint(v, _)) => u128_to_i128(v)?,
                         _ => return Err(ConstEvalError::NotSupported("this kind of operator")),
                     };
                     let r = match ty.kind(Interner) {
@@ -237,9 +239,7 @@ pub fn eval_const(
                 hir_def::expr::UnaryOp::Neg => {
                     let v = match ev {
                         ComputedExpr::Literal(Literal::Int(v, _)) => v,
-                        ComputedExpr::Literal(Literal::Uint(v, _)) => v
-                            .try_into()
-                            .map_err(|_| ConstEvalError::NotSupported("too big u128"))?,
+                        ComputedExpr::Literal(Literal::Uint(v, _)) => u128_to_i128(v)?,
                         _ => return Err(ConstEvalError::NotSupported("this kind of operator")),
                     };
                     Ok(ComputedExpr::Literal(Literal::Int(
@@ -258,16 +258,12 @@ pub fn eval_const(
             let op = op.ok_or(ConstEvalError::IncompleteExpr)?;
             let v1 = match lhs {
                 ComputedExpr::Literal(Literal::Int(v, _)) => v,
-                ComputedExpr::Literal(Literal::Uint(v, _)) => {
-                    v.try_into().map_err(|_| ConstEvalError::NotSupported("too big u128"))?
-                }
+                ComputedExpr::Literal(Literal::Uint(v, _)) => u128_to_i128(v)?,
                 _ => return Err(ConstEvalError::NotSupported("this kind of operator")),
             };
             let v2 = match rhs {
                 ComputedExpr::Literal(Literal::Int(v, _)) => v,
-                ComputedExpr::Literal(Literal::Uint(v, _)) => {
-                    v.try_into().map_err(|_| ConstEvalError::NotSupported("too big u128"))?
-                }
+                ComputedExpr::Literal(Literal::Uint(v, _)) => u128_to_i128(v)?,
                 _ => return Err(ConstEvalError::NotSupported("this kind of operator")),
             };
             match op {
@@ -380,7 +376,7 @@ pub fn eval_const(
                 }
                 ValueNs::EnumVariantId(id) => match ctx.db.const_eval_variant(id)? {
                     ComputedExpr::Literal(lit) => {
-                        Ok(ComputedExpr::Enum(get_name(id, ctx), id, lit))
+                        Ok(ComputedExpr::Enum(get_name(ctx, id), id, lit))
                     }
                     _ => Err(ConstEvalError::NotSupported(
                         "Enums can't evalute to anything but numbers",
@@ -389,6 +385,7 @@ pub fn eval_const(
                 _ => Err(ConstEvalError::NotSupported("path that are not const or local")),
             }
         }
+        // FIXME: Handle the cast target
         &Expr::Cast { expr, .. } => match eval_const(expr, ctx)? {
             ComputedExpr::Enum(_, _, lit) => Ok(ComputedExpr::Literal(lit)),
             _ => Err(ConstEvalError::NotSupported("Can't cast these types")),
@@ -463,7 +460,7 @@ pub(crate) fn const_eval_recover(
     Err(ConstEvalError::Loop)
 }
 
-pub(crate) fn const_eval_recover_variant(
+pub(crate) fn const_eval_variant_recover(
     _: &dyn HirDatabase,
     _: &[String],
     _: &EnumVariantId,
@@ -471,7 +468,7 @@ pub(crate) fn const_eval_recover_variant(
     Err(ConstEvalError::Loop)
 }
 
-pub(crate) fn const_eval_query(
+pub(crate) fn const_eval_variant_query(
     db: &dyn HirDatabase,
     const_id: ConstId,
 ) -> Result<ComputedExpr, ConstEvalError> {
diff --git a/crates/hir-ty/src/db.rs b/crates/hir-ty/src/db.rs
index e16530ecc1..72abcc2b4b 100644
--- a/crates/hir-ty/src/db.rs
+++ b/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;
 
@@ -44,12 +43,12 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
     #[salsa::invoke(crate::lower::const_param_ty_query)]
     fn const_param_ty(&self, def: ConstParamId) -> Ty;
 
-    #[salsa::invoke(crate::consteval::const_eval_query)]
+    #[salsa::invoke(crate::consteval::const_eval_variant_query)]
     #[salsa::cycle(crate::consteval::const_eval_recover)]
     fn const_eval(&self, def: ConstId) -> Result<ComputedExpr, ConstEvalError>;
 
     #[salsa::invoke(crate::consteval::const_eval_query_variant)]
-    #[salsa::cycle(crate::consteval::const_eval_recover_variant)]
+    #[salsa::cycle(crate::consteval::const_eval_variant_recover)]
     fn const_eval_variant(&self, def: EnumVariantId) -> Result<ComputedExpr, ConstEvalError>;
 
     #[salsa::invoke(crate::lower::impl_trait_query)]
@@ -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)
diff --git a/crates/hir-ty/src/infer.rs b/crates/hir-ty/src/infer.rs
index 7d1c982075..daf1e2f0c6 100644
--- a/crates/hir-ty/src/infer.rs
+++ b/crates/hir-ty/src/infer.rs
@@ -19,6 +19,7 @@ use std::sync::Arc;
 use chalk_ir::{cast::Cast, ConstValue, DebruijnIndex, Mutability, Safety, Scalar, TypeFlags};
 use hir_def::{
     body::Body,
+    builtin_type::BuiltinType,
     data::{ConstData, StaticData},
     expr::{BindingAnnotation, ExprId, PatId},
     lang_item::LangItemTarget,
@@ -68,8 +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) => {
-            // FIXME: This should return the `repr(...)` type of the enum
-            ctx.return_ty = TyBuilder::def_ty(db, v.parent.into()).fill_with_unknown().build()
+            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),
+            });
         }
     }
 
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 389e07db33..1c48d2ff08 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -39,7 +39,7 @@ use arrayvec::ArrayVec;
 use base_db::{CrateDisplayName, CrateId, CrateOrigin, Edition, FileId, ProcMacroKind};
 use either::Either;
 use hir_def::{
-    adt::{ReprKind, VariantData},
+    adt::{ReprData, VariantData},
     body::{BodyDiagnostic, SyntheticSyntax},
     expr::{BindingAnnotation, LabelId, Pat, PatId},
     generics::{TypeOrConstParamData, TypeParamProvenance},
@@ -874,7 +874,7 @@ impl Struct {
         Type::from_def(db, self.id)
     }
 
-    pub fn repr(self, db: &dyn HirDatabase) -> Option<ReprKind> {
+    pub fn repr(self, db: &dyn HirDatabase) -> Option<ReprData> {
         db.struct_data(self.id).repr.clone()
     }
 
@@ -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),
         }
     }
 
@@ -2964,7 +2975,7 @@ impl Type {
 
         let adt = adt_id.into();
         match adt {
-            Adt::Struct(s) => matches!(s.repr(db), Some(ReprKind::Packed)),
+            Adt::Struct(s) => matches!(s.repr(db), Some(ReprData { packed: true, .. })),
             _ => false,
         }
     }


Test patch:
diff --git a/crates/hir-ty/src/tests.rs b/crates/hir-ty/src/tests.rs
index be5ece9c5c..ebbc541014 100644
--- a/crates/hir-ty/src/tests.rs
+++ b/crates/hir-ty/src/tests.rs
@@ -461,6 +461,18 @@ fn visit_module(
                     let body = db.body(def);
                     visit_body(db, &body, cb);
                 }
+                ModuleDefId::AdtId(hir_def::AdtId::EnumId(it)) => {
+                    db.enum_data(it)
+                        .variants
+                        .iter()
+                        .map(|(id, _)| hir_def::EnumVariantId { parent: it, local_id: id })
+                        .for_each(|it| {
+                            let def = it.into();
+                            cb(def);
+                            let body = db.body(def);
+                            visit_body(db, &body, cb);
+                        });
+                }
                 ModuleDefId::TraitId(it) => {
                     let trait_data = db.trait_data(it);
                     for &(_, item) in trait_data.items.iter() {
diff --git a/crates/hir-ty/src/tests/simple.rs b/crates/hir-ty/src/tests/simple.rs
index 4ea103e5d9..1757327b92 100644
--- a/crates/hir-ty/src/tests/simple.rs
+++ b/crates/hir-ty/src/tests/simple.rs
@@ -1693,16 +1693,16 @@ fn infer_type_param() {
 fn infer_const() {
     check_infer(
         r#"
-        struct Foo;
-        impl Foo { const ASSOC_CONST: u32 = 0; }
-        const GLOBAL_CONST: u32 = 101;
-        fn test() {
-            const LOCAL_CONST: u32 = 99;
-            let x = LOCAL_CONST;
-            let z = GLOBAL_CONST;
-            let id = Foo::ASSOC_CONST;
-        }
-        "#,
+struct Foo;
+impl Foo { const ASSOC_CONST: u32 = 0; }
+const GLOBAL_CONST: u32 = 101;
+fn test() {
+    const LOCAL_CONST: u32 = 99;
+    let x = LOCAL_CONST;
+    let z = GLOBAL_CONST;
+    let id = Foo::ASSOC_CONST;
+}
+"#,
         expect![[r#"
             48..49 '0': u32
             79..82 '101': u32
@@ -1722,17 +1722,17 @@ fn infer_const() {
 fn infer_static() {
     check_infer(
         r#"
-        static GLOBAL_STATIC: u32 = 101;
-        static mut GLOBAL_STATIC_MUT: u32 = 101;
-        fn test() {
-            static LOCAL_STATIC: u32 = 99;
-            static mut LOCAL_STATIC_MUT: u32 = 99;
-            let x = LOCAL_STATIC;
-            let y = LOCAL_STATIC_MUT;
-            let z = GLOBAL_STATIC;
-            let w = GLOBAL_STATIC_MUT;
-        }
-        "#,
+static GLOBAL_STATIC: u32 = 101;
+static mut GLOBAL_STATIC_MUT: u32 = 101;
+fn test() {
+    static LOCAL_STATIC: u32 = 99;
+    static mut LOCAL_STATIC_MUT: u32 = 99;
+    let x = LOCAL_STATIC;
+    let y = LOCAL_STATIC_MUT;
+    let z = GLOBAL_STATIC;
+    let w = GLOBAL_STATIC_MUT;
+}
+"#,
         expect![[r#"
             28..31 '101': u32
             69..72 '101': u32
@@ -1751,6 +1751,41 @@ fn infer_static() {
     );
 }
 
+#[test]
+fn infer_enum_variant() {
+    check_infer(
+        r#"
+enum Foo {
+    A = 15,
+    B = Foo::A as isize + 1
+}
+"#,
+        expect![[r#"
+            19..21 '15': isize
+            31..37 'Foo::A': Foo
+            31..46 'Foo::A as isize': isize
+            31..50 'Foo::A...ze + 1': isize
+            49..50 '1': isize
+        "#]],
+    );
+    check_infer(
+        r#"
+#[repr(u32)]
+enum Foo {
+    A = 15,
+    B = Foo::A as u32 + 1
+}
+"#,
+        expect![[r#"
+            32..34 '15': u32
+            44..50 'Foo::A': Foo
+            44..57 'Foo::A as u32': u32
+            44..61 'Foo::A...32 + 1': u32
+            60..61 '1': u32
+        "#]],
+    );
+}
+
 #[test]
 fn shadowing_primitive() {
     check_types(