crates/hir_ty/src/utils.rs

   1 //! Helper functions for working with def, which don't need to be a separate
   2 //! query, but can't be computed directly from `*Data` (ie, which need a `db`).
   3
   4 use std::iter;
   5
   6 use base_db::CrateId;
   7 use chalk_ir::{fold::Shift, BoundVar, DebruijnIndex};
   8 use hir_def::{
   9     db::DefDatabase,
  10     generics::{
  11         GenericParams, TypeParamData, TypeParamProvenance, WherePredicate, WherePredicateTypeTarget,
  12     },
  13     intern::Interned,
  14     path::Path,
  15     resolver::{HasResolver, TypeNs},
  16     type_ref::{TraitBoundModifier, TypeRef},
  17     GenericDefId, ItemContainerId, Lookup, TraitId, TypeAliasId, TypeParamId,
  18 };
  19 use hir_expand::name::{name, Name};
  20 use rustc_hash::FxHashSet;
  21 use smallvec::{smallvec, SmallVec};
  22 use syntax::SmolStr;
  23
  24 use crate::{
  25     db::HirDatabase, ChalkTraitId, Interner, Substitution, TraitRef, TraitRefExt, TyKind,
  26     WhereClause,
  27 };
  28
  29 pub(crate) fn fn_traits(db: &dyn DefDatabase, krate: CrateId) -> impl Iterator<Item = TraitId> {
  30     [
  31         db.lang_item(krate, SmolStr::new_inline("fn")),
  32         db.lang_item(krate, SmolStr::new_inline("fn_mut")),
  33         db.lang_item(krate, SmolStr::new_inline("fn_once")),
  34     ]
  35     .into_iter()
  36     .flatten()
  37     .flat_map(|it| it.as_trait())
  38 }
  39
  40 fn direct_super_traits(db: &dyn DefDatabase, trait_: TraitId) -> SmallVec<[TraitId; 4]> {
  41     let resolver = trait_.resolver(db);
  42     // returning the iterator directly doesn't easily work because of
  43     // lifetime problems, but since there usually shouldn't be more than a
  44     // few direct traits this should be fine (we could even use some kind of
  45     // SmallVec if performance is a concern)
  46     let generic_params = db.generic_params(trait_.into());
  47     let trait_self = generic_params.find_trait_self_param();
  48     generic_params
  49         .where_predicates
  50         .iter()
  51         .filter_map(|pred| match pred {
  52             WherePredicate::ForLifetime { target, bound, .. }
  53             | WherePredicate::TypeBound { target, bound } => match target {
  54                 WherePredicateTypeTarget::TypeRef(type_ref) => match &**type_ref {
  55                     TypeRef::Path(p) if p == &Path::from(name![Self]) => bound.as_path(),
  56                     _ => None,
  57                 },
  58                 WherePredicateTypeTarget::TypeParam(local_id) if Some(*local_id) == trait_self => {
  59                     bound.as_path()
  60                 }
  61                 _ => None,
  62             },
  63             WherePredicate::Lifetime { .. } => None,
  64         })
  65         .filter_map(|(path, bound_modifier)| match bound_modifier {
  66             TraitBoundModifier::None => Some(path),
  67             TraitBoundModifier::Maybe => None,
  68         })
  69         .filter_map(|path| match resolver.resolve_path_in_type_ns_fully(db, path.mod_path()) {
  70             Some(TypeNs::TraitId(t)) => Some(t),
  71             _ => None,
  72         })
  73         .collect()
  74 }
  75
  76 fn direct_super_trait_refs(db: &dyn HirDatabase, trait_ref: &TraitRef) -> Vec<TraitRef> {
  77     // returning the iterator directly doesn't easily work because of
  78     // lifetime problems, but since there usually shouldn't be more than a
  79     // few direct traits this should be fine (we could even use some kind of
  80     // SmallVec if performance is a concern)
  81     let generic_params = db.generic_params(trait_ref.hir_trait_id().into());
  82     let trait_self = match generic_params.find_trait_self_param() {
  83         Some(p) => TypeParamId { parent: trait_ref.hir_trait_id().into(), local_id: p },
  84         None => return Vec::new(),
  85     };
  86     db.generic_predicates_for_param(trait_self, None)
  87         .iter()
  88         .filter_map(|pred| {
  89             pred.as_ref().filter_map(|pred| match pred.skip_binders() {
  90                 // FIXME: how to correctly handle higher-ranked bounds here?
  91                 WhereClause::Implemented(tr) => Some(
  92                     tr.clone()
  93                         .shifted_out_to(&Interner, DebruijnIndex::ONE)
  94                         .expect("FIXME unexpected higher-ranked trait bound"),
  95                 ),
  96                 _ => None,
  97             })
  98         })
  99         .map(|pred| pred.substitute(&Interner, &trait_ref.substitution))
 100         .collect()
 101 }
 102
 103 /// Returns an iterator over the whole super trait hierarchy (including the
 104 /// trait itself).
 105 pub fn all_super_traits(db: &dyn DefDatabase, trait_: TraitId) -> SmallVec<[TraitId; 4]> {
 106     // we need to take care a bit here to avoid infinite loops in case of cycles
 107     // (i.e. if we have `trait A: B; trait B: A;`)
 108
 109     let mut result = smallvec![trait_];
 110     let mut i = 0;
 111     while let Some(&t) = result.get(i) {
 112         // yeah this is quadratic, but trait hierarchies should be flat
 113         // enough that this doesn't matter
 114         for tt in direct_super_traits(db, t) {
 115             if !result.contains(&tt) {
 116                 result.push(tt);
 117             }
 118         }
 119         i += 1;
 120     }
 121     result
 122 }
 123
 124 /// Given a trait ref (`Self: Trait`), builds all the implied trait refs for
 125 /// super traits. The original trait ref will be included. So the difference to
 126 /// `all_super_traits` is that we keep track of type parameters; for example if
 127 /// we have `Self: Trait<u32, i32>` and `Trait<T, U>: OtherTrait<U>` we'll get
 128 /// `Self: OtherTrait<i32>`.
 129 pub(super) fn all_super_trait_refs(db: &dyn HirDatabase, trait_ref: TraitRef) -> SuperTraits {
 130     SuperTraits { db, seen: iter::once(trait_ref.trait_id).collect(), stack: vec![trait_ref] }
 131 }
 132
 133 pub(super) struct SuperTraits<'a> {
 134     db: &'a dyn HirDatabase,
 135     stack: Vec<TraitRef>,
 136     seen: FxHashSet<ChalkTraitId>,
 137 }
 138
 139 impl<'a> SuperTraits<'a> {
 140     fn elaborate(&mut self, trait_ref: &TraitRef) {
 141         let mut trait_refs = direct_super_trait_refs(self.db, trait_ref);
 142         trait_refs.retain(|tr| !self.seen.contains(&tr.trait_id));
 143         self.stack.extend(trait_refs);
 144     }
 145 }
 146
 147 impl<'a> Iterator for SuperTraits<'a> {
 148     type Item = TraitRef;
 149
 150     fn next(&mut self) -> Option<Self::Item> {
 151         if let Some(next) = self.stack.pop() {
 152             self.elaborate(&next);
 153             Some(next)
 154         } else {
 155             None
 156         }
 157     }
 158 }
 159
 160 pub(super) fn associated_type_by_name_including_super_traits(
 161     db: &dyn HirDatabase,
 162     trait_ref: TraitRef,
 163     name: &Name,
 164 ) -> Option<(TraitRef, TypeAliasId)> {
 165     all_super_trait_refs(db, trait_ref).find_map(|t| {
 166         let assoc_type = db.trait_data(t.hir_trait_id()).associated_type_by_name(name)?;
 167         Some((t, assoc_type))
 168     })
 169 }
 170
 171 pub(crate) fn generics(db: &dyn DefDatabase, def: GenericDefId) -> Generics {
 172     let parent_generics = parent_generic_def(db, def).map(|def| Box::new(generics(db, def)));
 173     Generics { def, params: db.generic_params(def), parent_generics }
 174 }
 175
 176 #[derive(Debug)]
 177 pub(crate) struct Generics {
 178     def: GenericDefId,
 179     pub(crate) params: Interned<GenericParams>,
 180     parent_generics: Option<Box<Generics>>,
 181 }
 182
 183 impl Generics {
 184     pub(crate) fn iter<'a>(
 185         &'a self,
 186     ) -> impl Iterator<Item = (TypeParamId, &'a TypeParamData)> + 'a {
 187         self.parent_generics
 188             .as_ref()
 189             .into_iter()
 190             .flat_map(|it| {
 191                 it.params
 192                     .types
 193                     .iter()
 194                     .map(move |(local_id, p)| (TypeParamId { parent: it.def, local_id }, p))
 195             })
 196             .chain(
 197                 self.params
 198                     .types
 199                     .iter()
 200                     .map(move |(local_id, p)| (TypeParamId { parent: self.def, local_id }, p)),
 201             )
 202     }
 203
 204     pub(crate) fn iter_parent<'a>(
 205         &'a self,
 206     ) -> impl Iterator<Item = (TypeParamId, &'a TypeParamData)> + 'a {
 207         self.parent_generics.as_ref().into_iter().flat_map(|it| {
 208             it.params
 209                 .types
 210                 .iter()
 211                 .map(move |(local_id, p)| (TypeParamId { parent: it.def, local_id }, p))
 212         })
 213     }
 214
 215     pub(crate) fn len(&self) -> usize {
 216         self.len_split().0
 217     }
 218
 219     /// (total, parents, child)
 220     pub(crate) fn len_split(&self) -> (usize, usize, usize) {
 221         let parent = self.parent_generics.as_ref().map_or(0, |p| p.len());
 222         let child = self.params.types.len();
 223         (parent + child, parent, child)
 224     }
 225
 226     /// (parent total, self param, type param list, impl trait)
 227     pub(crate) fn provenance_split(&self) -> (usize, usize, usize, usize) {
 228         let parent = self.parent_generics.as_ref().map_or(0, |p| p.len());
 229         let self_params = self
 230             .params
 231             .types
 232             .iter()
 233             .filter(|(_, p)| p.provenance == TypeParamProvenance::TraitSelf)
 234             .count();
 235         let list_params = self
 236             .params
 237             .types
 238             .iter()
 239             .filter(|(_, p)| p.provenance == TypeParamProvenance::TypeParamList)
 240             .count();
 241         let impl_trait_params = self
 242             .params
 243             .types
 244             .iter()
 245             .filter(|(_, p)| p.provenance == TypeParamProvenance::ArgumentImplTrait)
 246             .count();
 247         (parent, self_params, list_params, impl_trait_params)
 248     }
 249
 250     pub(crate) fn param_idx(&self, param: TypeParamId) -> Option<usize> {
 251         Some(self.find_param(param)?.0)
 252     }
 253
 254     fn find_param(&self, param: TypeParamId) -> Option<(usize, &TypeParamData)> {
 255         if param.parent == self.def {
 256             let (idx, (_local_id, data)) = self
 257                 .params
 258                 .types
 259                 .iter()
 260                 .enumerate()
 261                 .find(|(_, (idx, _))| *idx == param.local_id)
 262                 .unwrap();
 263             let (_total, parent_len, _child) = self.len_split();
 264             Some((parent_len + idx, data))
 265         } else {
 266             self.parent_generics.as_ref().and_then(|g| g.find_param(param))
 267         }
 268     }
 269
 270     /// Returns a Substitution that replaces each parameter by a bound variable.
 271     pub(crate) fn bound_vars_subst(&self, debruijn: DebruijnIndex) -> Substitution {
 272         Substitution::from_iter(
 273             &Interner,
 274             self.iter()
 275                 .enumerate()
 276                 .map(|(idx, _)| TyKind::BoundVar(BoundVar::new(debruijn, idx)).intern(&Interner)),
 277         )
 278     }
 279
 280     /// Returns a Substitution that replaces each parameter by itself (i.e. `Ty::Param`).
 281     pub(crate) fn type_params_subst(&self, db: &dyn HirDatabase) -> Substitution {
 282         Substitution::from_iter(
 283             &Interner,
 284             self.iter().map(|(id, _)| {
 285                 TyKind::Placeholder(crate::to_placeholder_idx(db, id)).intern(&Interner)
 286             }),
 287         )
 288     }
 289 }
 290
 291 fn parent_generic_def(db: &dyn DefDatabase, def: GenericDefId) -> Option<GenericDefId> {
 292     let container = match def {
 293         GenericDefId::FunctionId(it) => it.lookup(db).container,
 294         GenericDefId::TypeAliasId(it) => it.lookup(db).container,
 295         GenericDefId::ConstId(it) => it.lookup(db).container,
 296         GenericDefId::EnumVariantId(it) => return Some(it.parent.into()),
 297         GenericDefId::AdtId(_) | GenericDefId::TraitId(_) | GenericDefId::ImplId(_) => return None,
 298     };
 299
 300     match container {
 301         ItemContainerId::ImplId(it) => Some(it.into()),
 302         ItemContainerId::TraitId(it) => Some(it.into()),
 303         ItemContainerId::ModuleId(_) | ItemContainerId::ExternBlockId(_) => None,
 304     }
 305 }