Ability to relocate/shade in assembly

com-lihaoyi · Aug 19, 2020 · 55514ce · 55514ce
1 parent b8f48ce
commit 55514ce
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git 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
@@ -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 {