Refactor assembly

main/src/modules/Assembly.scala

@@ -3,10 +3,7 @@ package mill.modules
 import java.io.InputStream
 import java.util.jar.JarFile
 import java.util.regex.Pattern
-
-import geny.Generator
 import mill.Agg
-
 import scala.collection.JavaConverters._
 
 object Assembly {
@@ -53,22 +50,18 @@ object Assembly {
     }
 
     classpathIterator(inputPaths).foldLeft(Map.empty[String, GroupedEntry]) {
-      case (entries, entry) =>
-        val mapping = entry.mapping
-
+      case (entries, (mapping, entry)) =>
         rulesMap.get(mapping) match {
           case Some(_: Assembly.Rule.Exclude) =>
             entries
           case Some(a: Assembly.Rule.Append) =>
-            val newEntry = entries.getOrElse(mapping, AppendEntry(Nil, a.separator)).append(entry)
+            val newEntry = entries.getOrElse(mapping, AppendEntry(Seq.empty, a.separator)).append(entry)
             entries + (mapping -> newEntry)
-
           case _ if excludePatterns.exists(_(mapping)) =>
             entries
           case _ if appendPatterns.exists(_(mapping)) =>
             val newEntry = entries.getOrElse(mapping, AppendEntry.empty).append(entry)
             entries + (mapping -> newEntry)
-
           case _ if !entries.contains(mapping) =>
             entries + (mapping -> WriteOnceEntry(entry))
           case _ =>
@@ -77,52 +70,39 @@ object Assembly {
     }
   }
 
-  private def classpathIterator(inputPaths: Agg[os.Path]): Generator[AssemblyEntry] = {
-    Generator.from(inputPaths)
+  private def classpathIterator(inputPaths: Agg[os.Path]): Agg[(String, InputStream)] =
+    inputPaths
       .filter(os.exists)
-      .flatMap {
-        p =>
-          if (os.isFile(p)) {
-            val jf = new JarFile(p.toIO)
-            Generator.from(
-              for(entry <- jf.entries().asScala if !entry.isDirectory)
-                yield JarFileEntry(entry.getName, () => jf.getInputStream(entry))
-            )
-          }
-          else {
-            os.walk.stream(p)
-              .filter(os.isFile)
-              .map(sub => PathEntry(sub.relativeTo(p).toString, sub))
-          }
+      .flatMap { path =>
+        if (os.isFile(path)) {
+          val jarFile = new JarFile(path.toIO)
+          jarFile
+            .entries()
+            .asScala
+            .filterNot(_.isDirectory)
+            .map(entry => entry.getName -> jarFile.getInputStream(entry))
+        }
+        else {
+          os
+            .walk(path)
+            .filter(os.isFile)
+            .map(subPath => subPath.relativeTo(path).toString -> os.read.inputStream(subPath))
+        }
       }
-  }
 }
 
 private[modules] sealed trait GroupedEntry {
-  def append(entry: AssemblyEntry): GroupedEntry
+  def append(entry: InputStream): GroupedEntry
 }
 
 private[modules] object AppendEntry {
   val empty: AppendEntry = AppendEntry(Nil, Assembly.defaultSeparator)
 }
 
-private[modules] case class AppendEntry(entries: List[AssemblyEntry], separator: String) extends GroupedEntry {
-  def append(entry: AssemblyEntry): GroupedEntry = copy(entries = entry :: this.entries)
-}
-
-private[modules] case class WriteOnceEntry(entry: AssemblyEntry) extends GroupedEntry {
-  def append(entry: AssemblyEntry): GroupedEntry = this
-}
-
-private[this] sealed trait AssemblyEntry {
-  def mapping: String
-  def inputStream: InputStream
-}
-
-private[this] case class PathEntry(mapping: String, path: os.Path) extends AssemblyEntry {
-  def inputStream: InputStream = os.read.inputStream(path)
+private[modules] case class AppendEntry(entries: Seq[InputStream], separator: String) extends GroupedEntry {
+  def append(entry: InputStream): GroupedEntry = copy(entries = entry +: entries)
 }
 
-private[this] case class JarFileEntry(mapping: String, getIs: () => InputStream) extends AssemblyEntry {
-  def inputStream: InputStream = getIs()
+private[modules] case class WriteOnceEntry(entry: InputStream) extends GroupedEntry {
+  def append(entry: InputStream): GroupedEntry = this
 }

main/src/modules/Jvm.scala

@@ -7,19 +7,16 @@ import java.nio.file.{FileSystems, Files, StandardOpenOption}
 import java.nio.file.attribute.PosixFilePermission
 import java.util.Collections
 import java.util.jar.{Attributes, JarEntry, JarFile, JarOutputStream, Manifest}
-
-import coursier.{Dependency, Fetch, Repository, Resolution}
+import coursier.{Dependency, Repository, Resolution}
 import coursier.util.{Gather, Task}
-import geny.Generator
 import mill.main.client.InputPumper
 import mill.eval.{PathRef, Result}
 import mill.util.Ctx
 import mill.api.IO
 import mill.api.Loose.Agg
-
 import scala.collection.mutable
 import scala.collection.JavaConverters._
-import upickle.default.{macroRW, ReadWriter => RW}
+import upickle.default.{ReadWriter => RW}
 
 object Jvm {
   /**
@@ -299,21 +296,17 @@ object Jvm {
           val path = zipFs.getPath(mapping).toAbsolutePath
           val separated =
             if (entries.isEmpty) Nil
-            else
-              entries.head +: entries.tail.flatMap { e =>
-                List(JarFileEntry(e.mapping, () => new ByteArrayInputStream(separator.getBytes)), e)
-              }
-          val concatenated = new SequenceInputStream(
-            Collections.enumeration(separated.map(_.inputStream).asJava))
+            else entries.head +: entries.tail.flatMap(e => Seq(new ByteArrayInputStream(separator.getBytes), e))
+          val concatenated = new SequenceInputStream(Collections.enumeration(separated.asJava))
           writeEntry(path, concatenated, append = true)
         case (mapping, WriteOnceEntry(entry)) =>
           val path = zipFs.getPath(mapping).toAbsolutePath
-          writeEntry(path, entry.inputStream, append = false)
-      }
 
+          writeEntry(path, entry, append = false)
+      }
     zipFs.close()
-    val output = ctx.dest / "out.jar"
 
+    val output = ctx.dest / "out.jar"
     // Prepend shell script and make it executable
     if (prependShellScript.isEmpty) os.move(tmp, output)
     else{