Add ScalaPB integration (#395)

* Add ScalaPB integration

* Update ci scripts with new scalapblib module

* Move ScalaPB integration to contrib module

10 files changed, 331 insertions, 5 deletions

diff --git a/build.sc b/build.sc
index 4b142cc9..27914803 100755
--- a/build.sc
+++ b/build.sc
@@ -232,6 +232,15 @@ object twirllib extends MillModule {
 
 }
 
+object contrib extends MillModule {
+
+  object scalapblib extends MillModule {
+    def moduleDeps = Seq(scalalib)
+  }
+
+}
+
+
 object scalanativelib extends MillModule {
   def moduleDeps = Seq(scalalib)
 
@@ -341,7 +350,7 @@ def launcherScript(shellJvmArgs: Seq[String],
 }
 
 object dev extends MillModule{
-  def moduleDeps = Seq(scalalib, scalajslib, scalanativelib)
+  def moduleDeps = Seq(scalalib, scalajslib, scalanativelib, contrib.scalapblib)
   def forkArgs =
     (
       scalalib.testArgs() ++
diff --git a/ci/test-mill-0.sh b/ci/test-mill-0.sh
index e63760cf..8e44b912 100755
--- a/ci/test-mill-0.sh
+++ b/ci/test-mill-0.sh
@@ -6,4 +6,4 @@ set -eux
 git clean -xdf
 
 # Run tests
-mill -i all {main,scalalib,scalajslib,twirllib,main.client}.test
+mill -i all {main,scalalib,scalajslib,twirllib,main.client,contrib.scalapblib}.test
diff --git a/ci/test-mill-bootstrap.sh b/ci/test-mill-bootstrap.sh
index 8010e700..cd959f9e 100755
--- a/ci/test-mill-bootstrap.sh
+++ b/ci/test-mill-bootstrap.sh
@@ -27,4 +27,4 @@ git clean -xdf
 rm -rf ~/.mill
 
 # Use second build to run tests using Mill
-~/mill-2 -i all {main,scalalib,scalajslib}.test
\ No newline at end of file
+~/mill-2 -i all {main,scalalib,scalajslib,twirllib,contrib.scalapblib}.test
diff --git a/ci/test-mill-dev.sh b/ci/test-mill-dev.sh
index f5a8bfcd..deb48dca 100755
--- a/ci/test-mill-dev.sh
+++ b/ci/test-mill-dev.sh
@@ -11,5 +11,5 @@ mill -i dev.assembly
 rm -rf ~/.mill
 
 # Second build & run tests
-out/dev/assembly/dest/mill -i all {main,scalalib,scalajslib,twirllib}.test
+out/dev/assembly/dest/mill -i all {main,scalalib,scalajslib,twirllib,contrib.scalapblib}.test
 
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()
+}
diff --git a/contrib/scalapblib/test/protobuf/tutorial/Tutorial.proto b/contrib/scalapblib/test/protobuf/tutorial/Tutorial.proto
new file mode 100644
index 00000000..d66911b8
--- /dev/null
+++ b/contrib/scalapblib/test/protobuf/tutorial/Tutorial.proto
@@ -0,0 +1,29 @@
+syntax = "proto2";
+
+package tutorial;
+
+option java_package = "com.example.tutorial";
+option java_outer_classname = "AddressBookProtos";
+
+message Person {
+  required string name = 1;
+  required int32 id = 2;
+  optional string email = 3;
+
+  enum PhoneType {
+    MOBILE = 0;
+    HOME = 1;
+    WORK = 2;
+  }
+
+  message PhoneNumber {
+    required string number = 1;
+    optional PhoneType type = 2 [default = HOME];
+  }
+
+  repeated PhoneNumber phones = 4;
+}
+
+message AddressBook {
+  repeated Person people = 1;
+}
diff --git a/contrib/scalapblib/test/src/mill/contrib/scalapblib/TutorialTests.scala b/contrib/scalapblib/test/src/mill/contrib/scalapblib/TutorialTests.scala
new file mode 100644
index 00000000..f88d3a5f
--- /dev/null
+++ b/contrib/scalapblib/test/src/mill/contrib/scalapblib/TutorialTests.scala
@@ -0,0 +1,114 @@
+package mill.contrib.scalapblib
+
+import ammonite.ops.{Path, cp, ls, mkdir, pwd, rm, _}
+import mill.eval.Result
+import mill.util.{TestEvaluator, TestUtil}
+import utest.framework.TestPath
+import utest.{TestSuite, Tests, assert, _}
+
+object TutorialTests extends TestSuite {
+
+  trait TutorialBase extends TestUtil.BaseModule {
+    override def millSourcePath: Path = TestUtil.getSrcPathBase() / millOuterCtx.enclosing.split('.')
+  }
+
+  trait TutorialModule extends ScalaPBModule {
+    def scalaVersion = "2.12.4"
+    def scalaPBVersion = "0.7.4"
+    def scalaPBFlatPackage = true
+  }
+
+  object Tutorial extends TutorialBase {
+
+    object core extends TutorialModule {
+      override def scalaPBVersion = "0.7.4"
+    }
+  }
+
+  val resourcePath: Path = pwd / 'contrib / 'scalapblib / 'test / 'protobuf / 'tutorial
+
+  def protobufOutPath[M <: TestUtil.BaseModule](eval: TestEvaluator[M]): Path =
+    eval.outPath / 'core / 'compileScalaPB / 'dest / 'com / 'example / 'tutorial
+
+  def workspaceTest[T, M <: TestUtil.BaseModule](m: M)
+                                                (t: TestEvaluator[M] => T)
+                                                (implicit tp: TestPath): T = {
+    val eval = new TestEvaluator(m)
+    rm(m.millSourcePath)
+    println(m.millSourcePath)
+    rm(eval.outPath)
+    println(eval.outPath)
+    mkdir(m.millSourcePath / 'core / 'protobuf)
+    cp(resourcePath, m.millSourcePath / 'core / 'protobuf / 'tutorial)
+    t(eval)
+  }
+
+  def compiledSourcefiles: Seq[RelPath] = Seq[RelPath](
+    "AddressBook.scala",
+    "Person.scala",
+    "TutorialProto.scala"
+  )
+
+  def tests: Tests = Tests {
+    'scalapbVersion - {
+
+      'fromBuild - workspaceTest(Tutorial) { eval =>
+        val Right((result, evalCount)) = eval.apply(Tutorial.core.scalaPBVersion)
+
+        assert(
+          result == "0.7.4",
+          evalCount > 0
+        )
+      }
+    }
+
+    'compileScalaPB - {
+      'calledDirectly - workspaceTest(Tutorial) { eval =>
+        val Right((result, evalCount)) = eval.apply(Tutorial.core.compileScalaPB)
+
+        val outPath = protobufOutPath(eval)
+
+        val outputFiles = ls.rec(result.path).filter(_.isFile)
+
+        val expectedSourcefiles = compiledSourcefiles.map(outPath / _)
+
+        assert(
+          result.path == eval.outPath / 'core / 'compileScalaPB / 'dest,
+          outputFiles.nonEmpty,
+          outputFiles.forall(expectedSourcefiles.contains),
+          outputFiles.size == 3,
+          evalCount > 0
+        )
+
+        // don't recompile if nothing changed
+        val Right((_, unchangedEvalCount)) = eval.apply(Tutorial.core.compileScalaPB)
+
+        assert(unchangedEvalCount == 0)
+      }
+
+      // This throws a NullPointerException in coursier somewhere
+      //
+      // 'triggeredByScalaCompile - workspaceTest(Tutorial) { eval =>
+      //   val Right((_, evalCount)) = eval.apply(Tutorial.core.compile)
+
+      //   val outPath = protobufOutPath(eval)
+
+      //   val outputFiles = ls.rec(outPath).filter(_.isFile)
+
+      //   val expectedSourcefiles = compiledSourcefiles.map(outPath / _)
+
+      //   assert(
+      //     outputFiles.nonEmpty,
+      //     outputFiles.forall(expectedSourcefiles.contains),
+      //     outputFiles.size == 3,
+      //     evalCount > 0
+      //   )
+
+      //   // don't recompile if nothing changed
+      //   val Right((_, unchangedEvalCount)) = eval.apply(Tutorial.core.compile)
+
+      //   assert(unchangedEvalCount == 0)
+      // }
+    }
+  }
+}
diff --git a/scratch/build.sc b/scratch/build.sc
index db446c99..46c271b5 100644
--- a/scratch/build.sc
+++ b/scratch/build.sc
@@ -12,4 +12,9 @@ object core extends JavaModule{
 object app extends JavaModule{
   def moduleDeps = Seq(core)
   object test extends Tests with JUnitTests
-}
\ No newline at end of file
+}
+
+object scalapb extends mill.contrib.scalapblib.ScalaPBModule {
+  def scalaVersion = "2.12.4"
+  def scalaPBVersion = "0.7.4"
+}
diff --git a/scratch/scalapb/protobuf/scratch.proto b/scratch/scalapb/protobuf/scratch.proto
new file mode 100644
index 00000000..3153ad54
--- /dev/null
+++ b/scratch/scalapb/protobuf/scratch.proto
@@ -0,0 +1,29 @@
+syntax = "proto2";
+
+package tutorial;
+
+option java_package = "com.example.tutorial";
+option java_outer_classname = "AddressBookProtos";
+
+message Person {
+  required string name = 1;
+  required int32 id = 2;
+  optional string email = 3;
+
+  enum PhoneType {
+    MOBILE = 0;
+    HOME = 1;
+    WORK = 2;
+  }
+
+  message PhoneNumber {
+    required string number = 1;
+    optional PhoneType type = 2 [default = HOME];
+  }
+
+  repeated PhoneNumber phones = 4;
+}
+
+message AddressBook {
+  repeated Person people = 1;
+}
\ No newline at end of file