Skip to main content

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.
10
11use arrayvec::ArrayVec;
12use base_db::SourceDatabase;
13use hir_def::{
14    ConstParamId, GenericDefId, GenericParamId, ItemContainerId, LifetimeParamId, Lookup,
15    TypeOrConstParamId, TypeParamId,
16    expr_store::ExpressionStore,
17    hir::generics::{
18        GenericParamDataRef, GenericParams, LifetimeParamData, TypeOrConstParamData,
19        TypeParamProvenance, WherePredicate,
20    },
21};
22
23pub(crate) fn generics(db: &dyn SourceDatabase, def: GenericDefId) -> Generics<'_> {
24    let mut chain = ArrayVec::new();
25    let mut parent_params_len = 0;
26    if let Some(parent_def) = parent_generic_def(db, def) {
27        let (parent_params, parent_store) = GenericParams::with_store(db, parent_def);
28        chain.push(SingleGenerics {
29            def: parent_def,
30            params: parent_params,
31            store: parent_store,
32            preceding_params_len: 0,
33        });
34        parent_params_len = parent_params.len() as u32;
35    }
36    let (params, store) = GenericParams::with_store(db, def);
37    chain.push(SingleGenerics { def, params, store, preceding_params_len: parent_params_len });
38    Generics { chain }
39}
40
41#[derive(Debug)]
42pub struct Generics<'db> {
43    chain: ArrayVec<SingleGenerics<'db>, 2>,
44}
45
46#[derive(Debug)]
47pub(crate) struct SingleGenerics<'db> {
48    def: GenericDefId,
49    preceding_params_len: u32,
50    params: &'db GenericParams,
51    store: &'db ExpressionStore,
52}
53
54impl<'db> SingleGenerics<'db> {
55    pub(crate) fn def(&self) -> GenericDefId {
56        self.def
57    }
58
59    pub(crate) fn store(&self) -> &'db ExpressionStore {
60        self.store
61    }
62
63    pub(crate) fn where_predicates(&self) -> impl Iterator<Item = &WherePredicate> {
64        self.params.where_predicates().iter()
65    }
66
67    pub(crate) fn has_no_params(&self) -> bool {
68        self.params.is_empty()
69    }
70
71    pub(crate) fn len_lifetimes(&self) -> usize {
72        self.params.len_lifetimes()
73    }
74
75    pub(crate) fn len(&self, consider_late_bound: bool) -> usize {
76        if consider_late_bound {
77            self.params.len()
78        } else {
79            self.params.len() - self.params.len_late_bound_lifetimes()
80        }
81    }
82
83    fn iter_lifetimes(&self) -> impl Iterator<Item = (LifetimeParamId, &'db LifetimeParamData)> {
84        let parent = self.def;
85        self.params
86            .iter_early_bound_lt()
87            .map(move |(local_id, data)| (LifetimeParamId { parent, local_id }, data))
88    }
89
90    fn iter_late_bound_lifetimes(
91        &self,
92        consider_late_bound: bool,
93    ) -> impl Iterator<Item = (LifetimeParamId, &'db LifetimeParamData)> {
94        let parent = self.def;
95        self.params.iter_late_bound_lt().filter_map(move |(local_id, data)| {
96            consider_late_bound.then_some((LifetimeParamId { parent, local_id }, data))
97        })
98    }
99
100    pub(crate) fn iter_type_or_consts(
101        &self,
102    ) -> impl Iterator<Item = (TypeOrConstParamId, &'db TypeOrConstParamData)> {
103        let parent = self.def;
104        self.params
105            .iter_type_or_consts()
106            .map(move |(local_id, data)| (TypeOrConstParamId { parent, local_id }, data))
107    }
108
109    fn iter_type_or_consts_as_generic(
110        &self,
111    ) -> impl Iterator<Item = (GenericParamId, GenericParamDataRef<'db>)> {
112        self.iter_type_or_consts().map(|(id, data)| match data {
113            TypeOrConstParamData::TypeParamData(data) => (
114                GenericParamId::TypeParamId(TypeParamId::from_unchecked(id)),
115                GenericParamDataRef::TypeParamData(data),
116            ),
117            TypeOrConstParamData::ConstParamData(data) => (
118                GenericParamId::ConstParamId(ConstParamId::from_unchecked(id)),
119                GenericParamDataRef::ConstParamData(data),
120            ),
121        })
122    }
123
124    fn trait_self_and_others(
125        &self,
126    ) -> (
127        Option<(GenericParamId, GenericParamDataRef<'db>)>,
128        impl Iterator<Item = (GenericParamId, GenericParamDataRef<'db>)>,
129    ) {
130        let mut iter = self.iter_type_or_consts_as_generic();
131        let trait_self = if let GenericDefId::TraitId(_) = self.def { iter.next() } else { None };
132        (trait_self, iter)
133    }
134
135    pub(crate) fn iter(
136        &self,
137        consider_late_bound: bool,
138    ) -> impl Iterator<Item = (GenericParamId, GenericParamDataRef<'db>)> {
139        let lifetime_map = |(id, data)| {
140            (GenericParamId::LifetimeParamId(id), GenericParamDataRef::LifetimeParamData(data))
141        };
142        let lifetimes = self.iter_lifetimes().map(lifetime_map);
143        let late_bound_lifetimes =
144            self.iter_late_bound_lifetimes(consider_late_bound).map(lifetime_map);
145
146        let (trait_self, type_and_consts) = self.trait_self_and_others();
147        trait_self.into_iter().chain(lifetimes).chain(type_and_consts).chain(late_bound_lifetimes)
148    }
149
150    pub(crate) fn iter_with_idx(
151        &self,
152    ) -> impl Iterator<Item = (u32, GenericParamId, GenericParamDataRef<'db>)> {
153        std::iter::zip(self.preceding_params_len.., self.iter(false))
154            .map(|(index, (id, data))| (index, id, data))
155    }
156
157    pub(crate) fn iter_id(
158        &self,
159        consider_late_bound: bool,
160    ) -> impl Iterator<Item = GenericParamId> {
161        self.iter(consider_late_bound).map(|(id, _)| id)
162    }
163
164    pub(crate) fn iter_late_bound(
165        &self,
166    ) -> impl Iterator<Item = (GenericParamId, GenericParamDataRef<'db>)> {
167        // we don't handle late bound types or const now, so it is ignored for now
168        let parent = self.def;
169        self.params.iter_late_bound_lt().map(move |(local_id, data)| {
170            (
171                GenericParamId::LifetimeParamId(LifetimeParamId { parent, local_id }),
172                GenericParamDataRef::LifetimeParamData(data),
173            )
174        })
175    }
176}
177
178impl<'db> Generics<'db> {
179    pub(crate) fn iter_owners(&self) -> impl DoubleEndedIterator<Item = &SingleGenerics<'db>> {
180        self.chain.iter()
181    }
182
183    pub(crate) fn owner(&self) -> &SingleGenerics<'db> {
184        self.chain.last().expect("must have an owner params")
185    }
186
187    pub(crate) fn parent(&self) -> Option<&SingleGenerics<'db>> {
188        match &*self.chain {
189            [parent, _owner] => Some(parent),
190            _ => None,
191        }
192    }
193
194    pub(crate) fn has_no_params(&self) -> bool {
195        self.iter_owners().all(|owner| owner.has_no_params())
196    }
197
198    pub(crate) fn def(&self) -> GenericDefId {
199        self.owner().def
200    }
201
202    pub(crate) fn store(&self) -> &'db ExpressionStore {
203        self.owner().store
204    }
205
206    pub(crate) fn iter_self(
207        &self,
208    ) -> impl Iterator<Item = (GenericParamId, GenericParamDataRef<'db>)> {
209        self.owner().iter(false)
210    }
211
212    pub(crate) fn iter_self_with_idx(
213        &self,
214    ) -> impl Iterator<Item = (u32, GenericParamId, GenericParamDataRef<'db>)> {
215        self.owner().iter_with_idx()
216    }
217
218    pub(crate) fn iter_self_late_bound(
219        &self,
220    ) -> impl Iterator<Item = (GenericParamId, GenericParamDataRef<'db>)> {
221        self.owner().iter_late_bound()
222    }
223
224    pub(crate) fn iter_parent_id(&self) -> impl Iterator<Item = GenericParamId> {
225        self.parent().into_iter().flat_map(move |parent| parent.iter_id(false))
226    }
227
228    pub(crate) fn iter_self_type_or_consts(
229        &self,
230    ) -> impl Iterator<Item = (TypeOrConstParamId, &'db TypeOrConstParamData)> {
231        self.owner().iter_type_or_consts()
232    }
233
234    /// Iterate over the parent params followed by self params.
235    pub(crate) fn iter(
236        &self,
237        consider_late_bound: bool,
238    ) -> impl Iterator<Item = (GenericParamId, GenericParamDataRef<'db>)> {
239        self.iter_owners().flat_map(move |owner| owner.iter(consider_late_bound))
240    }
241
242    pub(crate) fn iter_id(
243        &self,
244        consider_late_bound: bool,
245    ) -> impl Iterator<Item = GenericParamId> {
246        self.iter_owners().flat_map(move |owner| owner.iter_id(consider_late_bound))
247    }
248
249    /// Returns total number of generic parameters in scope, including those from parent.
250    pub(crate) fn len(&self, consider_late_bound: bool) -> usize {
251        match &*self.chain {
252            [parent, owner] => parent.len(consider_late_bound) + owner.len(consider_late_bound),
253            [owner] => owner.len(consider_late_bound),
254            _ => unreachable!(),
255        }
256    }
257
258    #[inline]
259    pub(crate) fn len_parent(&self) -> usize {
260        // add `consider_late_bound` arg if needed in future, currently it's not needed.
261        self.parent().map_or(0, |p| p.len(true))
262    }
263
264    pub(crate) fn len_lifetimes_self(&self) -> usize {
265        self.owner().len_lifetimes()
266    }
267
268    pub(crate) fn provenance_split(&self) -> ProvenanceSplit {
269        let parent_total = self.len_parent();
270
271        let owner = self.owner();
272        let lifetimes = owner.params.len_lifetimes();
273
274        let mut has_self_param = false;
275        let mut non_impl_trait_type_params = 0;
276        let mut impl_trait_type_params = 0;
277        let mut const_params = 0;
278        owner.params.iter_type_or_consts().for_each(|(_, data)| match data {
279            TypeOrConstParamData::TypeParamData(p) => match p.provenance {
280                TypeParamProvenance::TypeParamList => non_impl_trait_type_params += 1,
281                TypeParamProvenance::TraitSelf => has_self_param |= true,
282                TypeParamProvenance::ArgumentImplTrait => impl_trait_type_params += 1,
283            },
284            TypeOrConstParamData::ConstParamData(_) => const_params += 1,
285        });
286
287        ProvenanceSplit {
288            parent_total,
289            has_self_param,
290            non_impl_trait_type_params,
291            const_params,
292            impl_trait_type_params,
293            lifetimes,
294        }
295    }
296
297    fn find_owner(&self, def: GenericDefId) -> &SingleGenerics<'db> {
298        match &*self.chain {
299            [parent, owner] => {
300                if parent.def == def {
301                    parent
302                } else {
303                    debug_assert_eq!(def, owner.def);
304                    owner
305                }
306            }
307            [owner] => {
308                debug_assert_eq!(def, owner.def);
309                owner
310            }
311            _ => unreachable!(),
312        }
313    }
314
315    pub(crate) fn type_or_const_param_idx(&self, param: TypeOrConstParamId) -> u32 {
316        let owner = self.find_owner(param.parent);
317        let has_trait_self = matches!(owner.def, GenericDefId::TraitId(_));
318        if has_trait_self && param.local_id == GenericParams::SELF_PARAM_ID_IN_SELF {
319            owner.preceding_params_len
320        } else {
321            owner.preceding_params_len
322                + owner.len_lifetimes() as u32
323                + param.local_id.into_raw().into_u32()
324        }
325    }
326
327    // Rename this?
328    pub(crate) fn lifetime_param_idx(
329        &self,
330        param: LifetimeParamId,
331        is_lowering_impl_trait_bounds: bool,
332    ) -> (u32, bool) {
333        let owner = self.find_owner(param.parent);
334        if is_lowering_impl_trait_bounds {
335            let idx = self.opaque_lifetime_idx(param);
336            return (owner.preceding_params_len + (idx as u32), false);
337        }
338
339        let has_trait_self = matches!(owner.def, GenericDefId::TraitId(_));
340        match owner.params.lifetime_param_idx(&param.local_id) {
341            Some((idx, is_late_bound)) => {
342                let idx = if is_late_bound {
343                    idx as u32
344                } else {
345                    owner.preceding_params_len + u32::from(has_trait_self) + (idx as u32)
346                };
347                (idx, is_late_bound)
348            }
349            _ => unreachable!(),
350        }
351    }
352
353    #[deprecated = "don't use this; it's easy to expose an erroneous `Generics` with this"]
354    pub(crate) fn empty(def: GenericDefId) -> Self {
355        let mut chain = ArrayVec::new();
356        chain.push(SingleGenerics {
357            def,
358            preceding_params_len: 0,
359            params: GenericParams::empty(),
360            store: ExpressionStore::empty(),
361        });
362        Generics { chain }
363    }
364
365    fn opaque_lifetime_idx(&self, param: LifetimeParamId) -> usize {
366        self.find_owner(param.parent)
367            .iter_id(true)
368            .position(|id| {
369                let GenericParamId::LifetimeParamId(id) = id else {
370                    return false;
371                };
372                param == id
373            })
374            .unwrap()
375    }
376}
377
378pub(crate) struct ProvenanceSplit {
379    pub(crate) parent_total: usize,
380    // The rest are about self.
381    pub(crate) has_self_param: bool,
382    pub(crate) non_impl_trait_type_params: usize,
383    pub(crate) const_params: usize,
384    pub(crate) impl_trait_type_params: usize,
385    pub(crate) lifetimes: usize,
386}
387
388fn parent_generic_def(db: &dyn SourceDatabase, def: GenericDefId) -> Option<GenericDefId> {
389    let container = match def {
390        GenericDefId::FunctionId(it) => it.lookup(db).container,
391        GenericDefId::TypeAliasId(it) => it.lookup(db).container,
392        GenericDefId::ConstId(it) => it.lookup(db).container,
393        GenericDefId::StaticId(_)
394        | GenericDefId::AdtId(_)
395        | GenericDefId::TraitId(_)
396        | GenericDefId::ImplId(_) => return None,
397    };
398
399    match container {
400        ItemContainerId::ImplId(it) => Some(it.into()),
401        ItemContainerId::TraitId(it) => Some(it.into()),
402        ItemContainerId::ModuleId(_) | ItemContainerId::ExternBlockId(_) => None,
403    }
404}