Ability to relocate/shade in assembly

build.sc

@@ -74,6 +74,7 @@ object Deps {
   val utest = ivy"com.lihaoyi::utest:0.7.4"
   val zinc = ivy"org.scala-sbt::zinc:1.4.0-M1"
   val bsp = ivy"ch.epfl.scala:bsp4j:2.0.0-M4"
+  val jarRelocator = ivy"me.lucko:jar-relocator:1.4"
 }
 
 trait MillPublishModule extends PublishModule{
@@ -167,7 +168,8 @@ object main extends MillModule {
       // Necessary so we can share the JNA classes throughout the build process
       Deps.jna,
       Deps.jnaPlatform,
-      Deps.coursier
+      Deps.coursier,
+      Deps.jarRelocator
     )
 
     def generatedSources = T {

main/src/modules/Assembly.scala

@@ -3,10 +3,8 @@ 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 {
@@ -32,6 +30,8 @@ object Assembly {
 
     case class Exclude(path: String) extends Rule
 
+    case class Relocate(from: String, to: String) extends Rule
+
     object ExcludePattern {
       def apply(pattern: String): ExcludePattern = ExcludePattern(Pattern.compile(pattern))
     }
@@ -52,77 +52,63 @@ object Assembly {
       case Rule.ExcludePattern(pattern) => pattern.asPredicate().test(_)
     }
 
-    classpathIterator(inputPaths).foldLeft(Map.empty[String, GroupedEntry]) {
-      case (entries, entry) =>
-        val mapping = entry.mapping
-
-        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)
-            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 _ =>
-            entries
-        }
-    }
+    classpathIterator(inputPaths)
+      .foldLeft(Map.empty[String, GroupedEntry]) {
+        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(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 _ =>
+              entries
+          }
+      }
   }
 
-  private def classpathIterator(inputPaths: Agg[os.Path]): Generator[AssemblyEntry] = {
-    Generator.from(inputPaths)
-      .filter(os.exists)
-      .flatMap {
-        p =>
-          if (os.isFile(p)) {
-            val jf = new JarFile(p.toIO)
+  private def classpathIterator(inputPaths: Agg[os.Path]): Generator[(String, InputStream)] =
+      Generator.from(inputPaths)
+        .filter(os.exists)
+        .flatMap { path =>
+          if (os.isFile(path)) {
+            val jarFile = new JarFile(path.toIO)
             Generator.from(
-              for(entry <- jf.entries().asScala if !entry.isDirectory)
-                yield JarFileEntry(entry.getName, () => jf.getInputStream(entry))
+              jarFile
+                .entries()
+                .asScala
+                .filter(!_.isDirectory)
+                .map(entry => entry.getName -> jarFile.getInputStream(entry))
             )
           }
           else {
-            os.walk.stream(p)
+            os.walk
+              .stream(path)
               .filter(os.isFile)
-              .map(sub => PathEntry(sub.relativeTo(p).toString, sub))
+              .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,16 +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.util.{Gather, Task}
 import geny.Generator
+import me.lucko.jarrelocator.{JarRelocator, Relocation}
 import mill.main.client.InputPumper
 import mill.eval.{PathRef, Result}
 import mill.util.Ctx
 import mill.api.IO
 import mill.api.Loose.Agg
-
+import mill.modules.Assembly.Rule
 import scala.collection.mutable
 import scala.collection.JavaConverters._
 import upickle.default.{macroRW, ReadWriter => RW}
@@ -299,27 +299,27 @@ 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 relocations = assemblyRules.collect { case Rule.Relocate(from, to) => new Relocation(from, to) }
+    val relocated = ctx.dest / "out-relocated.jar"
+    new JarRelocator(tmp.toIO, relocated.toIO, relocations.asJava).run()
+
+    val output = ctx.dest / "out.jar"
     // Prepend shell script and make it executable
-    if (prependShellScript.isEmpty) os.move(tmp, output)
+    if (prependShellScript.isEmpty) os.move(relocated, output)
     else{
       val lineSep = if (!prependShellScript.endsWith("\n")) "\n\r\n" else ""
       os.write(output, prependShellScript + lineSep)
-      os.write.append(output, os.read.inputStream(tmp))
+      os.write.append(output, os.read.inputStream(relocated))
 
       if (!scala.util.Properties.isWin) {
         os.perms.set(