Refactor `mill.Main` to avoid going through Ammonite's main-method-dispatch system

7 files changed, 297 insertions, 184 deletions

diff --git a/core/src/main/scala/mill/Main.scala b/core/src/main/scala/mill/Main.scala
index 761430e4..4b4af35c 100644
--- a/core/src/main/scala/mill/Main.scala
+++ b/core/src/main/scala/mill/Main.scala
@@ -13,97 +13,10 @@ import ammonite.repl.Repl
 import ammonite.util.Util.normalizeNewlines
 import mill.define.Task.TaskModule
 object Main {
-  def parseSelector(input: String) = {
-    import fastparse.all._
-    val segment = P( CharsWhileIn(('a' to 'z') ++ ('A' to 'Z') ++ ('0' to '9')).! ).map(
-      Mirror.Segment.Label
-    )
-    val crossSegment = P( "[" ~ CharsWhile(c => c != ',' && c != ']').!.rep(1, sep=",") ~ "]" ).map(
-      Mirror.Segment.Cross
-    )
-    val query = P( segment ~ ("." ~ segment | crossSegment).rep ~ End ).map{
-      case (h, rest) => h :: rest.toList
-    }
-    query.parse(input)
-  }
-
-
-
-  def parseArgs(selectorString: String): Either[String, List[Mirror.Segment]] = {
-    import fastparse.all.Parsed
-    if (selectorString.isEmpty) Left("Selector cannot be empty")
-    else parseSelector(selectorString) match {
-      case f: Parsed.Failure => Left(s"Parsing exception ${f.msg}")
-      case Parsed.Success(selector, _) => Right(selector)
-    }
-  }
-
-
-  def consistencyCheck[T](mapping: Discovered.Mapping[T]): Either[String, Unit] = {
-    val consistencyErrors = Discovered.consistencyCheck(mapping)
-    if (consistencyErrors.nonEmpty) {
-      Left(s"Failed Discovered.consistencyCheck: ${consistencyErrors.map(Mirror.renderSelector)}")
-    } else {
-      Right(())
-    }
-  }
-
-  def evaluate(evaluator: Evaluator,
-               target: Task[Any],
-               watch: Path => Unit): Option[String] = {
-    val evaluated = evaluator.evaluate(OSet(target))
-    evaluated.transitive.foreach {
-      case t: define.Source => watch(t.handle.path)
-      case _ => // do nothing
-    }
-
-    val errorStr =
-      (for((k, fs) <- evaluated.failing.items()) yield {
-        val ks = k match{
-          case Left(t) => t.toString
-          case Right(t) => Mirror.renderSelector(t.segments.toList)
-        }
-        val fss = fs.map{
-          case Result.Exception(t) => t.toString
-          case Result.Failure(t) => t
-        }
-        s"$ks ${fss.mkString(", ")}"
-      }).mkString("\n")
 
-    evaluated.failing.keyCount match {
-      case 0 => None
-      case n => Some(s"$n targets failed\n$errorStr")
-    }
-  }
 
-  def apply[T](args: Seq[String],
-               mapping: Discovered.Mapping[T],
-               watch: Path => Unit,
-               coloredOutput: Boolean): Int = {
 
-    val log = new PrintLogger(coloredOutput)
 
-    val Seq(selectorString, rest @_*) = args
-
-    val res = for {
-      sel <- parseArgs(selectorString)
-      _ <- consistencyCheck(mapping)
-      crossSelectors = sel.map{
-        case Mirror.Segment.Cross(x) => x.toList.map(_.toString)
-        case _ => Nil
-      }
-      target <- mill.main.Resolve.resolve(sel, mapping.mirror, mapping.base, rest, crossSelectors, Nil)
-      evaluator = new Evaluator(pwd / 'out, mapping.value, log, sel)
-      _ <- evaluate(evaluator, target, watch).toLeft(())
-    } yield ()
-
-    res match {
-      case Left(err) =>
-        log.error(err)
-        1
-      case Right(_) => 0
-    }
-  }
 
   case class Config(home: ammonite.ops.Path = pwd/'out/'ammonite,
                     colored: Option[Boolean] = None,
@@ -141,16 +54,7 @@ object Main {
             welcomeBanner = None
           )
 
-        val runner = new ammonite.MainRunner(
-          config,
-          System.out, System.err,
-          System.in, System.out, System.err
-        ){
-          override def initMain(isRepl: Boolean) = {
-            super.initMain(isRepl).copy(scriptCodeWrapper = customCodeWrapper)
-          }
-        }
-
+        val runner = new mill.main.MainRunner(config)
         if (repl){
           runner.printInfo("Loading...")
           runner.runRepl()
@@ -159,82 +63,6 @@ object Main {
         }
     }
   }
-  val customCodeWrapper = new Preprocessor.CodeWrapper {
-    def top(pkgName: Seq[Name], imports: Imports, indexedWrapperName: Name) = {
-      s"""
-      |package ${pkgName.head.encoded}
-      |package ${Util.encodeScalaSourcePath(pkgName.tail)}
-      |$imports
-      |import mill._
-      |sealed abstract class ${indexedWrapperName.backticked} extends mill.Module{\n
-      |""".stripMargin
-    }
-
-    def bottom(printCode: String, indexedWrapperName: Name, extraCode: String) = {
-      val wrapName = indexedWrapperName.backticked
-      val tmpName = ammonite.util.Name(indexedWrapperName.raw + "-Temp").backticked
-
-      // Define `discovered` in the `tmpName` trait, before mixing in `MainWrapper`,
-      // to ensure that `$tempName#discovered` is initialized before `MainWrapper` is.
-      //
-      // `import $wrapName._` is necessary too let Ammonite pick up all the
-      // members of class wrapper, which are inherited but otherwise not visible
-      // in the AST of the `$wrapName` object
-      //
-      // We need to duplicate the Ammonite predef as part of the wrapper because
-      // imports within the body of the class wrapper are not brought into scope
-      // by the `import $wrapName._`. Other non-Ammonite-predef imports are not
-      // made available, and that's just too bad
-      s"""
-        |}
-        |trait $tmpName{
-        |  val discovered = mill.discover.Discovered.make[$wrapName]
-        |  val interpApi = ammonite.interp.InterpBridge.value
-        |}
-        |
-        |object $wrapName
-        |extends $wrapName
-        |with $tmpName
-        |with mill.MainWrapper[$wrapName] {
-        |  @ammonite.main.Router.main
-        |  def idea() = mill.scalaplugin.GenIdea(mapping)
-        |  ${ammonite.main.Defaults.replPredef}
-        |  ${ammonite.main.Defaults.predefString}
-        |  ${ammonite.Main.extraPredefString}
-        |  import ammonite.repl.ReplBridge.{value => repl}
-        |  import ammonite.interp.InterpBridge.{value => interp}
-        |  import $wrapName._
-      """.stripMargin +
-      Preprocessor.CodeWrapper.bottom(printCode, indexedWrapperName, extraCode)
-    }
-  }
 }
 
-/**
-  * Class that wraps each Mill build file.
-  */
-trait MainWrapper[T]{
-  val discovered: mill.discover.Discovered[T]
-  val interpApi: ammonite.interp.InterpAPI
-  val mapping = discovered.mapping(this.asInstanceOf[T])
-
-  mill.Main.consistencyCheck(mapping).left.foreach(msg => throw new Exception(msg))
-
-  @ammonite.main.Router.main
-  def run(args: String*) = mill.Main(
-    args,
-    mapping,
-    interpApi.watch,
-    true
-  )
-
-
-  val evaluator = new mill.eval.Evaluator(
-    ammonite.ops.pwd / 'out,
-    mapping.value,
-    new mill.util.PrintLogger(true)
-  )
 
-  implicit val replApplyHandler: mill.main.ReplApplyHandler =
-    new mill.main.ReplApplyHandler(evaluator)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/mill/eval/Evaluator.scala b/core/src/main/scala/mill/eval/Evaluator.scala
index ea680ada..4f0ace0c 100644
--- a/core/src/main/scala/mill/eval/Evaluator.scala
+++ b/core/src/main/scala/mill/eval/Evaluator.scala
@@ -16,8 +16,7 @@ import scala.collection.mutable
 class Evaluator(workspacePath: Path,
                 labeling: Map[Target[_], LabelledTarget[_]],
                 log: Logger,
-                sel: List[Mirror.Segment] = List(),
-                classLoaderSig: Seq[(Path, Long)] = Evaluator.classLoaderSig){
+                val classLoaderSig: Seq[(Path, Long)] = Evaluator.classLoaderSig){
   val workerCache = mutable.Map.empty[Ctx.Loader[_], Any]
 
   def evaluate(goals: OSet[Task[_]]): Evaluator.Results = {
@@ -197,15 +196,9 @@ class Evaluator(workspacePath: Path,
   }
 
   def resolveLogger(targetDestPath: Option[Path]): Logger = {
-    if (targetDestPath.isEmpty && sel.isEmpty)
-      log
+    if (targetDestPath.isEmpty) log
     else {
-      val path = targetDestPath.getOrElse(
-        sel.foldLeft[Path](pwd / 'out) {
-          case (d, Label(s)) => d / s
-          case (d, Cross(args)) => d / args.map(_.toString)
-        }
-      )
+      val path = targetDestPath.getOrElse(pwd/ 'out / 'command)
       val dir = path / up
       mkdir(dir)
       val file = dir / (path.last + ".log")
diff --git a/core/src/main/scala/mill/main/CustomCodeWrapper.scala b/core/src/main/scala/mill/main/CustomCodeWrapper.scala
new file mode 100644
index 00000000..3d840128
--- /dev/null
+++ b/core/src/main/scala/mill/main/CustomCodeWrapper.scala
@@ -0,0 +1,53 @@
+package mill.main
+
+import ammonite.interp.Preprocessor
+import ammonite.util.{Imports, Name, Util}
+
+object CustomCodeWrapper extends Preprocessor.CodeWrapper {
+  def top(pkgName: Seq[Name], imports: Imports, indexedWrapperName: Name) = {
+    s"""
+       |package ${pkgName.head.encoded}
+       |package ${Util.encodeScalaSourcePath(pkgName.tail)}
+       |$imports
+       |import mill._
+       |sealed abstract class ${indexedWrapperName.backticked} extends mill.Module{\n
+       |""".stripMargin
+  }
+
+
+  def bottom(printCode: String, indexedWrapperName: Name, extraCode: String) = {
+    val wrapName = indexedWrapperName.backticked
+    val tmpName = ammonite.util.Name(indexedWrapperName.raw + "-Temp").backticked
+
+    // Define `discovered` in the `tmpName` trait, before mixing in `MainWrapper`,
+    // to ensure that `$tempName#discovered` is initialized before `MainWrapper` is.
+    //
+    // `import $wrapName._` is necessary too let Ammonite pick up all the
+    // members of class wrapper, which are inherited but otherwise not visible
+    // in the AST of the `$wrapName` object
+    //
+    // We need to duplicate the Ammonite predef as part of the wrapper because
+    // imports within the body of the class wrapper are not brought into scope
+    // by the `import $wrapName._`. Other non-Ammonite-predef imports are not
+    // made available, and that's just too bad
+    s"""
+       |}
+       |trait $tmpName{
+       |  val discovered = mill.discover.Discovered.make[$wrapName]
+       |  val interpApi = ammonite.interp.InterpBridge.value
+       |}
+       |
+       |object $wrapName
+       |extends $wrapName
+       |with $tmpName
+       |with mill.main.MainWrapper[$wrapName] {
+       |  ${ammonite.main.Defaults.replPredef}
+       |  ${ammonite.main.Defaults.predefString}
+       |  ${ammonite.Main.extraPredefString}
+       |  import ammonite.repl.ReplBridge.{value => repl}
+       |  import ammonite.interp.InterpBridge.{value => interp}
+       |  import $wrapName._
+      """.stripMargin +
+      Preprocessor.CodeWrapper.bottom(printCode, indexedWrapperName, extraCode)
+  }
+}
diff --git a/core/src/main/scala/mill/main/MainRunner.scala b/core/src/main/scala/mill/main/MainRunner.scala
new file mode 100644
index 00000000..9239463f
--- /dev/null
+++ b/core/src/main/scala/mill/main/MainRunner.scala
@@ -0,0 +1,44 @@
+package mill.main
+import ammonite.ops.Path
+import ammonite.util.Res
+import mill.discover.Discovered
+import mill.eval.Evaluator
+
+class MainRunner(config: ammonite.main.Cli.Config)
+  extends ammonite.MainRunner(
+    config,
+    System.out, System.err, System.in, System.out, System.err
+  ){
+  var lastEvaluator: Option[(Seq[(Path, Long)], Discovered.Mapping[_], Evaluator)] = None
+  override def runScript(scriptPath: Path, scriptArgs: List[String]) =
+    watchLoop(
+      isRepl = false,
+      printing = true,
+      mainCfg => {
+        mainCfg.instantiateInterpreter() match{
+          case Left(problems) => problems
+          case Right(interp) =>
+            val result = RunScript.runScript(
+              mainCfg.wd, scriptPath, interp, scriptArgs, lastEvaluator
+            )
+
+            val interpWatched = interp.watchedFiles
+            result match{
+              case Res.Success((mapping, eval, evaluationWatches, success)) =>
+                lastEvaluator = Some((interpWatched, mapping, eval))
+                (result, interpWatched ++ evaluationWatches)
+              case _ =>
+                (result, interpWatched)
+            }
+
+
+        }
+      }
+    )
+  override def initMain(isRepl: Boolean) = {
+    super.initMain(isRepl).copy(scriptCodeWrapper = mill.main.CustomCodeWrapper)
+  }
+  override def handleWatchRes[T](res: Res[T], printing: Boolean) = {
+    super.handleWatchRes(res, printing = false)
+  }
+}
diff --git a/core/src/main/scala/mill/main/MainWrapper.scala b/core/src/main/scala/mill/main/MainWrapper.scala
new file mode 100644
index 00000000..813cd06b
--- /dev/null
+++ b/core/src/main/scala/mill/main/MainWrapper.scala
@@ -0,0 +1,19 @@
+package mill.main
+
+/**
+  * Class that wraps each Mill build file.
+  */
+trait MainWrapper[T]{
+  val discovered: mill.discover.Discovered[T]
+  val interpApi: ammonite.interp.InterpAPI
+  val mapping = discovered.mapping(this.asInstanceOf[T])
+
+  implicit val replApplyHandler: mill.main.ReplApplyHandler =
+    new mill.main.ReplApplyHandler(
+      new mill.eval.Evaluator(
+        ammonite.ops.pwd / 'out,
+        mapping.value,
+        new mill.util.PrintLogger(true)
+      )
+    )
+}
diff --git a/core/src/main/scala/mill/main/RunScript.scala b/core/src/main/scala/mill/main/RunScript.scala
new file mode 100644
index 00000000..14d521fe
--- /dev/null
+++ b/core/src/main/scala/mill/main/RunScript.scala
@@ -0,0 +1,176 @@
+package mill.main
+
+import java.nio.file.NoSuchFileException
+
+import ammonite.interp.Interpreter
+import ammonite.ops.{Path, pwd, read}
+import ammonite.util.Util.CodeSource
+import ammonite.util.{Name, Res, Util}
+import mill.define
+import mill.define.Task
+import mill.discover.{Discovered, Mirror}
+import mill.eval.{Evaluator, Result}
+import mill.util.{OSet, PrintLogger}
+
+import scala.collection.mutable
+
+/**
+  * Custom version of ammonite.main.Scripts, letting us run the build.sc script
+  * directly without going through Ammonite's main-method/argument-parsing
+  * subsystem
+  */
+object RunScript{
+
+  def runScript(wd: Path,
+                path: Path,
+                interp: ammonite.interp.Interpreter,
+                scriptArgs: Seq[String],
+                lastEvaluator: Option[(Seq[(Path, Long)], Discovered.Mapping[_], Evaluator)])
+  : Res[(Discovered.Mapping[_], Evaluator, Seq[(Path, Long)], Boolean)] = {
+
+    val log = new PrintLogger(true)
+    for{
+      (mapping, evaluator) <- lastEvaluator match{
+        case Some((prevInterpWatchedSig, prevMapping, prevEvaluator))
+          if watchedSigUnchanged(prevInterpWatchedSig) =>
+          Res.Success((prevMapping, prevEvaluator))
+
+        case _ =>
+          interp.watch(path)
+          for(mapping <- evaluateMapping(wd, path, interp))
+            yield (mapping, new Evaluator(pwd / 'out, mapping.value, log))
+      }
+    } yield {
+      val evaluationWatches = mutable.Buffer.empty[(Path, Long)]
+      val res = evaluateTarget(
+        evaluator,
+        mapping,
+        scriptArgs,
+        p => evaluationWatches.append((p, Interpreter.pathSignature(p)))
+      )
+      (mapping, evaluator, evaluationWatches, res.isRight)
+    }
+  }
+  def watchedSigUnchanged(sig: Seq[(Path, Long)]) = {
+    sig.forall{case (p, l) => Interpreter.pathSignature(p) == l}
+  }
+  def evaluateMapping(wd: Path,
+                      path: Path,
+                      interp: ammonite.interp.Interpreter): Res[Discovered.Mapping[_]] = {
+
+    val (pkg, wrapper) = Util.pathToPackageWrapper(Seq(), path relativeTo wd)
+
+    for {
+      scriptTxt <- try Res.Success(Util.normalizeNewlines(read(path))) catch {
+        case e: NoSuchFileException => Res.Failure("Script file not found: " + path)
+      }
+      processed <- interp.processModule(
+        scriptTxt,
+        CodeSource(wrapper, pkg, Seq(Name("ammonite"), Name("$file")), Some(path)),
+        autoImport = true,
+        extraCode = "",
+        hardcoded = true
+      )
+
+      buildClsName <- processed.blockInfo.lastOption match {
+        case Some(meta) => Res.Success(meta.id.wrapperPath)
+        case None => Res.Skip
+      }
+
+      buildCls = interp
+        .evalClassloader
+        .loadClass(buildClsName)
+
+      mapping <- try {
+        Util.withContextClassloader(interp.evalClassloader) {
+          Res.Success(
+            buildCls.getDeclaredMethod("mapping")
+              .invoke(null)
+              .asInstanceOf[Discovered.Mapping[_]]
+          )
+        }
+      } catch {
+        case e: Throwable => Res.Exception(e, "")
+      }
+      _ <- Res(consistencyCheck(mapping))
+    } yield mapping
+  }
+  def evaluateTarget[T](evaluator: Evaluator,
+                        mapping: Discovered.Mapping[T],
+                        scriptArgs: Seq[String],
+                        watch: Path => Unit) = {
+
+    val Seq(selectorString, rest @_*) = scriptArgs
+    for {
+      sel <- parseArgs(selectorString)
+      crossSelectors = sel.map{
+        case Mirror.Segment.Cross(x) => x.toList.map(_.toString)
+        case _ => Nil
+      }
+      target <- mill.main.Resolve.resolve(sel, mapping.mirror, mapping.base, rest, crossSelectors, Nil)
+      _ <- evaluate(evaluator, target, watch).toLeft(())
+    } yield ()
+  }
+  def evaluate(evaluator: Evaluator,
+               target: Task[Any],
+               watch: Path => Unit): Option[String] = {
+    val evaluated = evaluator.evaluate(OSet(target))
+    evaluated.transitive.foreach {
+      case t: define.Source => watch(t.handle.path)
+      case _ => // do nothing
+    }
+
+    val errorStr =
+      (for((k, fs) <- evaluated.failing.items()) yield {
+        val ks = k match{
+          case Left(t) => t.toString
+          case Right(t) => Mirror.renderSelector(t.segments.toList)
+        }
+        val fss = fs.map{
+          case Result.Exception(t) => t.toString
+          case Result.Failure(t) => t
+        }
+        s"$ks ${fss.mkString(", ")}"
+      }).mkString("\n")
+
+    evaluated.failing.keyCount match {
+      case 0 => None
+      case n => Some(s"$n targets failed\n$errorStr")
+    }
+  }
+
+  def parseSelector(input: String) = {
+    import fastparse.all._
+    val segment = P( CharsWhileIn(('a' to 'z') ++ ('A' to 'Z') ++ ('0' to '9')).! ).map(
+      Mirror.Segment.Label
+    )
+    val crossSegment = P( "[" ~ CharsWhile(c => c != ',' && c != ']').!.rep(1, sep=",") ~ "]" ).map(
+      Mirror.Segment.Cross
+    )
+    val query = P( segment ~ ("." ~ segment | crossSegment).rep ~ End ).map{
+      case (h, rest) => h :: rest.toList
+    }
+    query.parse(input)
+  }
+
+
+
+  def parseArgs(selectorString: String): Either[String, List[Mirror.Segment]] = {
+    import fastparse.all.Parsed
+    if (selectorString.isEmpty) Left("Selector cannot be empty")
+    else parseSelector(selectorString) match {
+      case f: Parsed.Failure => Left(s"Parsing exception ${f.msg}")
+      case Parsed.Success(selector, _) => Right(selector)
+    }
+  }
+
+
+  def consistencyCheck[T](mapping: Discovered.Mapping[T]): Either[String, Unit] = {
+    val consistencyErrors = Discovered.consistencyCheck(mapping)
+    if (consistencyErrors.nonEmpty) {
+      Left(s"Failed Discovered.consistencyCheck: ${consistencyErrors.map(Mirror.renderSelector)}")
+    } else {
+      Right(())
+    }
+  }
+}
\ No newline at end of file
diff --git a/core/src/test/scala/mill/main/MainTests.scala b/core/src/test/scala/mill/main/MainTests.scala
index 73060479..3e0f606b 100644
--- a/core/src/test/scala/mill/main/MainTests.scala
+++ b/core/src/test/scala/mill/main/MainTests.scala
@@ -13,7 +13,7 @@ object MainTests extends TestSuite{
                expected: Either[String, Task[_]]) = {
 
     val resolved = for{
-      args <- mill.Main.parseArgs(selectorString)
+      args <- mill.main.RunScript.parseArgs(selectorString)
       val crossSelectors = args.map{case Mirror.Segment.Cross(x) => x.toList.map(_.toString) case _ => Nil}
       task <- mill.main.Resolve.resolve(args, mapping.mirror, mapping.base, Nil, crossSelectors, Nil)
     } yield task