Split out `upstreamAssembly` from `assembly`

Also re-write `Jvm.createAssembly` to allow incremental assembly construction. This should allow much faster assembly creation in the common case where upstream dependencies do not change

8 files changed, 157 insertions, 124 deletions

diff --git a/core/src/mill/eval/PathRef.scala b/core/src/mill/eval/PathRef.scala
index 168eaa9a..0fbc18f8 100644
--- a/core/src/mill/eval/PathRef.scala
+++ b/core/src/mill/eval/PathRef.scala
@@ -4,10 +4,11 @@ import java.io.IOException
 import java.nio.file.attribute.BasicFileAttributes
 import java.nio.file.{FileVisitResult, FileVisitor}
 import java.nio.{file => jnio}
-import java.security.MessageDigest
+import java.security.{DigestOutputStream, MessageDigest}
+
 import upickle.default.{ReadWriter => RW}
 import ammonite.ops.Path
-import mill.util.JsonFormatters
+import mill.util.{DummyOutputStream, IO, JsonFormatters}
 
 
 /**
@@ -23,8 +24,7 @@ object PathRef{
   def apply(path: ammonite.ops.Path, quick: Boolean = false) = {
     val sig = {
       val digest = MessageDigest.getInstance("MD5")
-
-      val buffer = new Array[Byte](16 * 1024)
+      val digestOut = new DigestOutputStream(DummyOutputStream, digest)
       jnio.Files.walkFileTree(
         path.toNIO,
         new FileVisitor[jnio.Path] {
@@ -43,16 +43,7 @@ object PathRef{
               digest.update(value.toByte)
             }else {
               val is = jnio.Files.newInputStream(file)
-
-              def rec(): Unit = {
-                val length = is.read(buffer)
-                if (length != -1) {
-                  digest.update(buffer, 0, length)
-                  rec()
-                }
-              }
-              rec()
-
+              IO.stream(is, digestOut)
               is.close()
             }
             FileVisitResult.CONTINUE
diff --git a/core/src/mill/util/DummyInputStream.scala b/core/src/mill/util/DummyInputStream.scala
deleted file mode 100644
index 310b358b..00000000
--- a/core/src/mill/util/DummyInputStream.scala
+++ /dev/null
@@ -1,5 +0,0 @@
-package mill.util
-
-import java.io.ByteArrayInputStream
-
-object DummyInputStream extends ByteArrayInputStream(Array())
\ No newline at end of file
diff --git a/core/src/mill/util/IO.scala b/core/src/mill/util/IO.scala
new file mode 100644
index 00000000..833e52c7
--- /dev/null
+++ b/core/src/mill/util/IO.scala
@@ -0,0 +1,32 @@
+package mill.util
+
+import java.io.{InputStream, OutputStream}
+
+import scala.tools.nsc.interpreter.OutputStream
+
+/**
+  * Misc IO utilities, eventually probably should be pushed upstream into
+  * ammonite-ops
+  */
+object IO {
+  def stream(src: InputStream, dest: OutputStream) = {
+    val buffer = new Array[Byte](4096)
+    while ( {
+      src.read(buffer) match {
+        case -1 => false
+        case n =>
+          dest.write(buffer, 0, n)
+          true
+      }
+    }) ()
+  }
+}
+
+import java.io.{ByteArrayInputStream, OutputStream}
+
+object DummyInputStream extends ByteArrayInputStream(Array())
+object DummyOutputStream extends OutputStream{
+  override def write(b: Int) = ()
+  override def write(b: Array[Byte]) = ()
+  override def write(b: Array[Byte], off: Int, len: Int) = ()
+}
diff --git a/main/src/mill/modules/Jvm.scala b/main/src/mill/modules/Jvm.scala
index 57e02dd4..c8473b65 100644
--- a/main/src/mill/modules/Jvm.scala
+++ b/main/src/mill/modules/Jvm.scala
@@ -2,22 +2,19 @@ package mill.modules
 
 import java.io.{ByteArrayInputStream, FileOutputStream}
 import java.lang.reflect.Modifier
-import java.net.URLClassLoader
+import java.net.{URI, URLClassLoader}
+import java.nio.file.{FileSystems, Files, OpenOption, StandardOpenOption}
 import java.nio.file.attribute.PosixFilePermission
 import java.util.jar.{JarEntry, JarFile, JarOutputStream}
 
 import ammonite.ops._
-import mill.clientserver.{ClientServer, InputPumper}
-import mill.define.Task
+import geny.Generator
+import mill.clientserver.InputPumper
 import mill.eval.PathRef
-import mill.util.{Ctx, Loose}
-import mill.util.Ctx.Log
+import mill.util.{Ctx, IO}
 import mill.util.Loose.Agg
-import upickle.default.{Reader, Writer}
 
-import scala.annotation.tailrec
 import scala.collection.mutable
-import scala.reflect.ClassTag
 
 
 object Jvm {
@@ -231,68 +228,90 @@ object Jvm {
     PathRef(outputPath)
   }
 
+  def newOutputStream(p: java.nio.file.Path) = Files.newOutputStream(
+    p,
+    StandardOpenOption.TRUNCATE_EXISTING,
+    StandardOpenOption.CREATE
+  )
+
   def createAssembly(inputPaths: Agg[Path],
                      mainClass: Option[String] = None,
-                     prependShellScript: String = "")
-                    (implicit ctx: Ctx.Dest): PathRef = {
-    val outputPath = ctx.dest / "out.jar"
-    rm(outputPath)
+                     prependShellScript: String = "",
+                     base: Option[Path] = None)
+                    (implicit ctx: Ctx.Dest) = {
+    val tmp = ctx.dest / "out-tmp.jar"
 
-    if(inputPaths.nonEmpty) {
+    val baseUri = "jar:file:" + tmp
+    val hm = new java.util.HashMap[String, String]()
 
-      val output = new FileOutputStream(outputPath.toIO)
+    base match{
+      case Some(b) => cp(b, tmp)
+      case None => hm.put("create", "true")
+    }
 
-      // Prepend shell script and make it executable
-      if (prependShellScript.nonEmpty) {
-        output.write((prependShellScript + "\n").getBytes)
-        val perms = java.nio.file.Files.getPosixFilePermissions(outputPath.toNIO)
-        perms.add(PosixFilePermission.GROUP_EXECUTE)
-        perms.add(PosixFilePermission.OWNER_EXECUTE)
-        perms.add(PosixFilePermission.OTHERS_EXECUTE)
-        java.nio.file.Files.setPosixFilePermissions(outputPath.toNIO, perms)
-      }
+    val zipFs = FileSystems.newFileSystem(URI.create(baseUri), hm)
+
+    val manifest = createManifest(mainClass)
+    val manifestPath = zipFs.getPath(JarFile.MANIFEST_NAME)
+    Files.createDirectories(manifestPath.getParent)
+    val manifestOut = newOutputStream(manifestPath)
+    manifest.write(manifestOut)
+    manifestOut.close()
+
+    for(v <- classpathIterator(inputPaths)){
+      val (file, mapping) = v
+      val p = zipFs.getPath(mapping)
+      if (p.getParent != null) Files.createDirectories(p.getParent)
+      val outputStream = newOutputStream(p)
+      IO.stream(file, outputStream)
+      outputStream.close()
+      file.close()
+    }
+    zipFs.close()
+    val output = ctx.dest / "out.jar"
+
+    // Prepend shell script and make it executable
+    if (prependShellScript.isEmpty) mv(tmp, output)
+    else{
+      val outputStream = newOutputStream(output.toNIO)
+      IO.stream(new ByteArrayInputStream((prependShellScript + "\n").getBytes()), outputStream)
+      IO.stream(read.getInputStream(tmp), outputStream)
+      outputStream.close()
+
+      val perms = Files.getPosixFilePermissions(output.toNIO)
+      perms.add(PosixFilePermission.GROUP_EXECUTE)
+      perms.add(PosixFilePermission.OWNER_EXECUTE)
+      perms.add(PosixFilePermission.OTHERS_EXECUTE)
+      Files.setPosixFilePermissions(output.toNIO, perms)
+    }
 
-      val jar = new JarOutputStream(
-        output,
-        createManifest(mainClass)
-      )
+    PathRef(output)
+  }
 
-      val seen = mutable.Set("META-INF/MANIFEST.MF")
-      try{
-
-
-        for{
-          p <- inputPaths
-          if exists(p)
-          (file, mapping) <-
-            if (p.isFile) {
-              val jf = new JarFile(p.toIO)
-              import collection.JavaConverters._
-              for(entry <- jf.entries().asScala if !entry.isDirectory) yield {
-                read.bytes(jf.getInputStream(entry)) -> entry.getName
-              }
-            }
-            else {
-              ls.rec(p).iterator
-                .filter(_.isFile)
-                .map(sub => read.bytes(sub) -> sub.relativeTo(p).toString)
-            }
-          if !seen(mapping)
-        } {
-          seen.add(mapping)
-          val entry = new JarEntry(mapping.toString)
-          jar.putNextEntry(entry)
-          jar.write(file)
-          jar.closeEntry()
-        }
-      } finally {
-        jar.close()
-        output.close()
+
+  def classpathIterator(inputPaths: Agg[Path]) = {
+    Generator.from(inputPaths)
+      .filter(exists)
+      .flatMap{
+        p =>
+          if (p.isFile) {
+            val jf = new JarFile(p.toIO)
+            import collection.JavaConverters._
+            Generator.selfClosing((
+              for(entry <- jf.entries().asScala if !entry.isDirectory)
+                yield (jf.getInputStream(entry), entry.getName),
+              () => jf.close()
+            ))
+          }
+          else {
+            ls.rec.iter(p)
+              .filter(_.isFile)
+              .map(sub => read.getInputStream(sub) -> sub.relativeTo(p).toString)
+          }
       }
 
-    }
-    PathRef(outputPath)
   }
+
   def launcherShellScript(mainClass: String,
                           classPath: Agg[String],
                           jvmArgs: Seq[String]) = {
@@ -309,11 +328,11 @@ object Jvm {
 
     write(outputPath, launcherShellScript(mainClass, classPath.map(_.toString), jvmArgs))
 
-    val perms = java.nio.file.Files.getPosixFilePermissions(outputPath.toNIO)
+    val perms = Files.getPosixFilePermissions(outputPath.toNIO)
     perms.add(PosixFilePermission.GROUP_EXECUTE)
     perms.add(PosixFilePermission.OWNER_EXECUTE)
     perms.add(PosixFilePermission.OTHERS_EXECUTE)
-    java.nio.file.Files.setPosixFilePermissions(outputPath.toNIO, perms)
+    Files.setPosixFilePermissions(outputPath.toNIO, perms)
     PathRef(outputPath)
   }
 
diff --git a/main/src/mill/modules/Util.scala b/main/src/mill/modules/Util.scala
index cef11859..3029411c 100644
--- a/main/src/mill/modules/Util.scala
+++ b/main/src/mill/modules/Util.scala
@@ -3,7 +3,7 @@ package mill.modules
 
 import ammonite.ops.{Path, RelPath, empty, mkdir, read}
 import mill.eval.PathRef
-import mill.util.Ctx
+import mill.util.{Ctx, IO}
 
 object Util {
   def download(url: String, dest: RelPath = "download")(implicit ctx: Ctx.Dest) = {
@@ -45,15 +45,7 @@ object Util {
             val entryDest = ctx.dest / dest / RelPath(entry.getName)
             mkdir(entryDest / ammonite.ops.up)
             val fileOut = new java.io.FileOutputStream(entryDest.toString)
-            val buffer = new Array[Byte](4096)
-            while ( {
-              zipStream.read(buffer) match {
-                case -1 => false
-                case n =>
-                  fileOut.write(buffer, 0, n)
-                  true
-              }
-            }) ()
+            IO.stream(zipStream, fileOut)
             fileOut.close()
           }
           zipStream.closeEntry()
diff --git a/main/test/src/mill/util/TestEvaluator.scala b/main/test/src/mill/util/TestEvaluator.scala
index 078254f1..0dd435eb 100644
--- a/main/test/src/mill/util/TestEvaluator.scala
+++ b/main/test/src/mill/util/TestEvaluator.scala
@@ -25,11 +25,11 @@ class TestEvaluator[T <: TestUtil.BaseModule](module: T)
                                              tp: TestPath){
   val outPath =  TestUtil.getOutPath()
 
-  val logger = DummyLogger
-//  val logger = new PrintLogger(
-//    true,
-//    ammonite.util.Colors.Default, System.out, System.out, System.err, System.in
-// )
+//  val logger = DummyLogger
+  val logger = new PrintLogger(
+    true,
+    ammonite.util.Colors.Default, System.out, System.out, System.err, System.in
+ )
   val evaluator = new Evaluator(outPath, TestEvaluator.externalOutPath, module, logger)
 
   def apply[T](t: Task[T]): Either[Result.Failing[T], (T, Int)] = {
diff --git a/scalalib/src/mill/scalalib/ScalaModule.scala b/scalalib/src/mill/scalalib/ScalaModule.scala
index c7dbc322..5a355bc6 100644
--- a/scalalib/src/mill/scalalib/ScalaModule.scala
+++ b/scalalib/src/mill/scalalib/ScalaModule.scala
@@ -2,16 +2,15 @@ package mill
 package scalalib
 
 import ammonite.ops._
-import coursier.{Cache, MavenRepository, Repository}
-import mill.define.{Cross, Task}
+import coursier.Repository
+import mill.define.Task
 import mill.define.TaskModule
 import mill.eval.{PathRef, Result}
 import mill.modules.Jvm
-import mill.modules.Jvm.{createAssembly, createJar, interactiveSubprocess, runLocal, subprocess}
+import mill.modules.Jvm.{createAssembly, createJar, subprocess}
 import Lib._
-import mill.define.Cross.Resolver
 import mill.util.Loose.Agg
-import mill.util.{DummyInputStream, Strict}
+import mill.util.DummyInputStream
 
 /**
   * Core configuration required to compile a single Scala compilation target
@@ -144,6 +143,7 @@ trait ScalaModule extends mill.Module with TaskModule { outer =>
       if path.isFile && (path.ext == "scala" || path.ext == "java")
     } yield PathRef(path)
   }
+
   def compile: T[CompilationResult] = T.persistent{
     scalaWorker.worker().compileScala(
       scalaVersion(),
@@ -165,19 +165,35 @@ trait ScalaModule extends mill.Module with TaskModule { outer =>
     resolveDeps(T.task{compileIvyDeps() ++ scalaLibraryIvyDeps() ++ transitiveIvyDeps()})()
   }
 
-  def runClasspath = T{
+  def upstreamAssemblyClasspath = T{
     upstreamRunClasspath() ++
-    Agg(compile().classes) ++
-    resources() ++
     unmanagedClasspath() ++
     resolveDeps(T.task{runIvyDeps() ++ scalaLibraryIvyDeps() ++ transitiveIvyDeps()})()
   }
 
+  def runClasspath = T{
+    Agg(compile().classes) ++
+    resources() ++
+    upstreamAssemblyClasspath()
+
+  }
+
+  /**
+    * Build the assembly for upstream dependencies separate from the current classpath
+    *
+    * This should allow much faster assembly creation in the common case where
+    * upstream dependencies do not change
+    */
+  def upstreamAssembly = T{
+    createAssembly(upstreamAssemblyClasspath().map(_.path), mainClass())
+  }
+
   def assembly = T{
     createAssembly(
-      runClasspath().map(_.path).filter(exists),
+      Agg.from(resources().map(_.path)) ++ Agg(compile().classes.path),
       mainClass(),
-      prependShellScript = prependShellScript()
+      prependShellScript(),
+      Some(upstreamAssembly().path)
     )
   }
 
diff --git a/scalalib/test/src/mill/scalalib/HelloWorldTests.scala b/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
index 4fb5f5a5..317f9bec 100644
--- a/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
+++ b/scalalib/test/src/mill/scalalib/HelloWorldTests.scala
@@ -50,19 +50,6 @@ object HelloWorldTests extends TestSuite {
     }
   }
 
-  object HelloWorldWithMainAssembly extends HelloBase {
-    object core extends HelloWorldModule{
-      def mainClass = Some("Main")
-      def assembly = T{
-        mill.modules.Jvm.createAssembly(
-          runClasspath().map(_.path).filter(exists),
-          prependShellScript = prependShellScript(),
-          mainClass = mainClass()
-        )
-      }
-    }
-  }
-
   object HelloWorldWarnUnused extends HelloBase{
     object core extends HelloWorldModule {
       def scalacOptions = T(Seq("-Ywarn-unused"))
@@ -392,8 +379,8 @@ object HelloWorldTests extends TestSuite {
     }
 
     'assembly - {
-      'assembly - workspaceTest(HelloWorldWithMainAssembly){ eval =>
-        val Right((result, evalCount)) = eval.apply(HelloWorldWithMainAssembly.core.assembly)
+      'assembly - workspaceTest(HelloWorldWithMain){ eval =>
+        val Right((result, evalCount)) = eval.apply(HelloWorldWithMain.core.assembly)
         assert(
           exists(result.path),
           evalCount > 0
@@ -401,14 +388,15 @@ object HelloWorldTests extends TestSuite {
         val jarFile = new JarFile(result.path.toIO)
         val entries = jarFile.entries().asScala.map(_.getName).toSet
 
-        assert(entries.contains("Main.class"))
+        val mainPresent = entries.contains("Main.class")
+        assert(mainPresent)
         assert(entries.exists(s => s.contains("scala/Predef.class")))
 
         val mainClass = jarMainClass(jarFile)
         assert(mainClass.contains("Main"))
       }
-      'run - workspaceTest(HelloWorldWithMainAssembly){eval =>
-        val Right((result, evalCount)) = eval.apply(HelloWorldWithMainAssembly.core.assembly)
+      'run - workspaceTest(HelloWorldWithMain){eval =>
+        val Right((result, evalCount)) = eval.apply(HelloWorldWithMain.core.assembly)
 
         assert(
           exists(result.path),