@@ -38,7 +38,6 @@ use crate::ty::{InferConst, ParamConst};
38
38
use crate :: ty:: GenericParamDefKind ;
39
39
use crate :: ty:: layout:: { LayoutDetails , TargetDataLayout , VariantIdx } ;
40
40
use crate :: ty:: query;
41
- use crate :: ty:: steal:: Steal ;
42
41
use crate :: ty:: subst:: { UserSubsts , GenericArgKind } ;
43
42
use crate :: ty:: { BoundVar , BindingMode } ;
44
43
use crate :: ty:: CanonicalPolyFnSig ;
@@ -49,12 +48,13 @@ use crate::util::nodemap::{FxHashMap, FxHashSet};
49
48
use errors:: DiagnosticBuilder ;
50
49
use arena:: SyncDroplessArena ;
51
50
use smallvec:: SmallVec ;
51
+ use rustc_index:: vec:: { Idx , IndexVec } ;
52
52
use rustc_data_structures:: stable_hasher:: {
53
53
HashStable , StableHasher , StableVec , hash_stable_hashmap,
54
54
} ;
55
- use rustc_index:: vec:: { Idx , IndexVec } ;
56
55
use rustc_data_structures:: sharded:: ShardedHashMap ;
57
56
use rustc_data_structures:: sync:: { Lrc , Lock , WorkerLocal } ;
57
+ use rustc_data_structures:: steal:: Steal ;
58
58
use std:: any:: Any ;
59
59
use std:: borrow:: Borrow ;
60
60
use std:: cmp:: Ordering ;
@@ -2859,6 +2859,44 @@ fn ptr_eq<T, U>(t: *const T, u: *const U) -> bool {
2859
2859
t as * const ( ) == u as * const ( )
2860
2860
}
2861
2861
2862
+ /// In interpreter mode, locates the `DefId` of the user fn (a closure marked by an attribute named
2863
+ /// `rustc_interp_user_fn`) by visiting the local crate's HIR.
2864
+ fn find_interp_user_fn ( tcx : TyCtxt < ' _ > ) -> DefId {
2865
+ use hir:: intravisit:: { self , Visitor , NestedVisitorMap } ;
2866
+
2867
+ struct InterpUserFnVisitor < ' tcx > {
2868
+ tcx : TyCtxt < ' tcx > ,
2869
+ def_id : Option < DefId > ,
2870
+ }
2871
+
2872
+ impl < ' tcx > Visitor < ' tcx > for InterpUserFnVisitor < ' tcx > {
2873
+ fn nested_visit_map < ' this > ( & ' this mut self ) -> NestedVisitorMap < ' this , ' tcx > {
2874
+ NestedVisitorMap :: OnlyBodies ( & self . tcx . hir ( ) )
2875
+ }
2876
+
2877
+ fn visit_expr ( & mut self , ex : & ' tcx hir:: Expr ) {
2878
+ if syntax:: attr:: contains_name ( & ex. attrs , sym:: rustc_interp_user_fn) {
2879
+ self . def_id = Some ( self . tcx . hir ( ) . local_def_id ( ex. hir_id ) ) ;
2880
+ return ;
2881
+ }
2882
+
2883
+ intravisit:: walk_expr ( self , ex) ;
2884
+ }
2885
+ }
2886
+
2887
+ let mut visitor = InterpUserFnVisitor {
2888
+ tcx,
2889
+ def_id : None ,
2890
+ } ;
2891
+ tcx. hir ( ) . krate ( ) . visit_all_item_likes ( & mut visitor. as_deep_visitor ( ) ) ;
2892
+ visitor. def_id
2893
+ . unwrap_or_else ( || tcx. sess . fatal ( & format ! (
2894
+ "could not find interpreter user fn in HIR; it should be a closure expression \
2895
+ marked with the `#[{}]` attribute",
2896
+ sym:: rustc_interp_user_fn
2897
+ ) ) )
2898
+ }
2899
+
2862
2900
pub fn provide ( providers : & mut ty:: query:: Providers < ' _ > ) {
2863
2901
providers. in_scope_traits_map = |tcx, id| tcx. gcx . trait_map . get ( & id) ;
2864
2902
providers. module_exports = |tcx, id| tcx. gcx . export_map . get ( & id) . map ( |v| & v[ ..] ) ;
@@ -2936,4 +2974,8 @@ pub fn provide(providers: &mut ty::query::Providers<'_>) {
2936
2974
assert_eq ! ( cnum, LOCAL_CRATE ) ;
2937
2975
attr:: contains_name ( tcx. hir ( ) . krate_attrs ( ) , sym:: compiler_builtins)
2938
2976
} ;
2977
+ providers. interp_user_fn = |tcx, cnum| {
2978
+ assert_eq ! ( cnum, LOCAL_CRATE ) ;
2979
+ find_interp_user_fn ( tcx)
2980
+ } ;
2939
2981
}
0 commit comments