Skip to content

Commit 9d34bf5

Browse files
lcnrlcnr
committed
impl type_of for const args in ty dependent paths
1 parent d8ed1b0 commit 9d34bf5

File tree

3 files changed

+134
-88
lines changed

3 files changed

+134
-88
lines changed

‎src/librustc_middle/query/mod.rs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,21 @@ rustc_queries! {
8989
desc { |tcx| "HIR owner items in `{}`", tcx.def_path_str(key.to_def_id()) }
9090
}
9191

92+
/// Computes the `DefId` of the corresponding const parameter of a const argument.
93+
/// Returns `None` if `def_id` is not a const argument.
94+
///
95+
/// ```rust
96+
/// let a = foo::<7>();
97+
/// ^ Calling `const_param_of` for this argument,
98+
///
99+
/// fn foo<const N: usize>()
100+
/// ^ returns this `DefId`.
101+
/// ```
102+
query const_param_of(key: DefId) -> Option<DefId> {
103+
cache_on_disk_if { key.is_local() }
104+
desc { |tcx| "computing the const parameter of `{}`", tcx.def_path_str(key) }
105+
}
106+
92107
/// Records the type of every item.
93108
query type_of(key: DefId) -> Ty<'tcx> {
94109
desc { |tcx| "computing type of `{}`", tcx.def_path_str(key) }

‎src/librustc_typeck/collect.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ fn collect_mod_item_types(tcx: TyCtxt<'_>, module_def_id: DefId) {
6464

6565
pub fn provide(providers: &mut Providers<'_>) {
6666
*providers = Providers {
67+
const_param_of: type_of::const_param_of,
6768
type_of: type_of::type_of,
6869
generics_of,
6970
predicates_of,

‎src/librustc_typeck/collect/type_of.rs

Lines changed: 118 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,118 @@ use rustc_trait_selection::traits;
1717
use super::ItemCtxt;
1818
use super::{bad_placeholder_type, is_suggestable_infer_ty};
1919

20+
pub(super) fn const_param_of(tcx: TyCtxt<'_>, def_id: DefId) -> Option<DefId> {
21+
use hir::*;
22+
23+
// We can just exit here, as `const_param_of` is called for
24+
// all generic arguments, meaning that we would return `None` anyways
25+
// if `const_param_of` is not cached
26+
let hir_id = tcx.hir().as_local_hir_id(def_id.as_local()?);
27+
28+
if let Node::AnonConst(_) = tcx.hir().get(hir_id) {
29+
let parent_node_id = tcx.hir().get_parent_node(hir_id);
30+
let parent_node = tcx.hir().get(parent_node_id);
31+
32+
match parent_node {
33+
Node::Expr(&Expr {
34+
kind:
35+
ExprKind::MethodCall(segment, ..) | ExprKind::Path(QPath::TypeRelative(_, segment)),
36+
..
37+
}) => {
38+
let body_owner = tcx.hir().get_parent_did(parent_node_id);
39+
let tables = tcx.typeck_tables_of(body_owner.to_def_id());
40+
// This may fail in case the method/path does not actually exist.
41+
// As there is no relevant param for `def_id`, we simply return
42+
// `None` here.
43+
let type_dependent_def = tables.type_dependent_def_id(parent_node_id)?;
44+
let idx = segment
45+
.args
46+
.and_then(|args| {
47+
args.args
48+
.iter()
49+
.filter(|arg| arg.is_const())
50+
.position(|arg| arg.id() == hir_id)
51+
})
52+
.unwrap_or_else(|| {
53+
bug!("no arg matching AnonConst in segment");
54+
});
55+
56+
tcx.generics_of(type_dependent_def)
57+
.params
58+
.iter()
59+
.filter(|param| matches!(param.kind, ty::GenericParamDefKind::Const))
60+
.nth(idx)
61+
.map(|param| param.def_id)
62+
}
63+
64+
Node::Ty(&Ty { kind: TyKind::Path(_), .. })
65+
| Node::Expr(&Expr { kind: ExprKind::Struct(..), .. })
66+
| Node::Expr(&Expr { kind: ExprKind::Path(_), .. })
67+
| Node::TraitRef(..) => {
68+
let path = match parent_node {
69+
Node::Ty(&Ty { kind: TyKind::Path(QPath::Resolved(_, path)), .. })
70+
| Node::TraitRef(&TraitRef { path, .. }) => &*path,
71+
Node::Expr(&Expr {
72+
kind:
73+
ExprKind::Path(QPath::Resolved(_, path))
74+
| ExprKind::Struct(&QPath::Resolved(_, path), ..),
75+
..
76+
}) => {
77+
let body_owner = tcx.hir().get_parent_did(parent_node_id);
78+
let _tables = tcx.typeck_tables_of(body_owner.to_def_id());
79+
&*path
80+
}
81+
_ => span_bug!(DUMMY_SP, "unexpected const parent path {:?}", parent_node),
82+
};
83+
84+
// We've encountered an `AnonConst` in some path, so we need to
85+
// figure out which generic parameter it corresponds to and return
86+
// the relevant type.
87+
88+
let (arg_index, segment) = path
89+
.segments
90+
.iter()
91+
.filter_map(|seg| seg.args.map(|args| (args.args, seg)))
92+
.find_map(|(args, seg)| {
93+
args.iter()
94+
.filter(|arg| arg.is_const())
95+
.position(|arg| arg.id() == hir_id)
96+
.map(|index| (index, seg))
97+
})
98+
.unwrap_or_else(|| {
99+
bug!("no arg matching AnonConst in path");
100+
});
101+
102+
// Try to use the segment resolution if it is valid, otherwise we
103+
// default to the path resolution.
104+
let res = segment.res.filter(|&r| r != Res::Err).unwrap_or(path.res);
105+
let generics = match res {
106+
Res::Def(DefKind::Ctor(..), def_id) => {
107+
tcx.generics_of(tcx.parent(def_id).unwrap())
108+
}
109+
Res::Def(_, def_id) => tcx.generics_of(def_id),
110+
res => span_bug!(
111+
DUMMY_SP,
112+
"unexpected anon const res {:?} in path: {:?}",
113+
res,
114+
path,
115+
),
116+
};
117+
118+
generics
119+
.params
120+
.iter()
121+
.filter(|param| matches!(param.kind, ty::GenericParamDefKind::Const))
122+
.nth(arg_index)
123+
.map(|param| param.def_id)
124+
}
125+
_ => return None,
126+
}
127+
} else {
128+
None
129+
}
130+
}
131+
20132
pub(super) fn type_of(tcx: TyCtxt<'_>, def_id: DefId) -> Ty<'_> {
21133
use rustc_hir::*;
22134

@@ -187,6 +299,12 @@ pub(super) fn type_of(tcx: TyCtxt<'_>, def_id: DefId) -> Ty<'_> {
187299
}
188300

189301
Node::AnonConst(_) => {
302+
if let Some(param) = tcx.const_param_of(def_id) {
303+
// We defer to `type_of` of the corresponding parameter
304+
// for generic arguments.
305+
return tcx.type_of(param);
306+
}
307+
190308
let parent_node = tcx.hir().get(tcx.hir().get_parent_node(hir_id));
191309
match parent_node {
192310
Node::Ty(&Ty { kind: TyKind::Array(_, ref constant), .. })
@@ -203,94 +321,6 @@ pub(super) fn type_of(tcx: TyCtxt<'_>, def_id: DefId) -> Ty<'_> {
203321
.discr_type()
204322
.to_ty(tcx),
205323

206-
Node::Ty(&Ty { kind: TyKind::Path(_), .. })
207-
| Node::Expr(&Expr { kind: ExprKind::Struct(..) | ExprKind::Path(_), .. })
208-
| Node::TraitRef(..) => {
209-
let path = match parent_node {
210-
Node::Ty(&Ty { kind: TyKind::Path(QPath::Resolved(_, path)), .. })
211-
| Node::Expr(&Expr {
212-
kind:
213-
ExprKind::Path(QPath::Resolved(_, path))
214-
| ExprKind::Struct(&QPath::Resolved(_, path), ..),
215-
..
216-
})
217-
| Node::TraitRef(&TraitRef { path, .. }) => &*path,
218-
_ => {
219-
return tcx.ty_error_with_message(
220-
DUMMY_SP,
221-
&format!("unexpected const parent path {:?}", parent_node),
222-
);
223-
}
224-
};
225-
226-
// We've encountered an `AnonConst` in some path, so we need to
227-
// figure out which generic parameter it corresponds to and return
228-
// the relevant type.
229-
230-
let (arg_index, segment) = path
231-
.segments
232-
.iter()
233-
.filter_map(|seg| seg.args.as_ref().map(|args| (args.args, seg)))
234-
.find_map(|(args, seg)| {
235-
args.iter()
236-
.filter(|arg| arg.is_const())
237-
.enumerate()
238-
.filter(|(_, arg)| arg.id() == hir_id)
239-
.map(|(index, _)| (index, seg))
240-
.next()
241-
})
242-
.unwrap_or_else(|| {
243-
bug!("no arg matching AnonConst in path");
244-
});
245-
246-
// Try to use the segment resolution if it is valid, otherwise we
247-
// default to the path resolution.
248-
let res = segment.res.filter(|&r| r != Res::Err).unwrap_or(path.res);
249-
let generics = match res {
250-
Res::Def(DefKind::Ctor(..), def_id) => {
251-
tcx.generics_of(tcx.parent(def_id).unwrap())
252-
}
253-
Res::Def(_, def_id) => tcx.generics_of(def_id),
254-
res => {
255-
return tcx.ty_error_with_message(
256-
DUMMY_SP,
257-
&format!(
258-
"unexpected anon const res {:?} in path: {:?}",
259-
res, path,
260-
),
261-
);
262-
}
263-
};
264-
265-
let ty = generics
266-
.params
267-
.iter()
268-
.filter(|param| {
269-
if let ty::GenericParamDefKind::Const = param.kind {
270-
true
271-
} else {
272-
false
273-
}
274-
})
275-
.nth(arg_index)
276-
.map(|param| tcx.type_of(param.def_id));
277-
278-
if let Some(ty) = ty {
279-
ty
280-
} else {
281-
// This is no generic parameter associated with the arg. This is
282-
// probably from an extra arg where one is not needed.
283-
tcx.ty_error_with_message(
284-
DUMMY_SP,
285-
&format!(
286-
"missing generic parameter for `AnonConst`, \
287-
parent: {:?}, res: {:?}",
288-
parent_node, res
289-
),
290-
)
291-
}
292-
}
293-
294324
x => tcx.ty_error_with_message(
295325
DUMMY_SP,
296326
&format!("unexpected const parent in type_of_def_id(): {:?}", x),

0 commit comments

Comments
 (0)