Java 9 optimization by caching rt.jar (#239)

* Upgraded Ammonite to 1.0.5-7-f032887 that includes a Java 9 performance optimization for mill's client-server mode (lihaoyi/Ammonite#773).

* Java 9 optimization for interactive mode by caching rt.jar.

* Refactored based on @lihaoyi's comments on #239.

* Removed client side rt.jar caching.

* Propagatio of home path from Cli.

* Propagation, propagation, propagation!

17 files changed, 74 insertions, 46 deletions

diff --git a/build.sc b/build.sc
index c012c77b..d0d3ea3a 100755
--- a/build.sc
+++ b/build.sc
@@ -69,7 +69,7 @@ trait MillModule extends MillPublishModule{ outer =>
 
 object clientserver extends MillModule{
   def ivyDeps = Agg(
-    ivy"com.lihaoyi:::ammonite:1.0.5-4-c0cdbaf",
+    ivy"com.lihaoyi:::ammonite:1.0.5-7-f032887",
     ivy"org.scala-sbt.ipcsocket:ipcsocket:1.0.0"
   )
   val test = new Tests(implicitly)
@@ -84,7 +84,7 @@ object core extends MillModule {
 
   def ivyDeps = Agg(
     ivy"com.lihaoyi::sourcecode:0.1.4",
-    ivy"com.lihaoyi:::ammonite:1.0.5-4-c0cdbaf",
+    ivy"com.lihaoyi:::ammonite:1.0.5-7-f032887",
     ivy"jline:jline:2.14.5"
   )
 
diff --git a/clientserver/src/mill/clientserver/Client.java b/clientserver/src/mill/clientserver/Client.java
index eab8c5b5..261f69b1 100644
--- a/clientserver/src/mill/clientserver/Client.java
+++ b/clientserver/src/mill/clientserver/Client.java
@@ -1,6 +1,5 @@
 package mill.clientserver;
 
-import io.github.retronym.java9rtexport.Export;
 import org.scalasbt.ipcsocket.*;
 
 import java.io.*;
@@ -8,7 +7,6 @@ import java.net.Socket;
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.nio.channels.FileChannel;
-import java.nio.file.Files;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Iterator;
@@ -29,11 +27,6 @@ public class Client {
         }
         if (!System.getProperty("java.specification.version").startsWith("1.")) {
             selfJars.addAll(Arrays.asList(System.getProperty("java.class.path").split(File.pathSeparator)));
-            File rtFile = new File(lockBase + "/rt-" + System.getProperty("java.version") + ".jar");
-            if (!rtFile.exists()) {
-                Files.copy(Export.export().toPath(), rtFile.toPath());
-            }
-            selfJars.add(rtFile.getCanonicalPath());
         }
         ArrayList<String> l = new java.util.ArrayList<String>();
         l.add("java");
diff --git a/core/src/mill/eval/Evaluator.scala b/core/src/mill/eval/Evaluator.scala
index e61c2c4c..1e1b8b35 100644
--- a/core/src/mill/eval/Evaluator.scala
+++ b/core/src/mill/eval/Evaluator.scala
@@ -26,7 +26,8 @@ case class Labelled[T](task: NamedTask[T],
     case _ => None
   }
 }
-case class Evaluator[T](outPath: Path,
+case class Evaluator[T](home: Path,
+                        outPath: Path,
                         externalOutPath: Path,
                         rootModule: mill.define.BaseModule,
                         log: Logger,
@@ -265,7 +266,8 @@ case class Evaluator[T](outPath: Path,
                     throw new Exception("No `dest` folder available here")
                 }
             },
-            multiLogger
+            multiLogger,
+            home
           )
 
           val out = System.out
diff --git a/core/src/mill/util/ClassLoader.scala b/core/src/mill/util/ClassLoader.scala
index a1b42158..b53150c2 100644
--- a/core/src/mill/util/ClassLoader.scala
+++ b/core/src/mill/util/ClassLoader.scala
@@ -5,7 +5,9 @@ import java.net.{URL, URLClassLoader}
 import io.github.retronym.java9rtexport.Export
 
 object ClassLoader {
-  def create(urls: Seq[URL], parent: java.lang.ClassLoader): URLClassLoader = {
+  def create(urls: Seq[URL],
+             parent: java.lang.ClassLoader)
+            (implicit ctx: Ctx.Home): URLClassLoader = {
     val cl = new URLClassLoader(urls.toArray, parent)
     if (!ammonite.util.Util.java9OrAbove) return cl
     try {
@@ -13,7 +15,12 @@ object ClassLoader {
       cl
     } catch {
       case _: ClassNotFoundException =>
-        new URLClassLoader((urls ++ Some(Export.export().toURI.toURL)).toArray, parent)
+        val path = ctx.home
+        val rtFile = new java.io.File(path.toIO, s"rt-${System.getProperty("java.version")}.jar")
+        if (!rtFile.exists) {
+          java.nio.file.Files.copy(Export.export().toPath, rtFile.toPath)
+        }
+        new URLClassLoader((urls ++ Some(rtFile.toURI.toURL)).toArray, parent)
     }
   }
 }
diff --git a/core/src/mill/util/Ctx.scala b/core/src/mill/util/Ctx.scala
index 55cc05ab..99818194 100644
--- a/core/src/mill/util/Ctx.scala
+++ b/core/src/mill/util/Ctx.scala
@@ -20,6 +20,9 @@ object Ctx{
   trait Log{
     def log: Logger
   }
+  trait Home{
+    def home: Path
+  }
   object Log{
     implicit def logToCtx(l: Logger): Log = new Log { def log = l }
   }
@@ -27,13 +30,17 @@ object Ctx{
     def args: IndexedSeq[_]
   }
 
+  def defaultHome = ammonite.ops.home / ".mill" / "ammonite"
+
 }
 class Ctx(val args: IndexedSeq[_],
           dest0: () => Path,
-          val log: Logger)
+          val log: Logger,
+          val home: Path)
   extends Ctx.Dest
   with Ctx.Log
-  with Ctx.Args{
+  with Ctx.Args
+  with Ctx.Home{
 
   def dest = dest0()
   def length = args.length
diff --git a/main/src/mill/Main.scala b/main/src/mill/Main.scala
index 6c10f6f2..a899c8c6 100644
--- a/main/src/mill/Main.scala
+++ b/main/src/mill/Main.scala
@@ -59,7 +59,7 @@ object Main {
     val millArgSignature =
       Cli.genericSignature.filter(a => !removed(a.name)) :+ interactiveSignature
 
-    val millHome = home / ".mill" / "ammonite"
+    val millHome = mill.util.Ctx.defaultHome
 
     Cli.groupArgs(
       args.toList,
@@ -88,12 +88,14 @@ object Main {
           stderr.println("Build repl needs to be run with the -i/--interactive flag")
           (false, stateCache)
         }else{
+          val tqs = "\"\"\""
           val config =
             if(!repl) cliConfig
             else cliConfig.copy(
               predefCode =
-                """import $file.build, build._
+                s"""import $$file.build, build._
                   |implicit val replApplyHandler = mill.main.ReplApplyHandler(
+                  |  ammonite.ops.Path($tqs${cliConfig.home.toIO.getCanonicalPath.replaceAllLiterally("$", "$$")}$tqs),
                   |  interp.colors(),
                   |  repl.pprinter(),
                   |  build.millSelf.get,
diff --git a/main/src/mill/main/MainRunner.scala b/main/src/mill/main/MainRunner.scala
index ff7bcbbd..efebd5a5 100644
--- a/main/src/mill/main/MainRunner.scala
+++ b/main/src/mill/main/MainRunner.scala
@@ -62,6 +62,7 @@ class MainRunner(val config: ammonite.main.Cli.Config,
       printing = true,
       mainCfg => {
         val (result, interpWatched) = RunScript.runScript(
+          config.home,
           mainCfg.wd,
           scriptPath,
           mainCfg.instantiateInterpreter(),
diff --git a/main/src/mill/main/ReplApplyHandler.scala b/main/src/mill/main/ReplApplyHandler.scala
index a2b042ad..9ce30142 100644
--- a/main/src/mill/main/ReplApplyHandler.scala
+++ b/main/src/mill/main/ReplApplyHandler.scala
@@ -1,6 +1,7 @@
 package mill.main
 
 
+import ammonite.ops.Path
 import mill.define.Applicative.ApplyHandler
 import mill.define.Segment.Label
 import mill.define._
@@ -9,13 +10,15 @@ import mill.util.Strict.Agg
 
 import scala.collection.mutable
 object ReplApplyHandler{
-  def apply[T](colors: ammonite.util.Colors,
+  def apply[T](home: Path,
+               colors: ammonite.util.Colors,
                pprinter0: pprint.PPrinter,
                rootModule: mill.define.BaseModule,
                discover: Discover[_]) = {
     new ReplApplyHandler(
       pprinter0,
       new Evaluator(
+        home,
         ammonite.ops.pwd / 'out,
         ammonite.ops.pwd / 'out,
         rootModule,
diff --git a/main/src/mill/main/RunScript.scala b/main/src/mill/main/RunScript.scala
index 553f5b69..4d8653fb 100644
--- a/main/src/mill/main/RunScript.scala
+++ b/main/src/mill/main/RunScript.scala
@@ -23,7 +23,8 @@ import scala.reflect.ClassTag
   * subsystem
   */
 object RunScript{
-  def runScript(wd: Path,
+  def runScript(home: Path,
+                wd: Path,
                 path: Path,
                 instantiateInterpreter: => Either[(Res.Failing, Seq[(Path, Long)]), ammonite.interp.Interpreter],
                 scriptArgs: Seq[String],
@@ -52,7 +53,7 @@ object RunScript{
 
     val evalRes =
       for(s <- evalState)
-      yield new Evaluator[Any](wd / 'out, wd / 'out, s.rootModule, log, s.classLoaderSig, s.workerCache)
+      yield new Evaluator[Any](home, wd / 'out, wd / 'out, s.rootModule, log, s.classLoaderSig, s.workerCache)
 
     val evaluated = for{
       evaluator <- evalRes
diff --git a/main/src/mill/modules/Jvm.scala b/main/src/mill/modules/Jvm.scala
index d71cf75c..2ded95db 100644
--- a/main/src/mill/modules/Jvm.scala
+++ b/main/src/mill/modules/Jvm.scala
@@ -101,7 +101,8 @@ object Jvm {
 
   def inprocess[T](classPath: Agg[Path],
                    classLoaderOverrideSbtTesting: Boolean,
-                   body: ClassLoader => T): T = {
+                   body: ClassLoader => T)
+                  (implicit ctx: Ctx.Home): T = {
     val cl = if (classLoaderOverrideSbtTesting) {
       val outerClassLoader = getClass.getClassLoader
       new URLClassLoader(classPath.map(_.toIO.toURI.toURL).toArray, mill.util.ClassLoader.create(Seq(), null)){
diff --git a/main/test/src/mill/util/TestEvaluator.scala b/main/test/src/mill/util/TestEvaluator.scala
index 078254f1..ffff4b99 100644
--- a/main/test/src/mill/util/TestEvaluator.scala
+++ b/main/test/src/mill/util/TestEvaluator.scala
@@ -30,7 +30,7 @@ class TestEvaluator[T <: TestUtil.BaseModule](module: T)
 //    true,
 //    ammonite.util.Colors.Default, System.out, System.out, System.err, System.in
 // )
-  val evaluator = new Evaluator(outPath, TestEvaluator.externalOutPath, module, logger)
+  val evaluator = new Evaluator(Ctx.defaultHome, outPath, TestEvaluator.externalOutPath, module, logger)
 
   def apply[T](t: Task[T]): Either[Result.Failing[T], (T, Int)] = {
     val evaluated = evaluator.evaluate(Agg(t))
diff --git a/scalajslib/src/mill/scalajslib/ScalaJSBridge.scala b/scalajslib/src/mill/scalajslib/ScalaJSBridge.scala
index 507d8af7..59ed6450 100644
--- a/scalajslib/src/mill/scalajslib/ScalaJSBridge.scala
+++ b/scalajslib/src/mill/scalajslib/ScalaJSBridge.scala
@@ -5,6 +5,7 @@ import java.net.URLClassLoader
 
 import ammonite.ops.Path
 import mill.define.Discover
+import mill.util.Ctx
 import mill.{Agg, T}
 
 sealed trait OptimizeMode
@@ -21,7 +22,8 @@ object ModuleKind{
 class ScalaJSWorker {
   private var scalaInstanceCache = Option.empty[(Long, ScalaJSBridge)]
 
-  private def bridge(toolsClasspath: Agg[Path]) = {
+  private def bridge(toolsClasspath: Agg[Path])
+                    (implicit ctx: Ctx.Home) = {
     val classloaderSig =
       toolsClasspath.map(p => p.toString().hashCode + p.mtime.toMillis).sum
     scalaInstanceCache match {
@@ -47,7 +49,8 @@ class ScalaJSWorker {
            dest: File,
            main: Option[String],
            fullOpt: Boolean,
-           moduleKind: ModuleKind): Unit = {
+           moduleKind: ModuleKind)
+          (implicit ctx: Ctx.Home): Unit = {
     bridge(toolsClasspath).link(
       sources.items.map(_.toIO).toArray,
       libraries.items.map(_.toIO).toArray,
@@ -58,14 +61,16 @@ class ScalaJSWorker {
     )
   }
 
-  def run(toolsClasspath: Agg[Path], config: NodeJSConfig, linkedFile: File): Unit = {
+  def run(toolsClasspath: Agg[Path], config: NodeJSConfig, linkedFile: File)
+         (implicit ctx: Ctx.Home): Unit = {
     bridge(toolsClasspath).run(config, linkedFile)
   }
 
   def getFramework(toolsClasspath: Agg[Path],
                    config: NodeJSConfig,
                    frameworkName: String,
-                   linkedFile: File): (() => Unit, sbt.testing.Framework) = {
+                   linkedFile: File)
+                  (implicit ctx: Ctx.Home): (() => Unit, sbt.testing.Framework) = {
     bridge(toolsClasspath).getFramework(config, frameworkName, linkedFile)
   }
 
diff --git a/scalajslib/src/mill/scalajslib/ScalaJSModule.scala b/scalajslib/src/mill/scalajslib/ScalaJSModule.scala
index 92c071c2..099b8bc2 100644
--- a/scalajslib/src/mill/scalajslib/ScalaJSModule.scala
+++ b/scalajslib/src/mill/scalajslib/ScalaJSModule.scala
@@ -117,7 +117,7 @@ trait ScalaJSModule extends scalalib.ScalaModule { outer =>
            runClasspath: Agg[PathRef],
            mainClass: Option[String],
            mode: OptimizeMode,
-           moduleKind: ModuleKind)(implicit ctx: Ctx.Dest): PathRef = {
+           moduleKind: ModuleKind)(implicit ctx: Ctx): PathRef = {
     val outputPath = ctx.dest / "out.js"
 
     mkdir(ctx.dest)
diff --git a/scalalib/src/mill/scalalib/GenIdea.scala b/scalalib/src/mill/scalalib/GenIdea.scala
index b096ec67..30456c05 100644
--- a/scalalib/src/mill/scalalib/GenIdea.scala
+++ b/scalalib/src/mill/scalalib/GenIdea.scala
@@ -6,9 +6,10 @@ import coursier.maven.MavenRepository
 import mill.define._
 import mill.eval.{Evaluator, PathRef, Result}
 import mill.{T, scalalib}
-import mill.util.Ctx.Log
+import mill.util.Ctx.{Home, Log}
 import mill.util.{Loose, PrintLogger, Strict}
 import mill.util.Strict.Agg
+
 import scala.util.Try
 
 
@@ -27,7 +28,7 @@ object GenIdeaModule extends ExternalModule {
 }
 object GenIdea {
 
-  def apply(ctx: Log,
+  def apply(ctx: Log with Home,
             rootModule: BaseModule,
             discover: Discover[_]): Unit = {
     val pp = new scala.xml.PrettyPrinter(999, 4)
@@ -38,7 +39,7 @@ object GenIdea {
     rm! pwd/".idea_modules"
 
 
-    val evaluator = new Evaluator(pwd / 'out, pwd / 'out, rootModule, ctx.log)
+    val evaluator = new Evaluator(ctx.home, pwd / 'out, pwd / 'out, rootModule, ctx.log)
 
     for((relPath, xml) <- xmlFileLayout(evaluator, rootModule, jdkInfo)){
       write.over(pwd/relPath, pp.format(xml))
diff --git a/scalalib/src/mill/scalalib/ScalaModule.scala b/scalalib/src/mill/scalalib/ScalaModule.scala
index 58a325d0..06b6f0c6 100644
--- a/scalalib/src/mill/scalalib/ScalaModule.scala
+++ b/scalalib/src/mill/scalalib/ScalaModule.scala
@@ -319,7 +319,7 @@ trait ScalaModule extends mill.Module with TaskModule { outer =>
   }
 
   def ammoniteReplClasspath = T{
-    resolveDeps(T.task{Agg(ivy"com.lihaoyi:::ammonite:1.0.5-4-c0cdbaf")})()
+    resolveDeps(T.task{Agg(ivy"com.lihaoyi:::ammonite:1.0.5-7-f032887")})()
   }
   def repl() = T.command{
     if (T.ctx().log.inStream == DummyInputStream){
@@ -379,6 +379,7 @@ trait TestModule extends ScalaModule with TaskModule {
       jvmArgs = forkArgs(),
       envArgs = forkEnv(),
       mainArgs = Seq(
+        T.ctx().home.toString,
         testFrameworks().mkString(" "),
         runClasspath().map(_.path).mkString(" "),
         Seq(compile().classes.path).mkString(" "),
diff --git a/scalalib/src/mill/scalalib/ScalaWorkerApi.scala b/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
index 91222524..84db1dd8 100644
--- a/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
+++ b/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
@@ -67,7 +67,7 @@ trait ScalaWorkerApi {
                entireClasspath: Agg[Path],
                testClassfilePath: Agg[Path],
                args: Seq[String])
-              (implicit ctx: mill.util.Ctx.Log): (String, Seq[Result])
+              (implicit ctx: mill.util.Ctx.Log with mill.util.Ctx.Home): (String, Seq[Result])
 
   def discoverMainClasses(compilationResult: CompilationResult)
                          (implicit ctx: mill.util.Ctx): Seq[String]
diff --git a/scalaworker/src/mill/scalaworker/ScalaWorker.scala b/scalaworker/src/mill/scalaworker/ScalaWorker.scala
index ca20ccf1..c74768a5 100644
--- a/scalaworker/src/mill/scalaworker/ScalaWorker.scala
+++ b/scalaworker/src/mill/scalaworker/ScalaWorker.scala
@@ -36,24 +36,28 @@ object ScalaWorker{
 
   def main(args: Array[String]): Unit = {
     try{
-      val (frameworks, classpath, testCp, arguments, outputPath, colored) = args match {
-        case Array(fs, cp, tcp, op, c) => (fs, cp, tcp, "", op, c)
-        case Array(fs, cp, tcp, as, op, c) => (fs, cp, tcp, as, op, c)
+      val (hm, frameworks, classpath, testCp, arguments, outputPath, colored) = args match {
+        case Array(h, fs, cp, tcp, op, c) => (h, fs, cp, tcp, "", op, c)
+        case Array(h, fs, cp, tcp, as, op, c) => (h, fs, cp, tcp, as, op, c)
+      }
+      val ctx = new Ctx.Log with Ctx.Home {
+        val log = PrintLogger(
+          colored == "true",
+          if(colored == "true") Colors.Default
+          else Colors.BlackWhite,
+          System.out,
+          System.err,
+          System.err,
+          System.in
+        )
+        val home = Path(hm)
       }
       val result = new ScalaWorker(null, null).runTests(
         frameworkInstances = TestRunner.frameworks(frameworks.split(" ")),
         entireClasspath = Agg.from(classpath.split(" ").map(Path(_))),
         testClassfilePath = Agg.from(testCp.split(" ").map(Path(_))),
         args = arguments match{ case "" => Nil case x => x.split(" ").toList }
-      )(new PrintLogger(
-        colored == "true",
-        if(colored == "true") Colors.Default
-        else Colors.BlackWhite,
-        System.out,
-        System.err,
-        System.err,
-        System.in
-      ))
+      )(ctx)
 
       ammonite.ops.write(Path(outputPath), upickle.default.write(result))
     }catch{case e: Throwable =>
@@ -89,7 +93,7 @@ class ScalaWorker(ctx0: mill.util.Ctx,
         .get
 
       val sourceFolder = mill.modules.Util.unpackZip(sourceJar)(workingDir)
-      val classloader = mill.util.ClassLoader.create(compilerJars.map(_.toURI.toURL), null)
+      val classloader = mill.util.ClassLoader.create(compilerJars.map(_.toURI.toURL), null)(ctx0)
       val scalacMain = classloader.loadClass("scala.tools.nsc.Main")
       val argsArray = Array[String](
         "-d", compiledDest.toString,
@@ -250,7 +254,7 @@ class ScalaWorker(ctx0: mill.util.Ctx,
                entireClasspath: Agg[Path],
                testClassfilePath: Agg[Path],
                args: Seq[String])
-              (implicit ctx: mill.util.Ctx.Log): (String, Seq[Result]) = {
+              (implicit ctx: Ctx.Log with Ctx.Home): (String, Seq[Result]) = {
 
     Jvm.inprocess(entireClasspath, classLoaderOverrideSbtTesting = true, cl => {
       val frameworks = frameworkInstances(cl)