run start and main language item if provided

src/bin/miri.rs

@@ -19,11 +19,7 @@ use miri::{
 use rustc::session::Session;
 use rustc_driver::{driver, CompilerCalls};
 use rustc::ty::{TyCtxt, subst};
-use rustc::mir::mir_map::MirMap;
-use rustc::mir::repr::Mir;
 use rustc::hir::def_id::DefId;
-use rustc::hir::{map, ItemFn, Item};
-use syntax::codemap::Span;
 
 struct MiriCompilerCalls;
 
@@ -40,9 +36,12 @@ impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
 
             let tcx = state.tcx.unwrap();
             let mir_map = state.mir_map.unwrap();
-            let (span, mir, def_id) = get_main(tcx, mir_map);
+
+            let (node_id, span) = state.session.entry_fn.borrow().expect("no main or start function found");
             println!("found `main` function at: {:?}", span);
 
+            let mir = mir_map.map.get(&node_id).expect("no mir for main function");
+            let def_id = tcx.map.local_def_id(node_id);
             let mut ecx = EvalContext::new(tcx, mir_map);
             let substs = tcx.mk_substs(subst::Substs::empty());
             let return_ptr = ecx.alloc_ret_ptr(mir.return_ty, substs).expect("main function should not be diverging");
@@ -66,19 +65,6 @@ impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
     }
 }
 
-fn get_main<'a, 'b, 'tcx: 'b>(tcx: TyCtxt<'a, 'tcx, 'tcx>, mir_map: &'b MirMap<'tcx>) -> (Span, &'b Mir<'tcx>, DefId) {
-    for (&id, mir) in &mir_map.map {
-        if let map::Node::NodeItem(&Item { name, span, ref node, .. }) = tcx.map.get(id) {
-            if let ItemFn(..) = *node {
-                if name.as_str() == "main" {
-                    return (span, mir, tcx.map.local_def_id(id));
-                }
-            }
-        }
-    }
-    panic!("no main function found");
-}
-
 fn report(tcx: TyCtxt, ecx: &EvalContext, e: EvalError) {
     let frame = ecx.stack().last().expect("stackframe was empty");
     let block = &frame.mir.basic_blocks()[frame.next_block];

tests/run-pass/main_fn.rs

@@ -0,0 +1,5 @@
+#![feature(main)]
+
+#[main]
+fn foo() {
+}

tests/run-pass/start_fn.rs

@@ -0,0 +1,6 @@
+#![feature(start)]
+
+#[start]
+fn foo(_nargs: isize, _args: *const *const u8) -> isize {
+    return 0;
+}