From 0d46e00d71b3a68757529780e35da4593ce7583e Mon Sep 17 00:00:00 2001
From: Joan Goyeau <joan@goyeau.com>
Date: Tue, 18 Aug 2020 17:07:34 -0700
Subject: [PATCH] Ability to relocate/shade in assembly

---
 build.sc                        |  4 +++-
 main/src/modules/Assembly.scala | 23 +++++++++++++++++------
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/build.sc b/build.sc
index 3cbac66b5bf..46133b26a22 100755
--- a/build.sc
+++ b/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 jarjarabrams = ivy"com.eed3si9n.jarjarabrams::jarjar-abrams-core:0.2.0"
 }
 
 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.jarjarabrams
     )
 
     def generatedSources = T {
diff --git a/main/src/modules/Assembly.scala b/main/src/modules/Assembly.scala
index 2d104dfa8d5..30239257baf 100644
--- a/main/src/modules/Assembly.scala
+++ b/main/src/modules/Assembly.scala
@@ -1,5 +1,6 @@
 package mill.modules
 
+import com.eed3si9n.jarjarabrams.{ShadePattern, Shader}
 import java.io.InputStream
 import java.util.jar.JarFile
 import java.util.regex.Pattern
@@ -29,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))
     }
@@ -49,7 +52,7 @@ object Assembly {
       case Rule.ExcludePattern(pattern) => pattern.asPredicate().test(_)
     }
 
-    classpathIterator(inputPaths).foldLeft(Map.empty[String, GroupedEntry]) {
+    classpathIterator(inputPaths, assemblyRules).foldLeft(Map.empty[String, GroupedEntry]) {
       case (entries, (mapping, entry)) =>
         rulesMap.get(mapping) match {
           case Some(_: Assembly.Rule.Exclude) =>
@@ -70,25 +73,33 @@ object Assembly {
     }
   }
 
-  private def classpathIterator(inputPaths: Agg[os.Path]): Agg[(String, InputStream)] =
+  private def classpathIterator(inputPaths: Agg[os.Path], assemblyRules: Seq[Assembly.Rule]): Agg[(String, InputStream)] = {
+    val shadeRules = assemblyRules.collect {
+      case Rule.Relocate(from, to) => ShadePattern.Rename(List(from -> to)).inAll
+    }
+
     inputPaths
       .filter(os.exists)
       .flatMap { path =>
         if (os.isFile(path)) {
           val jarFile = new JarFile(path.toIO)
-          jarFile
+          val mappings = jarFile
             .entries()
             .asScala
             .filterNot(_.isDirectory)
-            .map(entry => entry.getName -> jarFile.getInputStream(entry))
+            .map(entry => jarFile.getInputStream(entry) -> entry.getName)
+          Shader.shadeInputStreams(shadeRules, mappings.toSeq, verbose = false)
         }
         else {
-          os
+          val pathsWithMappings = os
             .walk(path)
             .filter(os.isFile)
-            .map(subPath => subPath.relativeTo(path).toString -> os.read.inputStream(subPath))
+              .map(subPath => os.read.inputStream(subPath) -> subPath.relativeTo(path).toString)
+          Shader.shadeInputStreams(shadeRules, pathsWithMappings, verbose = false)
         }
       }
+      .map { case (inputStream, mapping) => mapping -> inputStream }
+  }
 }
 
 private[modules] sealed trait GroupedEntry {