Add ScalaPB integration (#395)

* Add ScalaPB integration

* Update ci scripts with new scalapblib module

* Move ScalaPB integration to contrib module

2 files changed, 140 insertions, 0 deletions

diff --git a/contrib/scalapblib/src/mill/contrib/scalapblib/ScalaPBModule.scala b/contrib/scalapblib/src/mill/contrib/scalapblib/ScalaPBModule.scala
new file mode 100644
index 00000000..9aa5b833
--- /dev/null
+++ b/contrib/scalapblib/src/mill/contrib/scalapblib/ScalaPBModule.scala
@@ -0,0 +1,70 @@
+package mill
+package contrib.scalapblib
+
+import coursier.{Cache, MavenRepository}
+import coursier.core.Version
+import mill.define.Sources
+import mill.eval.PathRef
+import mill.scalalib.Lib.resolveDependencies
+import mill.scalalib._
+import mill.util.Loose
+
+trait ScalaPBModule extends ScalaModule {
+
+  override def generatedSources = T { super.generatedSources() :+ compileScalaPB() }
+
+  override def ivyDeps = T {
+    super.ivyDeps() ++
+      Agg(ivy"com.thesamet.scalapb::scalapb-runtime:${scalaPBVersion()}") ++
+      (if (!scalaPBGrpc()) Agg() else Agg(ivy"com.thesamet.scalapb::scalapb-runtime-grpc:${scalaPBVersion()}"))
+  }
+
+  def scalaPBVersion: T[String]
+
+  def scalaPBFlatPackage: T[Boolean] = T { false }
+
+  def scalaPBJavaConversions: T[Boolean] = T { false }
+
+  def scalaPBGrpc: T[Boolean] = T { true }
+
+  def scalaPBSingleLineToProtoString: T[Boolean] = T { false }
+
+  def scalaPBSources: Sources = T.sources {
+    millSourcePath / 'protobuf
+  }
+
+  def scalaPBOptions: T[String] = T {
+    (
+      (if (scalaPBFlatPackage()) Seq("flat_package") else Seq.empty) ++
+      (if (scalaPBJavaConversions()) Seq("java_conversions") else Seq.empty) ++
+      (if (scalaPBGrpc()) Seq("grpc") else Seq.empty) ++ (
+        if (!scalaPBSingleLineToProtoString()) Seq.empty else {
+          if (Version(scalaPBVersion()) >= Version("0.7.0"))
+            Seq("single_line_to_proto_string")
+          else
+            Seq("single_line_to_string")
+        }
+      )
+    ).mkString(",")
+  }
+
+  def scalaPBClasspath: T[Loose.Agg[PathRef]] = T {
+    resolveDependencies(
+      Seq(
+        Cache.ivy2Local,
+        MavenRepository("https://repo1.maven.org/maven2")
+      ),
+      Lib.depToDependency(_, "2.12.4"),
+      Seq(ivy"com.thesamet.scalapb::scalapbc:${scalaPBVersion()}")
+    )
+  }
+
+  def compileScalaPB: T[PathRef] = T.persistent {
+    ScalaPBWorkerApi.scalaPBWorker
+      .compile(
+        scalaPBClasspath().map(_.path),
+        scalaPBSources().map(_.path),
+        scalaPBOptions(),
+        T.ctx().dest)
+  }
+}
diff --git a/contrib/scalapblib/src/mill/contrib/scalapblib/ScalaPBWorker.scala b/contrib/scalapblib/src/mill/contrib/scalapblib/ScalaPBWorker.scala
new file mode 100644
index 00000000..ea11a624
--- /dev/null
+++ b/contrib/scalapblib/src/mill/contrib/scalapblib/ScalaPBWorker.scala
@@ -0,0 +1,70 @@
+package mill
+package contrib.scalapblib
+
+import java.io.File
+import java.lang.reflect.Method
+import java.net.URLClassLoader
+
+import ammonite.ops.{Path, ls}
+import mill.eval.PathRef
+
+class ScalaPBWorker {
+
+  private var scalaPBInstanceCache = Option.empty[(Long, ScalaPBWorkerApi)]
+
+  private def scalaPB(scalaPBClasspath: Agg[Path]) = {
+    val classloaderSig = scalaPBClasspath.map(p => p.toString().hashCode + p.mtime.toMillis).sum
+    scalaPBInstanceCache match {
+      case Some((sig, instance)) if sig == classloaderSig => instance
+      case _ =>
+        val cl = new URLClassLoader(scalaPBClasspath.map(_.toIO.toURI.toURL).toArray)
+        val scalaPBCompilerClass = cl.loadClass("scalapb.ScalaPBC")
+        val mainMethod = scalaPBCompilerClass.getMethod("main", classOf[Array[java.lang.String]])
+
+        val instance = new ScalaPBWorkerApi {
+          override def compileScalaPB(source: File, scalaPBOptions: String, generatedDirectory: File) {
+            val opts = if (scalaPBOptions.isEmpty) "" else scalaPBOptions + ":"
+            mainMethod.invoke(
+              null,
+              Array(
+                "--throw",
+                s"--scala_out=${opts}${generatedDirectory.getCanonicalPath}",
+                s"--proto_path=${source.getParentFile.getCanonicalPath}",
+                source.getCanonicalPath
+              )
+            )
+          }
+        }
+        scalaPBInstanceCache = Some((classloaderSig, instance))
+        instance
+    }
+  }
+
+  def compile(scalaPBClasspath: Agg[Path], scalaPBSources: Seq[Path], scalaPBOptions: String, dest: Path)
+             (implicit ctx: mill.util.Ctx): mill.eval.Result[PathRef] = {
+    val compiler = scalaPB(scalaPBClasspath)
+
+    def compileScalaPBDir(inputDir: Path) {
+      // ls throws if the path doesn't exist
+      if (inputDir.toIO.exists) {
+        ls.rec(inputDir).filter(_.name.matches(".*.proto"))
+          .foreach { proto =>
+            compiler.compileScalaPB(proto.toIO, scalaPBOptions, dest.toIO)
+          }
+      }
+    }
+
+    scalaPBSources.foreach(compileScalaPBDir)
+
+    mill.eval.Result.Success(PathRef(dest))
+  }
+}
+
+trait ScalaPBWorkerApi {
+  def compileScalaPB(source: File, scalaPBOptions: String, generatedDirectory: File)
+}
+
+object ScalaPBWorkerApi {
+
+  def scalaPBWorker = new ScalaPBWorker()
+}