Run Main.main in current process (#115)

* no message

* fix compile

* fix compile again

* finishing in process running

* remove comment

* fix bug

* fix bug

4 files changed, 94 insertions, 25 deletions

diff --git a/core/src/main/scala/mill/modules/Jvm.scala b/core/src/main/scala/mill/modules/Jvm.scala
index 48f100b2..888a687b 100644
--- a/core/src/main/scala/mill/modules/Jvm.scala
+++ b/core/src/main/scala/mill/modules/Jvm.scala
@@ -1,6 +1,7 @@
 package mill.modules
 
 import java.io.FileOutputStream
+import java.lang.reflect.Modifier
 import java.net.URLClassLoader
 import java.nio.file.attribute.PosixFilePermission
 import java.util.jar.{JarEntry, JarFile, JarOutputStream}
@@ -9,6 +10,7 @@ import ammonite.ops._
 import mill.define.Task
 import mill.eval.PathRef
 import mill.util.Ctx
+import mill.util.Ctx.LogCtx
 import mill.util.Loose.Agg
 
 import scala.annotation.tailrec
@@ -36,6 +38,61 @@ object Jvm {
     %("java", "-cp", classPath.mkString(":"), mainClass, options)
   }
 
+  def inprocess(mainClass: String,
+    classPath: Agg[Path],
+    options: Seq[String] = Seq.empty)
+    (implicit ctx: Ctx): Unit = {
+    inprocess(classPath, classLoaderOverrideSbtTesting = false, cl => {
+      getMainMethod(mainClass, cl).invoke(null, options.toArray)
+    })
+  }
+
+  private def getMainMethod(mainClassName: String, cl: ClassLoader) = {
+    val mainClass = cl.loadClass(mainClassName)
+    val method = mainClass.getMethod("main", classOf[Array[String]])
+    // jvm allows the actual main class to be non-public and to run a method in the non-public class,
+    //  we need to make it accessible
+    method.setAccessible(true)
+    val modifiers = method.getModifiers
+    if (!Modifier.isPublic(modifiers))
+      throw new NoSuchMethodException(mainClassName + ".main is not public")
+    if (!Modifier.isStatic(modifiers))
+      throw new NoSuchMethodException(mainClassName + ".main is not static")
+    method
+  }
+
+
+  def inprocess[T](classPath: Agg[Path],
+    classLoaderOverrideSbtTesting: Boolean,
+    body: ClassLoader => T): T = {
+    val cl = if (classLoaderOverrideSbtTesting) {
+      val outerClassLoader = getClass.getClassLoader
+      new URLClassLoader(
+        classPath.map(_.toIO.toURI.toURL).toArray,
+        ClassLoader.getSystemClassLoader().getParent()){
+        override def findClass(name: String) = {
+          if (name.startsWith("sbt.testing.")){
+            outerClassLoader.loadClass(name)
+          }else{
+            super.findClass(name)
+          }
+        }
+      }
+    } else {
+      new URLClassLoader(
+        classPath.map(_.toIO.toURI.toURL).toArray,
+        ClassLoader.getSystemClassLoader().getParent())
+    }
+    val oldCl = Thread.currentThread().getContextClassLoader
+    Thread.currentThread().setContextClassLoader(cl)
+    try {
+      body(cl)
+    }finally{
+      Thread.currentThread().setContextClassLoader(oldCl)
+      cl.close()
+    }
+  }
+
   def subprocess(mainClass: String,
                  classPath: Agg[Path],
                  jvmOptions: Seq[String] = Seq.empty,
diff --git a/scalalib/src/main/scala/mill/scalalib/ScalaModule.scala b/scalalib/src/main/scala/mill/scalalib/ScalaModule.scala
index 5bae1749..1b2bd28d 100644
--- a/scalalib/src/main/scala/mill/scalalib/ScalaModule.scala
+++ b/scalalib/src/main/scala/mill/scalalib/ScalaModule.scala
@@ -7,7 +7,7 @@ import mill.define.{Cross, Task}
 import mill.define.TaskModule
 import mill.eval.{PathRef, Result}
 import mill.modules.Jvm
-import mill.modules.Jvm.{createAssembly, createJar, interactiveSubprocess, subprocess}
+import mill.modules.Jvm.{createAssembly, createJar, interactiveSubprocess, subprocess, inprocess}
 import Lib._
 import mill.define.Cross.Resolver
 import mill.util.Loose.Agg
@@ -217,7 +217,15 @@ trait ScalaModule extends mill.Module with TaskModule { outer =>
 
   def forkArgs = T{ Seq.empty[String] }
 
-  def run(args: String*) = T.command{
+
+  def run(args: String*) = T.command {
+    inprocess(
+      mainClass().getOrElse(throw new RuntimeException("No mainClass provided!")),
+      runClasspath().map(_.path),
+      args)
+  }
+
+  def forkRun(args: String*) = T.command{
     subprocess(
       mainClass().getOrElse(throw new RuntimeException("No mainClass provided!")),
       runClasspath().map(_.path),
diff --git a/scalalib/src/main/scala/mill/scalalib/TestRunner.scala b/scalalib/src/main/scala/mill/scalalib/TestRunner.scala
index 11d9c79f..fa387eef 100644
--- a/scalalib/src/main/scala/mill/scalalib/TestRunner.scala
+++ b/scalalib/src/main/scala/mill/scalalib/TestRunner.scala
@@ -7,6 +7,7 @@ import java.util.zip.ZipInputStream
 
 import ammonite.ops.{Path, ls, pwd}
 import ammonite.util.Colors
+import mill.modules.Jvm
 import mill.util.Ctx.LogCtx
 import mill.util.{PrintLogger}
 import mill.util.Loose.Agg
@@ -74,23 +75,7 @@ object TestRunner {
             testClassfilePath: Agg[Path],
             args: Seq[String])
            (implicit ctx: LogCtx): (String, Seq[Result]) = {
-    val outerClassLoader = getClass.getClassLoader
-
-    val cl = new URLClassLoader(
-      entireClasspath.map(_.toIO.toURI.toURL).toArray,
-      ClassLoader.getSystemClassLoader().getParent()){
-      override def findClass(name: String) = {
-        if (name.startsWith("sbt.testing.")){
-          outerClassLoader.loadClass(name)
-        }else{
-          super.findClass(name)
-        }
-      }
-    }
-
-    val oldCl = Thread.currentThread().getContextClassLoader
-    Thread.currentThread().setContextClassLoader(cl)
-    try {
+    Jvm.inprocess(entireClasspath, classLoaderOverrideSbtTesting = true, cl => {
       val framework = cl.loadClass(frameworkName)
         .newInstance()
         .asInstanceOf[sbt.testing.Framework]
@@ -145,12 +130,7 @@ object TestRunner {
         )
       }
       (doneMsg, results)
-    }finally{
-      Thread.currentThread().setContextClassLoader(oldCl)
-      cl.close()
-    }
-
-
+    })
   }
 
   case class Result(fullyQualifiedName: String,
diff --git a/scalalib/src/test/scala/mill/scalalib/HelloWorldTests.scala b/scalalib/src/test/scala/mill/scalalib/HelloWorldTests.scala
index 70c78cee..7444604b 100644
--- a/scalalib/src/test/scala/mill/scalalib/HelloWorldTests.scala
+++ b/scalalib/src/test/scala/mill/scalalib/HelloWorldTests.scala
@@ -262,6 +262,30 @@ object HelloWorldTests extends TestSuite {
         )
       }
     }
+
+    'forkRun - {
+      'runIfMainClassProvided - {
+        val runResult = basePath / 'out / 'run / 'dest / "hello-mill"
+        val Right((_, evalCount)) = helloWorldWithMainEvaluator(
+          HelloWorldWithMain.forkRun(runResult.toString)
+        )
+
+        assert(evalCount > 0)
+
+
+        assert(
+          exists(runResult),
+          read(runResult) == "hello rockjam, your age is: 25"
+        )
+      }
+      'notRunWithoutMainClass - {
+        val Left(Result.Exception(err, _)) = helloWorldEvaluator(HelloWorld.forkRun())
+
+        assert(
+          err.isInstanceOf[RuntimeException]
+        )
+      }
+    }
     'run - {
       'runIfMainClassProvided - {
         val runResult = basePath / 'out / 'run / 'dest / "hello-mill"