1010
1111
1212use front:: map as ast_map;
13+ use middle:: def_id:: { CRATE_DEF_INDEX } ;
1314use session:: { config, Session } ;
1415use syntax:: ast:: NodeId ;
1516use syntax:: attr;
1617use syntax:: codemap:: Span ;
1718use syntax:: entry:: EntryPointType ;
1819use rustc_front:: hir:: { Item , ItemFn } ;
19- use rustc_front:: visit;
20- use rustc_front:: visit:: Visitor ;
20+ use rustc_front:: intravisit:: Visitor ;
2121
22- struct EntryContext < ' a > {
22+ struct EntryContext < ' a , ' tcx : ' a > {
2323 session : & ' a Session ,
2424
25- // The current depth in the ast
26- depth : usize ,
25+ map : & ' a ast_map:: Map < ' tcx > ,
2726
2827 // The top-level function called 'main'
2928 main_fn : Option < ( NodeId , Span ) > ,
@@ -39,11 +38,12 @@ struct EntryContext<'a> {
3938 non_main_fns : Vec < ( NodeId , Span ) > ,
4039}
4140
42- impl < ' a , ' v > Visitor < ' v > for EntryContext < ' a > {
43- fn visit_item ( & mut self , item : & Item ) {
44- self . depth += 1 ;
45- find_item ( item, self ) ;
46- self . depth -= 1 ;
41+ impl < ' a , ' tcx > Visitor < ' tcx > for EntryContext < ' a , ' tcx > {
42+ fn visit_item ( & mut self , item : & ' tcx Item ) {
43+ let def_id = self . map . local_def_id ( item. id ) ;
44+ let def_key = self . map . def_key ( def_id) ;
45+ let at_root = def_key. parent == Some ( CRATE_DEF_INDEX ) ;
46+ find_item ( item, self , at_root) ;
4747 }
4848}
4949
@@ -64,29 +64,29 @@ pub fn find_entry_point(session: &Session, ast_map: &ast_map::Map) {
6464
6565 let mut ctxt = EntryContext {
6666 session : session,
67- depth : 0 ,
67+ map : ast_map ,
6868 main_fn : None ,
6969 attr_main_fn : None ,
7070 start_fn : None ,
7171 non_main_fns : Vec :: new ( ) ,
7272 } ;
7373
74- visit :: walk_crate ( & mut ctxt , ast_map. krate ( ) ) ;
74+ ast_map. krate ( ) . visit_all_items ( & mut ctxt ) ;
7575
7676 configure_main ( & mut ctxt) ;
7777}
7878
7979// Beware, this is duplicated in libsyntax/entry.rs, make sure to keep
8080// them in sync.
81- fn entry_point_type ( item : & Item , depth : usize ) -> EntryPointType {
81+ fn entry_point_type ( item : & Item , at_root : bool ) -> EntryPointType {
8282 match item. node {
8383 ItemFn ( ..) => {
8484 if attr:: contains_name ( & item. attrs , "start" ) {
8585 EntryPointType :: Start
8686 } else if attr:: contains_name ( & item. attrs , "main" ) {
8787 EntryPointType :: MainAttr
8888 } else if item. name . as_str ( ) == "main" {
89- if depth == 1 {
89+ if at_root {
9090 // This is a top-level function so can be 'main'
9191 EntryPointType :: MainNamed
9292 } else {
@@ -101,8 +101,8 @@ fn entry_point_type(item: &Item, depth: usize) -> EntryPointType {
101101}
102102
103103
104- fn find_item ( item : & Item , ctxt : & mut EntryContext ) {
105- match entry_point_type ( item, ctxt . depth ) {
104+ fn find_item ( item : & Item , ctxt : & mut EntryContext , at_root : bool ) {
105+ match entry_point_type ( item, at_root ) {
106106 EntryPointType :: MainNamed => {
107107 if ctxt. main_fn . is_none ( ) {
108108 ctxt. main_fn = Some ( ( item. id , item. span ) ) ;
@@ -132,8 +132,6 @@ fn find_item(item: &Item, ctxt: &mut EntryContext) {
132132 } ,
133133 EntryPointType :: None => ( )
134134 }
135-
136- visit:: walk_item ( ctxt, item) ;
137135}
138136
139137fn configure_main ( this : & mut EntryContext ) {
0 commit comments