1use 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 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 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 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 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 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 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}