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