Simplify `MainRunner` and `RunScript` code

4 files changed, 49 insertions, 54 deletions

diff --git a/core/src/main/scala/mill/main/CustomCodeWrapper.scala b/core/src/main/scala/mill/main/CustomCodeWrapper.scala
deleted file mode 100644
index f92850f4..00000000
--- a/core/src/main/scala/mill/main/CustomCodeWrapper.scala
+++ /dev/null
@@ -1,30 +0,0 @@
-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) = {
-    val wrapName = indexedWrapperName.backticked
-    s"""
-       |package ${pkgName.head.encoded}
-       |package ${Util.encodeScalaSourcePath(pkgName.tail)}
-       |$imports
-       |import mill._
-       |
-       |object $wrapName extends $wrapName with mill.main.MainWrapper[$wrapName]{
-       |  lazy val discovered = mill.discover.Discovered.make[$wrapName]
-       |}
-       |
-       |sealed abstract class $wrapName extends mill.Module{
-       |""".stripMargin
-  }
-
-
-  def bottom(printCode: String, indexedWrapperName: Name, extraCode: String) = {
-    // We need to disable the `$main` method definition inside the wrapper
-    // class, because otherwise it might get picked up by Ammonite and run as
-    // a static class, which naturally blows up
-    "\n}"
-  }
-}
diff --git a/core/src/main/scala/mill/main/MainRunner.scala b/core/src/main/scala/mill/main/MainRunner.scala
index b99c4dc4..fb0f3d7c 100644
--- a/core/src/main/scala/mill/main/MainRunner.scala
+++ b/core/src/main/scala/mill/main/MainRunner.scala
@@ -1,15 +1,21 @@
 package mill.main
+import ammonite.interp.Preprocessor
 import ammonite.ops.Path
-import ammonite.util.Res
+import ammonite.util.{Imports, Name, Res, Util}
 import mill.discover.Discovered
 import mill.eval.Evaluator
 
+/**
+  * Customized version of [[ammonite.MainRunner]], allowing us to run Mill
+  * `build.sc` scripts with mill-specific tweaks such as a custom
+  * `scriptCodeWrapper` or with a persistent evaluator between runs.
+  */
 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
+  var lastEvaluator: Option[(Seq[(Path, Long)], Evaluator[_])] = None
   override def runScript(scriptPath: Path, scriptArgs: List[String]) =
     watchLoop(
       isRepl = false,
@@ -24,8 +30,8 @@ class MainRunner(config: ammonite.main.Cli.Config)
 
             val interpWatched = interp.watchedFiles
             result match{
-              case Res.Success((mapping, eval, evaluationWatches, success)) =>
-                lastEvaluator = Some((interpWatched, mapping, eval))
+              case Res.Success((eval, evaluationWatches, success)) =>
+                lastEvaluator = Some((interpWatched, eval))
                 (result, interpWatched ++ evaluationWatches)
               case _ =>
                 (result, interpWatched)
@@ -36,9 +42,40 @@ class MainRunner(config: ammonite.main.Cli.Config)
       }
     )
   override def initMain(isRepl: Boolean) = {
-    super.initMain(isRepl).copy(scriptCodeWrapper = mill.main.CustomCodeWrapper)
+    super.initMain(isRepl).copy(scriptCodeWrapper = mill.main.MainRunner.CustomCodeWrapper)
   }
   override def handleWatchRes[T](res: Res[T], printing: Boolean) = {
     super.handleWatchRes(res, printing = false)
   }
 }
+
+object MainRunner{
+  object CustomCodeWrapper extends Preprocessor.CodeWrapper {
+    def top(pkgName: Seq[Name], imports: Imports, indexedWrapperName: Name) = {
+      val wrapName = indexedWrapperName.backticked
+      s"""
+         |package ${pkgName.head.encoded}
+         |package ${Util.encodeScalaSourcePath(pkgName.tail)}
+         |$imports
+         |import mill._
+         |
+         |object $wrapName extends $wrapName{
+         |  // Stub to make sure Ammonite has something to call after it evaluates a script,
+         |  // even if it does nothing...
+         |  def $$main() = Iterator[String]()
+         |  lazy val mapping = mill.discover.Discovered.make[$wrapName].mapping(this)
+         |}
+         |
+         |sealed abstract class $wrapName extends mill.Module{
+         |""".stripMargin
+    }
+
+
+    def bottom(printCode: String, indexedWrapperName: Name, extraCode: String) = {
+      // We need to disable the `$main` method definition inside the wrapper class,
+      // because otherwise it might get picked up by Ammonite and run as a static
+      // class method, which blows up since it's defined as an instance method
+      "\n}"
+    }
+  }
+}
\ No newline at end of file
diff --git a/core/src/main/scala/mill/main/MainWrapper.scala b/core/src/main/scala/mill/main/MainWrapper.scala
deleted file mode 100644
index 16fd58c8..00000000
--- a/core/src/main/scala/mill/main/MainWrapper.scala
+++ /dev/null
@@ -1,12 +0,0 @@
-package mill.main
-
-/**
-  * Class that wraps each Mill build file.
-  */
-trait MainWrapper[T]{
-  val discovered: mill.discover.Discovered[T]
-  // Stub to make sure Ammonite has something to call after it evaluates a script,
-  // even if it does nothing...
-  def $main() = Iterator[String]()
-  lazy val mapping = discovered.mapping(this.asInstanceOf[T])
-}
diff --git a/core/src/main/scala/mill/main/RunScript.scala b/core/src/main/scala/mill/main/RunScript.scala
index e201beaf..f4ad375b 100644
--- a/core/src/main/scala/mill/main/RunScript.scala
+++ b/core/src/main/scala/mill/main/RunScript.scala
@@ -25,20 +25,20 @@ object RunScript{
                 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)] = {
+                lastEvaluator: Option[(Seq[(Path, Long)], Evaluator[_])])
+  : Res[(Evaluator[_], Seq[(Path, Long)], Boolean)] = {
 
     val log = new PrintLogger(true)
     for{
-      (mapping, evaluator) <- lastEvaluator match{
-        case Some((prevInterpWatchedSig, prevMapping, prevEvaluator))
+      evaluator <- lastEvaluator match{
+        case Some((prevInterpWatchedSig, prevEvaluator))
           if watchedSigUnchanged(prevInterpWatchedSig) =>
-          Res.Success((prevMapping, prevEvaluator))
+          Res.Success(prevEvaluator)
 
         case _ =>
           interp.watch(path)
           for(mapping <- evaluateMapping(wd, path, interp))
-          yield (mapping, new Evaluator(pwd / 'out, mapping, log))
+          yield new Evaluator(pwd / 'out, mapping, log)
       }
     } yield {
       val evaluationWatches = mutable.Buffer.empty[(Path, Long)]
@@ -47,7 +47,7 @@ object RunScript{
         scriptArgs,
         p => evaluationWatches.append((p, Interpreter.pathSignature(p)))
       )
-      (mapping, evaluator, evaluationWatches, res.isRight)
+      (evaluator, evaluationWatches, res.isRight)
     }
   }
   def watchedSigUnchanged(sig: Seq[(Path, Long)]) = {