Merge branch 'master' of github.com:lihaoyi/mill

1 files changed, 53 insertions, 15 deletions

diff --git a/core/src/mill/util/ClassLoader.scala b/core/src/mill/util/ClassLoader.scala
index b53150c2..d650aceb 100644
--- a/core/src/mill/util/ClassLoader.scala
+++ b/core/src/mill/util/ClassLoader.scala
@@ -2,25 +2,63 @@ package mill.util
 
 import java.net.{URL, URLClassLoader}
 
+import ammonite.ops._
 import io.github.retronym.java9rtexport.Export
 
 object ClassLoader {
+
+  def create(urls: Seq[URL], parent: java.lang.ClassLoader)(
+      implicit ctx: Ctx.Home): URLClassLoader = {
+    new URLClassLoader(
+      makeUrls(urls).toArray,
+      refinePlatformParent(parent)
+    )
+  }
+
   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 {
-      cl.loadClass("javax.script.ScriptEngineManager")
-      cl
-    } catch {
-      case _: ClassNotFoundException =>
-        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)
+             parent: java.lang.ClassLoader,
+             customFindClass: String => Option[Class[_]])(
+      implicit ctx: Ctx.Home): URLClassLoader = {
+    new URLClassLoader(
+      makeUrls(urls).toArray,
+      refinePlatformParent(parent)
+    ) {
+      override def findClass(name: String): Class[_] = {
+        customFindClass(name).getOrElse(super.findClass(name))
+      }
+    }
+  }
+
+  /**
+    *  Return `ClassLoader.getPlatformClassLoader` for java 9 and above, if parent class loader is null,
+    *  otherwise return same parent class loader.
+    *  More details: https://docs.oracle.com/javase/9/migrate/toc.htm#JSMIG-GUID-A868D0B9-026F-4D46-B979-901834343F9E
+    *
+    *  `ClassLoader.getPlatformClassLoader` call is implemented via runtime reflection, cause otherwise
+    *  mill could be compiled only with jdk 9 or above. We don't want to introduce this restriction now.
+    */
+  private def refinePlatformParent(parent: java.lang.ClassLoader): ClassLoader = {
+    if (ammonite.util.Util.java9OrAbove) {
+      if (parent == null)
+        classOf[ClassLoader]
+          .getMethod("getPlatformClassLoader")
+          .invoke(null)
+          .asInstanceOf[ClassLoader]
+      else parent
+    } else {
+      parent
+    }
+  }
+
+  private def makeUrls(urls: Seq[URL])(implicit ctx: Ctx.Home): Seq[URL] = {
+    if (ammonite.util.Util.java9OrAbove) {
+      val rtFile = ctx.home / s"rt-${System.getProperty("java.version")}.jar"
+      if (!exists(rtFile)) {
+        cp(Path(Export.export()), rtFile)
+      }
+      urls :+ rtFile.toNIO.toUri.toURL
+    } else {
+      urls
     }
   }
 }