hir_ty/
generics.rs

1//! Utilities for working with generics.
2//!
3//! The layout for generics as expected by chalk are as follows:
4//! - Parent parameters
5//! - Optional Self parameter
6//! - Lifetime parameters
7//! - Type or Const parameters
8//!
9//! where parent follows the same scheme.
10use std::ops;
11
12use hir_def::{
13    ConstParamId, GenericDefId, GenericParamId, ItemContainerId, LifetimeParamId, Lookup,
14    TypeOrConstParamId, TypeParamId,
15    db::DefDatabase,
16    expr_store::ExpressionStore,
17    hir::generics::{
18        GenericParamDataRef, GenericParams, LifetimeParamData, LocalLifetimeParamId,
19        LocalTypeOrConstParamId, TypeOrConstParamData, TypeParamProvenance, WherePredicate,
20    },
21};
22use itertools::chain;
23use triomphe::Arc;
24
25pub fn generics(db: &dyn DefDatabase, def: GenericDefId) -> Generics {
26    let parent_generics = parent_generic_def(db, def).map(|def| Box::new(generics(db, def)));
27    let (params, store) = db.generic_params_and_store(def);
28    let has_trait_self_param = params.trait_self_param().is_some();
29    Generics { def, params, parent_generics, has_trait_self_param, store }
30}
31#[derive(Clone, Debug)]
32pub struct Generics {
33    def: GenericDefId,
34    params: Arc<GenericParams>,
35    store: Arc<ExpressionStore>,
36    parent_generics: Option<Box<Generics>>,
37    has_trait_self_param: bool,
38}
39
40impl<T> ops::Index<T> for Generics
41where
42    GenericParams: ops::Index<T>,
43{
44    type Output = <GenericParams as ops::Index<T>>::Output;
45    fn index(&self, index: T) -> &Self::Output {
46        &self.params[index]
47    }
48}
49
50impl Generics {
51    pub(crate) fn def(&self) -> GenericDefId {
52        self.def
53    }
54
55    pub(crate) fn store(&self) -> &ExpressionStore {
56        &self.store
57    }
58
59    pub(crate) fn where_predicates(&self) -> impl Iterator<Item = &WherePredicate> {
60        self.params.where_predicates().iter()
61    }
62
63    pub(crate) fn is_empty(&self) -> bool {
64        self.params.is_empty() && self.parent_generics.as_ref().is_none_or(|g| g.params.is_empty())
65    }
66
67    pub(crate) fn iter_id(&self) -> impl Iterator<Item = GenericParamId> + '_ {
68        self.iter_parent_id().chain(self.iter_self_id())
69    }
70
71    pub(crate) fn iter_self_id(&self) -> impl Iterator<Item = GenericParamId> + '_ {
72        self.iter_self().map(|(id, _)| id)
73    }
74
75    pub(crate) fn iter_parent_id(&self) -> impl Iterator<Item = GenericParamId> + '_ {
76        self.iter_parent().map(|(id, _)| id)
77    }
78
79    pub(crate) fn iter_self_type_or_consts(
80        &self,
81    ) -> impl DoubleEndedIterator<Item = (LocalTypeOrConstParamId, &TypeOrConstParamData)> + '_
82    {
83        let mut toc = self.params.iter_type_or_consts();
84        let trait_self_param = self.has_trait_self_param.then(|| toc.next()).flatten();
85        chain!(trait_self_param, toc)
86    }
87
88    /// Iterate over the parent params followed by self params.
89    pub(crate) fn iter(
90        &self,
91    ) -> impl DoubleEndedIterator<Item = (GenericParamId, GenericParamDataRef<'_>)> + '_ {
92        self.iter_parent().chain(self.iter_self())
93    }
94
95    pub(crate) fn iter_parents_with_store(
96        &self,
97    ) -> impl Iterator<Item = ((GenericParamId, GenericParamDataRef<'_>), &ExpressionStore)> + '_
98    {
99        self.iter_parent()
100            .zip(self.parent_generics().into_iter().flat_map(|it| std::iter::repeat(&*it.store)))
101    }
102
103    /// Iterate over the params without parent params.
104    pub(crate) fn iter_self(
105        &self,
106    ) -> impl DoubleEndedIterator<Item = (GenericParamId, GenericParamDataRef<'_>)> + '_ {
107        let mut toc = self.params.iter_type_or_consts().map(from_toc_id(self));
108        let trait_self_param = self.has_trait_self_param.then(|| toc.next()).flatten();
109        chain!(trait_self_param, self.params.iter_lt().map(from_lt_id(self)), toc)
110    }
111
112    /// Iterator over types and const params of parent.
113    pub(crate) fn iter_parent(
114        &self,
115    ) -> impl DoubleEndedIterator<Item = (GenericParamId, GenericParamDataRef<'_>)> + '_ {
116        self.parent_generics().into_iter().flat_map(|it| {
117            let mut toc = it.params.iter_type_or_consts().map(from_toc_id(it));
118            let trait_self_param = it.has_trait_self_param.then(|| toc.next()).flatten();
119            chain!(trait_self_param, it.params.iter_lt().map(from_lt_id(it)), toc)
120        })
121    }
122
123    /// Returns total number of generic parameters in scope, including those from parent.
124    pub(crate) fn len(&self) -> usize {
125        let parent = self.len_parent();
126        let child = self.params.len();
127        parent + child
128    }
129
130    #[inline]
131    pub(crate) fn len_parent(&self) -> usize {
132        self.parent_generics().map_or(0, Generics::len)
133    }
134
135    /// Returns numbers of generic parameters excluding those from parent.
136    pub(crate) fn len_self(&self) -> usize {
137        self.params.len()
138    }
139
140    pub(crate) fn len_lifetimes_self(&self) -> usize {
141        self.params.len_lifetimes()
142    }
143
144    /// (parent total, self param, type params, const params, impl trait list, lifetimes)
145    pub(crate) fn provenance_split(&self) -> (usize, bool, usize, usize, usize, usize) {
146        let mut self_param = false;
147        let mut type_params = 0;
148        let mut impl_trait_params = 0;
149        let mut const_params = 0;
150        self.params.iter_type_or_consts().for_each(|(_, data)| match data {
151            TypeOrConstParamData::TypeParamData(p) => match p.provenance {
152                TypeParamProvenance::TypeParamList => type_params += 1,
153                TypeParamProvenance::TraitSelf => self_param |= true,
154                TypeParamProvenance::ArgumentImplTrait => impl_trait_params += 1,
155            },
156            TypeOrConstParamData::ConstParamData(_) => const_params += 1,
157        });
158
159        let lifetime_params = self.params.len_lifetimes();
160
161        let parent_len = self.parent_generics().map_or(0, Generics::len);
162        (parent_len, self_param, type_params, const_params, impl_trait_params, lifetime_params)
163    }
164
165    pub(crate) fn type_or_const_param(
166        &self,
167        param: TypeOrConstParamId,
168    ) -> Option<(usize, TypeOrConstParamData)> {
169        let idx = self.find_type_or_const_param(param)?;
170        self.iter().nth(idx).and_then(|p| {
171            let data = match p.1 {
172                GenericParamDataRef::TypeParamData(p) => p.clone().into(),
173                GenericParamDataRef::ConstParamData(p) => p.clone().into(),
174                _ => return None,
175            };
176            Some((idx, data))
177        })
178    }
179
180    pub fn type_or_const_param_idx(&self, param: TypeOrConstParamId) -> Option<usize> {
181        self.find_type_or_const_param(param)
182    }
183
184    fn find_type_or_const_param(&self, param: TypeOrConstParamId) -> Option<usize> {
185        if param.parent == self.def {
186            let idx = param.local_id.into_raw().into_u32() as usize;
187            debug_assert!(
188                idx <= self.params.len_type_or_consts(),
189                "idx: {} len: {}",
190                idx,
191                self.params.len_type_or_consts()
192            );
193            if self.params.trait_self_param() == Some(param.local_id) {
194                return Some(idx);
195            }
196            Some(self.parent_generics().map_or(0, |g| g.len()) + self.params.len_lifetimes() + idx)
197        } else {
198            debug_assert_eq!(self.parent_generics().map(|it| it.def), Some(param.parent));
199            self.parent_generics().and_then(|g| g.find_type_or_const_param(param))
200        }
201    }
202
203    pub fn lifetime_idx(&self, lifetime: LifetimeParamId) -> Option<usize> {
204        self.find_lifetime(lifetime)
205    }
206
207    fn find_lifetime(&self, lifetime: LifetimeParamId) -> Option<usize> {
208        if lifetime.parent == self.def {
209            let idx = lifetime.local_id.into_raw().into_u32() as usize;
210            debug_assert!(idx <= self.params.len_lifetimes());
211            Some(
212                self.parent_generics().map_or(0, |g| g.len())
213                    + self.params.trait_self_param().is_some() as usize
214                    + idx,
215            )
216        } else {
217            debug_assert_eq!(self.parent_generics().map(|it| it.def), Some(lifetime.parent));
218            self.parent_generics().and_then(|g| g.find_lifetime(lifetime))
219        }
220    }
221
222    pub(crate) fn parent_generics(&self) -> Option<&Generics> {
223        self.parent_generics.as_deref()
224    }
225}
226
227pub(crate) fn trait_self_param_idx(db: &dyn DefDatabase, def: GenericDefId) -> Option<usize> {
228    match def {
229        GenericDefId::TraitId(_) => {
230            let params = db.generic_params(def);
231            params.trait_self_param().map(|idx| idx.into_raw().into_u32() as usize)
232        }
233        GenericDefId::ImplId(_) => None,
234        _ => {
235            let parent_def = parent_generic_def(db, def)?;
236            let parent_params = db.generic_params(parent_def);
237            let parent_self_idx = parent_params.trait_self_param()?.into_raw().into_u32() as usize;
238            Some(parent_self_idx)
239        }
240    }
241}
242
243pub(crate) fn parent_generic_def(db: &dyn DefDatabase, def: GenericDefId) -> Option<GenericDefId> {
244    let container = match def {
245        GenericDefId::FunctionId(it) => it.lookup(db).container,
246        GenericDefId::TypeAliasId(it) => it.lookup(db).container,
247        GenericDefId::ConstId(it) => it.lookup(db).container,
248        GenericDefId::StaticId(_)
249        | GenericDefId::AdtId(_)
250        | GenericDefId::TraitId(_)
251        | GenericDefId::ImplId(_) => return None,
252    };
253
254    match container {
255        ItemContainerId::ImplId(it) => Some(it.into()),
256        ItemContainerId::TraitId(it) => Some(it.into()),
257        ItemContainerId::ModuleId(_) | ItemContainerId::ExternBlockId(_) => None,
258    }
259}
260
261fn from_toc_id<'a>(
262    it: &'a Generics,
263) -> impl Fn(
264    (LocalTypeOrConstParamId, &'a TypeOrConstParamData),
265) -> (GenericParamId, GenericParamDataRef<'a>) {
266    move |(local_id, p): (_, _)| {
267        let id = TypeOrConstParamId { parent: it.def, local_id };
268        match p {
269            TypeOrConstParamData::TypeParamData(p) => (
270                GenericParamId::TypeParamId(TypeParamId::from_unchecked(id)),
271                GenericParamDataRef::TypeParamData(p),
272            ),
273            TypeOrConstParamData::ConstParamData(p) => (
274                GenericParamId::ConstParamId(ConstParamId::from_unchecked(id)),
275                GenericParamDataRef::ConstParamData(p),
276            ),
277        }
278    }
279}
280
281fn from_lt_id<'a>(
282    it: &'a Generics,
283) -> impl Fn((LocalLifetimeParamId, &'a LifetimeParamData)) -> (GenericParamId, GenericParamDataRef<'a>)
284{
285    move |(local_id, p): (_, _)| {
286        (
287            GenericParamId::LifetimeParamId(LifetimeParamId { parent: it.def, local_id }),
288            GenericParamDataRef::LifetimeParamData(p),
289        )
290    }
291}