Handle return types for in type const bodies

Author: HKalbasi

crates/hir-def/src/db.rs

@@ -25,9 +25,9 @@ use crate::{
     AnonymousConstId, AttrDefId, BlockId, BlockLoc, ConstId, ConstLoc, DefWithBodyId, EnumId,
     EnumLoc, ExternBlockId, ExternBlockLoc, FunctionId, FunctionLoc, GenericDefId, ImplId, ImplLoc,
     InTypeConstId, LocalEnumVariantId, LocalFieldId, Macro2Id, Macro2Loc, MacroRulesId,
-    MacroRulesLoc, ProcMacroId, ProcMacroLoc, StaticId, StaticLoc, StructId, StructLoc,
-    TraitAliasId, TraitAliasLoc, TraitId, TraitLoc, TypeAliasId, TypeAliasLoc, TypeOwnerId,
-    UnionId, UnionLoc, VariantId,
+    MacroRulesLoc, OpaqueInternableThing, ProcMacroId, ProcMacroLoc, StaticId, StaticLoc, StructId,
+    StructLoc, TraitAliasId, TraitAliasLoc, TraitId, TraitLoc, TypeAliasId, TypeAliasLoc,
+    TypeOwnerId, UnionId, UnionLoc, VariantId,
 };
 
 #[salsa::query_group(InternDatabaseStorage)]
@@ -65,7 +65,10 @@ pub trait InternDatabase: SourceDatabase {
     #[salsa::interned]
     fn intern_anonymous_const(&self, id: (DefWithBodyId, ExprId)) -> AnonymousConstId;
     #[salsa::interned]
-    fn intern_in_type_const(&self, id: (AstId<ast::Expr>, TypeOwnerId)) -> InTypeConstId;
+    fn intern_in_type_const(
+        &self,
+        id: (AstId<ast::Expr>, TypeOwnerId, Box<dyn OpaqueInternableThing>),
+    ) -> InTypeConstId;
 }
 
 #[salsa::query_group(DefDatabaseStorage)]

crates/hir-def/src/lib.rs

@@ -57,7 +57,10 @@ mod test_db;
 mod macro_expansion_tests;
 mod pretty;
 
-use std::hash::{Hash, Hasher};
+use std::{
+    hash::{Hash, Hasher},
+    panic::{RefUnwindSafe, UnwindSafe},
+};
 
 use base_db::{
     impl_intern_key,
@@ -501,9 +504,40 @@ impl TypeOwnerId {
 
 impl_from!(ModuleId, DefWithBodyId(FunctionId, ConstId, StaticId), GenericDefId(AdtId, TypeAliasId, ImplId) for TypeOwnerId);
 
+/// A thing that we want to store in interned ids, but we don't know its type in `hir-def`
+pub trait OpaqueInternableThing:
+    std::any::Any + std::fmt::Debug + Sync + Send + UnwindSafe + RefUnwindSafe
+{
+    fn as_any(&self) -> &dyn std::any::Any;
+    fn box_any(&self) -> Box<dyn std::any::Any>;
+    fn dyn_hash(&self, state: &mut dyn Hasher);
+    fn dyn_eq(&self, other: &dyn OpaqueInternableThing) -> bool;
+    fn dyn_clone(&self) -> Box<dyn OpaqueInternableThing>;
+}
+
+impl Hash for dyn OpaqueInternableThing {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.dyn_hash(state);
+    }
+}
+
+impl PartialEq for dyn OpaqueInternableThing {
+    fn eq(&self, other: &Self) -> bool {
+        self.dyn_eq(other)
+    }
+}
+
+impl Eq for dyn OpaqueInternableThing {}
+
+impl Clone for Box<dyn OpaqueInternableThing> {
+    fn clone(&self) -> Self {
+        self.dyn_clone()
+    }
+}
+
 #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
 pub struct InTypeConstId(InternId);
-type InTypeConstLoc = (AstId<ast::Expr>, TypeOwnerId);
+type InTypeConstLoc = (AstId<ast::Expr>, TypeOwnerId, Box<dyn OpaqueInternableThing>);
 impl_intern!(InTypeConstId, InTypeConstLoc, intern_in_type_const, lookup_intern_in_type_const);
 
 impl InTypeConstId {
@@ -535,7 +569,7 @@ impl GeneralConstId {
                 parent.as_generic_def_id()
             }
             GeneralConstId::InTypeConstId(x) => {
-                let (_, parent) = x.lookup(db);
+                let (_, parent, _) = x.lookup(db);
                 parent.as_generic_def_id()
             }
         }

crates/hir-ty/src/infer.rs

@@ -41,8 +41,13 @@ use stdx::{always, never};
 use triomphe::Arc;
 
 use crate::{
-    db::HirDatabase, fold_tys, infer::coerce::CoerceMany, lower::ImplTraitLoweringMode,
-    static_lifetime, to_assoc_type_id, traits::FnTrait, utils::UnevaluatedConstEvaluatorFolder,
+    db::HirDatabase,
+    fold_tys,
+    infer::coerce::CoerceMany,
+    lower::ImplTraitLoweringMode,
+    static_lifetime, to_assoc_type_id,
+    traits::FnTrait,
+    utils::{InTypeConstIdMetadata, UnevaluatedConstEvaluatorFolder},
     AliasEq, AliasTy, ClosureId, DomainGoal, GenericArg, Goal, ImplTraitId, InEnvironment,
     Interner, ProjectionTy, RpitId, Substitution, TraitEnvironment, TraitRef, Ty, TyBuilder, TyExt,
 };
@@ -102,9 +107,9 @@ pub(crate) fn infer_query(db: &dyn HirDatabase, def: DefWithBodyId) -> Arc<Infer
                 },
             });
         }
-        DefWithBodyId::InTypeConstId(_) => {
-            // FIXME: We should know the expected type here.
-            ctx.return_ty = ctx.table.new_type_var();
+        DefWithBodyId::InTypeConstId(c) => {
+            ctx.return_ty =
+                c.lookup(db.upcast()).2.box_any().downcast::<InTypeConstIdMetadata>().unwrap().0;
         }
     }
 

crates/hir-ty/src/lower.rs

@@ -53,7 +53,10 @@ use crate::{
     mapping::{from_chalk_trait_id, ToChalk},
     static_lifetime, to_assoc_type_id, to_chalk_trait_id, to_placeholder_idx,
     utils::Generics,
-    utils::{all_super_trait_refs, associated_type_by_name_including_super_traits, generics},
+    utils::{
+        all_super_trait_refs, associated_type_by_name_including_super_traits, generics,
+        InTypeConstIdMetadata,
+    },
     AliasEq, AliasTy, Binders, BoundVar, CallableSig, Const, ConstScalar, DebruijnIndex, DynTy,
     FnPointer, FnSig, FnSubst, GenericArgData, ImplTraitId, Interner, ParamKind, PolyFnSig,
     ProjectionTy, QuantifiedWhereClause, QuantifiedWhereClauses, ReturnTypeImplTrait,
@@ -2052,7 +2055,13 @@ pub(crate) fn const_or_path_to_chalk(
                 // that are unlikely to be edited.
                 return unknown_const(expected_ty);
             }
-            let c = db.intern_in_type_const((x, owner)).into();
+            let c = db
+                .intern_in_type_const((
+                    x,
+                    owner,
+                    Box::new(InTypeConstIdMetadata(expected_ty.clone())),
+                ))
+                .into();
             intern_const_scalar(
                 ConstScalar::UnevaluatedConst(c, Substitution::empty(Interner)),
                 expected_ty,

crates/hir-ty/src/tests/simple.rs

@@ -1828,6 +1828,38 @@ impl Foo for u8 {
     );
 }
 
+#[test]
+fn const_eval_in_function_signature() {
+    check_types(
+        r#"
+const fn foo() -> usize {
+    5
+}
+
+fn f() -> [u8; foo()] {
+    loop {}
+}
+
+fn main() {
+    let t = f();
+      //^ [u8; 5]
+}"#,
+    );
+    check_types(
+        r#"
+//- minicore: default, builtin_impls
+fn f() -> [u8; Default::default()] {
+    loop {}
+}
+
+fn main() {
+    let t = f();
+      //^ [u8; 0]
+}
+    "#,
+    );
+}
+
 #[test]
 fn shadowing_primitive_with_inner_items() {
     check_types(

crates/hir-ty/src/utils.rs

@@ -1,7 +1,7 @@
 //! Helper functions for working with def, which don't need to be a separate
 //! query, but can't be computed directly from `*Data` (ie, which need a `db`).
 
-use std::iter;
+use std::{hash::Hash, iter};
 
 use base_db::CrateId;
 use chalk_ir::{
@@ -20,7 +20,8 @@ use hir_def::{
     resolver::{HasResolver, TypeNs},
     type_ref::{TraitBoundModifier, TypeRef},
     ConstParamId, EnumId, EnumVariantId, FunctionId, GenericDefId, ItemContainerId,
-    LocalEnumVariantId, Lookup, TraitId, TypeAliasId, TypeOrConstParamId, TypeParamId,
+    LocalEnumVariantId, Lookup, OpaqueInternableThing, TraitId, TypeAliasId, TypeOrConstParamId,
+    TypeParamId,
 };
 use hir_expand::name::Name;
 use intern::Interned;
@@ -464,3 +465,28 @@ pub(crate) fn detect_variant_from_bytes<'a>(
     };
     Some((var_id, var_layout))
 }
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub(crate) struct InTypeConstIdMetadata(pub(crate) Ty);
+
+impl OpaqueInternableThing for InTypeConstIdMetadata {
+    fn dyn_hash(&self, mut state: &mut dyn std::hash::Hasher) {
+        self.hash(&mut state);
+    }
+
+    fn dyn_eq(&self, other: &dyn OpaqueInternableThing) -> bool {
+        other.as_any().downcast_ref::<Self>().map_or(false, |x| self == x)
+    }
+
+    fn dyn_clone(&self) -> Box<dyn OpaqueInternableThing> {
+        Box::new(self.clone())
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn box_any(&self) -> Box<dyn std::any::Any> {
+        Box::new(self.clone())
+    }
+}