1 files changed, 13 insertions, 9 deletions
diff --git a/stage1/resolver.scala b/stage1/resolver.scala
index 0e5d221..be4d278 100644
--- a/stage1/resolver.scala
+++ b/stage1/resolver.scala
@@ -77,24 +77,28 @@ trait DependencyImplementation extends Dependency{
     )
   }
   */
-  def flatClassLoader: Boolean = false
 
-  def runMain( className: String, args: Seq[String] ): ExitCode = lib.trapExitCode{
-    lib.runMain( classLoader.loadClass( className ), args )
+  def runMain( className: String, args: Seq[String] ): ExitCode = {
+      lib.getMain( classLoader.loadClass( className ) )( args )
   }
 
-  def runMain( args: Seq[String] ): ExitCode = lib.trapExitCode{
-    mainMethod.getOrElse(
+  def runMain( args: Seq[String] ): ExitCode = {
+    val c = mainClass.getOrElse(
       throw new RuntimeException( "No main class found in " + this )
-    )( args )
+    )
+    runMain( c.getName, args )
   }
 
-  def mainMethod = lib.pickOne( "Which one do you want to run?", mainMethods )( _.name )
+  def mainClass = lib.pickOne(
+    "Which one do you want to run?",
+    classes.filter( lib.findMain(_).nonEmpty )
+  )( _.name.stripSuffix( "$" ) )
 
   def classes = exportedClasspath.files.flatMap(
-    lib.iterateClasses( _, classLoader, false )
+    lib.topLevelClasses( _, classLoader, false )
   )
-  def mainMethods = classes.flatMap( lib.discoverMain )
+
+  def flatClassLoader: Boolean = false
 
   def classLoader: ClassLoader = taskCache[DependencyImplementation]( "classLoader" ).memoize{
     if( flatClassLoader ){