Simplify `MainRunner` and `RunScript` code

1 files changed, 42 insertions, 5 deletions

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